[mypyc] Fix reference leak on list set item failure

mypyc/lib-rt/list_ops.c

@@ -14,6 +14,12 @@ PyObject *CPyList_Build(Py_ssize_t len, ...) {
 
     PyObject *res = PyList_New(len);
     if (res == NULL) {
+        va_list args;
+        va_start(args, len);
+        for (i = 0; i < len; i++) {
+            Py_DECREF(va_arg(args, PyObject *));
+        }
+        va_end(args);
         return NULL;
     }
 
@@ -123,22 +129,26 @@ bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value) {
         Py_ssize_t size = PyList_GET_SIZE(list);
         if (n >= 0) {
             if (n >= size) {
+                Py_DECREF(value);
                 PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
                 return false;
             }
         } else {
             n += size;
             if (n < 0) {
+                Py_DECREF(value);
                 PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
                 return false;
             }
         }
-        // PyList_SET_ITEM doesn't decref the old element, so we do
-        Py_DECREF(PyList_GET_ITEM(list, n));
+        // Set the replacement before decref, which may run reentrant Python code.
+        PyObject *old = PyList_GET_ITEM(list, n);
         // N.B: Steals reference
         PyList_SET_ITEM(list, n, value);
+        Py_DECREF(old);
         return true;
     } else {
+        Py_DECREF(value);
         PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
         return false;
     }
@@ -147,20 +157,23 @@ bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value) {
 bool CPyList_SetItemInt64(PyObject *list, int64_t index, PyObject *value) {
     size_t size = PyList_GET_SIZE(list);
     if (unlikely((uint64_t)index >= size)) {
-        if (index > 0) {
+        if (index >= 0) {
+            Py_DECREF(value);
             PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
             return false;
         }
         index += size;
         if (index < 0) {
+            Py_DECREF(value);
             PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
             return false;
         }
     }
-    // PyList_SET_ITEM doesn't decref the old element, so we do
-    Py_DECREF(PyList_GET_ITEM(list, index));
+    // Set the replacement before decref, which may run reentrant Python code.
+    PyObject *old = PyList_GET_ITEM(list, index);
     // N.B: Steals reference
     PyList_SET_ITEM(list, index, value);
+    Py_DECREF(old);
     return true;
 }
 
@@ -188,6 +201,7 @@ bool CPyList_SetItem_(PyObject *list, CPyTagged index, PyObject *value) {
         }
         return PyList_SetItem(list, n, value) >= 0;
     } else {
+        Py_DECREF(value);
         PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
         return false;
     }

mypyc/test-data/run-i64.test

@@ -1581,3 +1581,49 @@ def test_i64_definedness_across_yield() -> None:
     assert list(conditional(True)) == [0, 5]
     with assertRaises(AttributeError):
         list(conditional(False))
+
+[case testI64ListSetErrorOwnership]
+from typing import Any, List
+from mypy_extensions import i64
+def assign(items: List[Any], index: i64, value: Any) -> None:
+    items[index] = value
+
+[file driver.py]
+import gc
+import weakref
+from native import assign
+from testutil import assertRaises
+
+class Sentinel:
+    pass
+
+for index in (1, -2, 2**63 - 1, -2**63):
+    value = Sentinel()
+    ref = weakref.ref(value)
+    try:
+        assign([None], index, value)
+    except (IndexError, OverflowError):
+        pass
+    else:
+        assert False, index
+    del value
+    gc.collect()
+    assert ref() is None, index
+with assertRaises(IndexError):
+    assign([], 0, None)
+items = [None]
+value = Sentinel()
+assign(items, -1, value)
+assert items[0] is value
+assign(items, 0, None)
+assert items == [None]
+
+# Replacing an element may run Python code that mutates the same list.
+class Reentrant:
+    def __del__(self):
+        items.clear()
+
+for index in (0, -1):
+    items = [Reentrant()]
+    assign(items, index, None)
+    assert items == []

mypyc/test-data/run-lists.test

@@ -671,3 +671,48 @@ class subc(list[Any]):
 [file userdefinedlist.py]
 class list:
     pass
+
+[case testListSetErrorOwnership]
+from typing import Any, List
+def assign(items: List[Any], index: int, value: Any) -> None:
+    items[index] = value
+
+[file driver.py]
+import gc
+import weakref
+from native import assign
+from testutil import assertRaises
+
+class Sentinel:
+    pass
+
+for index in (1, -2, 2**100, -2**100):
+    value = Sentinel()
+    ref = weakref.ref(value)
+    try:
+        assign([None], index, value)
+    except (IndexError, OverflowError):
+        pass
+    else:
+        assert False, index
+    del value
+    gc.collect()
+    assert ref() is None, index
+with assertRaises(IndexError):
+    assign([], 0, None)
+items = [None]
+value = Sentinel()
+assign(items, -1, value)
+assert items[0] is value
+assign(items, 0, None)
+assert items == [None]
+
+# Replacing an element may run Python code that mutates the same list.
+class Reentrant:
+    def __del__(self):
+        items.clear()
+
+for index in (0, -1):
+    items = [Reentrant()]
+    assign(items, index, None)
+    assert items == []