Repository: python/mypy
Pull request: [mypyc] Add incref in CPyDict_SetDefault
### Description

Closes https://github.com/mypyc/mypyc/issues/831

## Test Plan


Problem statement:
Bug: Segmentation fault using setdefault
code following will trigger segmentation fault on my computer:
``` python
from typing import Set, Dict

def test_dict() -> None:
    d3: Dict[str, Set[int]] = {'a': {1}}
    assert d3.setdefault('b', set()) == set()

test_dict()
```
environment: gcc (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
python: 3.8

It's interesting that following code runs well:
``` python
from typing import Set, Dict

d3: Dict[str, Set[int]] = {'a': {1}}
assert d3.setdefault('b', set()) == set()
```

Hints from issue discussion:
can you verify that before commit f463a3921fcd5cb360c12a84650880a2a92e0566 the code work?
I cannot reproduce this on my mac.
Thank you for your work, I will figure out what happened after finishing my gsoc proposal. I may not be quick to respond this time.
Sorry to respond late. The problem still exists on master branch and after checkout before commit f463a3921fcd5cb360c12a84650880a2a92e0566 the problem disappears.(In addition, my python version is 3.8)
you can try pasting the compiled C code here so we can see if there anything abnormal here.
```C
char CPyDef_test_dict(void) {
    PyObject *cpy_r_r0;
    PyObject *cpy_r_r1;
    PyObject *cpy_r_r2;
    int32_t cpy_r_r3;
    char cpy_r_r4;
    PyObject *cpy_r_r5;
    PyObject *cpy_r_d3;
    PyObject *cpy_r_r6;
    PyObject *cpy_r_r7;
    PyObject *cpy_r_r8;
    PyObject *cpy_r_r9;
    PyObject *cpy_r_r10;
    PyObject *cpy_r_r11;
    char cpy_r_r12;
    char cpy_r_r13;
    char cpy_r_r14;
CPyL0: ;
    cpy_r_r0 = CPyStatics[3]; /* 'a' */
    cpy_r_r1 = PySet_New(NULL);
    if (unlikely(cpy_r_r1 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 4, CPyStatic_globals);
        goto CPyL13;
    }
CPyL1: ;
    cpy_r_r2 = CPyTagged_StealAsObject(2);
    cpy_r_r3 = PySet_Add(cpy_r_r1, cpy_r_r2);
    CPy_DecRef(cpy_r_r2);
    cpy_r_r4 = cpy_r_r3 >= 0;
    if (unlikely(!cpy_r_r4)) {
        CPy_AddTraceback("spec.py", "test_dict", 4, CPyStatic_globals);
        goto CPyL14;
    }
CPyL2: ;
    cpy_r_r5 = CPyDict_Build(1, cpy_r_r0, cpy_r_r1);
    CPy_DecRef(cpy_r_r1);
    if (unlikely(cpy_r_r5 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 4, CPyStatic_globals);
        goto CPyL13;
    }
CPyL3: ;
    cpy_r_d3 = cpy_r_r5;
    cpy_r_r6 = CPyStatics[4]; /* 'b' */
    cpy_r_r7 = PySet_New(NULL);
    if (unlikely(cpy_r_r7 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL15;
    }
CPyL4: ;
    cpy_r_r8 = CPyDict_SetDefault(cpy_r_d3, cpy_r_r6, cpy_r_r7);
    CPy_DecRef(cpy_r_d3);
    CPy_DecRef(cpy_r_r7);
    if (unlikely(cpy_r_r8 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL13;
    }
CPyL5: ;
    CPy_INCREF(cpy_r_r8);
    if (likely(PySet_Check(cpy_r_r8)))
        cpy_r_r9 = cpy_r_r8;
    else {
        CPy_TypeError("set", cpy_r_r8);
        cpy_r_r9 = NULL;
    }
    if (unlikely(cpy_r_r9 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL13;
    }
CPyL6: ;
    cpy_r_r10 = PySet_New(NULL);
    if (unlikely(cpy_r_r10 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL16;
    }
CPyL7: ;
    cpy_r_r11 = PyObject_RichCompare(cpy_r_r9, cpy_r_r10, 2);
    CPy_DecRef(cpy_r_r9);
    CPy_DecRef(cpy_r_r10);
    if (unlikely(cpy_r_r11 == NULL)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL13;
    }
CPyL8: ;
    if (unlikely(!PyBool_Check(cpy_r_r11))) {
        CPy_TypeError("bool", cpy_r_r11);
        cpy_r_r12 = 2;
    } else
        cpy_r_r12 = cpy_r_r11 == Py_True;
    CPy_DecRef(cpy_r_r11);
    if (unlikely(cpy_r_r12 == 2)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL13;
    }
CPyL9: ;
    if (cpy_r_r12) goto CPyL12;
CPyL10: ;
    PyErr_SetNone(PyExc_AssertionError);
    cpy_r_r13 = 0;
    if (unlikely(!cpy_r_r13)) {
        CPy_AddTraceback("spec.py", "test_dict", 5, CPyStatic_globals);
        goto CPyL13;
    }
CPyL11: ;
    CPy_Unreachable();
CPyL12: ;
    return 1;
CPyL13: ;
    cpy_r_r14 = 2;
    return cpy_r_r14;
CPyL14: ;
    CPy_DecRef(cpy_r_r1);
    goto CPyL13;
CPyL15: ;
    CPy_DecRef(cpy_r_d3);
    goto CPyL13;
CPyL16: ;
    CPy_DecRef(cpy_r_r9);
    goto CPyL13;
}


```
The generated C code on my Mac is exactly the same as yours. However I didn't encounter segmentation fault when running compiled module.

macOS 11.1 (Intel) 
Apple Clang 12.0.0 - x86_64-apple-darwin20.2.0
Python 3.7.9
Mypyc: Latest master branch
I changed to Python 3.8 and encountered the bug.

```
* thread #1, queue = 'com.apple.main-thread', stop reason = EXC_BAD_ACCESS (code=1, address=0x0)
    frame #0: 0x000000010009a9c6 python`set_dealloc + 38
python`set_dealloc:
->  0x10009a9c6 <+38>: movq   %rax, (%rcx)
    0x10009a9c9 <+41>: movq   0x8(%rax), %rdx
    0x10009a9cd <+45>: andl   $0x3, %edx
    0x10009a9d0 <+48>: orq    %rcx, %rdx
Target 0: (python) stopped.
```
back trace

```
* thread #1, queue = 'com.apple.main-thread', stop reason = EXC_BAD_ACCESS (code=1, address=0x0)
  * frame #0: 0x000000010009a9c6 python`set_dealloc + 38
    frame #1: 0x00000001014a8bb9 tmp5.cpython-38-darwin.so`CPyDef_test_dict at __native.c:127:5 [opt]
    frame #2: 0x00000001014a8a1b tmp5.cpython-38-darwin.so`CPyDef___top_level__ at __native.c:281:17 [opt]
    frame #3: 0x00000001014a8804 tmp5.cpython-38-darwin.so`PyInit_tmp5 at __native.c:44:19 [opt]
```

It seems that `CPy_DecRef(cpy_r_r9)` after `PyObject_RichCompare` triggers a set_dealloc operation.

I believe it's because PyDict_SetDefault returns a borrowed reference (https://docs.python.org/3/c-api/dict.html#c.PyDict_SetDefault). That means the reference is still owned by the caller, and if we want to hold onto the return value after the containing dict is removed, we need to INCREF it.

However, in https://github.com/python/mypy/blob/master/mypyc/lib-rt/dict_ops.c#L72 we return the reference directly, without INCREF. If my hypothesis is correct, adding an INCREF there should fix the segfault.
Interestingly the code generated by the two different segments are locally the same, they only differ in whether the operation resides in top-level or not. My guess is the difference in heap structure leads to this result.

Commit message: [mypyc] Add incref in CPyDict_SetDefault

Patch to review:
diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c
index c4dbe8d31c32..1416c9ba953a 100644
--- a/mypyc/lib-rt/dict_ops.c
+++ b/mypyc/lib-rt/dict_ops.c
@@ -69,7 +69,9 @@ PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key) {
 
 PyObject *CPyDict_SetDefault(PyObject *dict, PyObject *key, PyObject *value) {
     if (PyDict_CheckExact(dict)){
-        return PyDict_SetDefault(dict, key, value);
+        PyObject* ret = PyDict_SetDefault(dict, key, value);
+        Py_INCREF(ret);
+        return ret;
     }
     return PyObject_CallMethod(dict, "setdefault", "(OO)", key, value);
 }
diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index 267f9d79179d..e176ffcac316 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -124,7 +124,6 @@
     arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],
     return_type=object_rprimitive,
     c_function_name='CPyDict_SetDefault',
-    is_borrowed=True,
     error_kind=ERR_MAGIC)
 
 # dict.setdefault(key)


Code review comments:
diff_hunk: @@ -69,7 +69,9 @@ PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key) {
 
 PyObject *CPyDict_SetDefault(PyObject *dict, PyObject *key, PyObject *value) {
     if (PyDict_CheckExact(dict)){
-        return PyDict_SetDefault(dict, key, value);
+        PyObject* ret = PyDict_SetDefault(dict, key, value);
+        Py_INCREF(ret);
path: mypyc/lib-rt/dict_ops.c
text: @user1:
I think this should be Py_XINCREF in case SetDefault returned NULL.