[mypyc] Fix dict iteration memory safety on free-threaded builds (#21617)

`PyDict_Next()` returns a borrowed reference, so we need to add critical
sections around the calls to ensure memory safety.

I used coding agent assist.

Author: JukkaL

mypyc/lib-rt/dict_ops.c

@@ -346,16 +346,24 @@ tuple_T3CIO CPyDict_NextKey(PyObject *dict_or_iter, CPyTagged offset) {
     PyObject *dummy;
 
     if (PyDict_CheckExact(dict_or_iter)) {
-        ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &ret.f2, &dummy);
+        PyObject *key;
+        // PyDict_Next() returns a borrowed reference. On free-threaded builds,
+        // hold the dict lock until we have converted it to a strong reference.
+        Py_BEGIN_CRITICAL_SECTION(dict_or_iter);
+        ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &key, &dummy);
+        if (ret.f0) {
+            ret.f2 = Py_NewRef(key);
+        }
+        Py_END_CRITICAL_SECTION();
+
         if (ret.f0) {
             ret.f1 = CPyTagged_FromSsize_t(py_offset);
         } else {
             // Set key to None, so mypyc can manage refcounts.
             ret.f1 = 0;
             ret.f2 = Py_None;
+            Py_INCREF(ret.f2);
         }
-        // PyDict_Next() returns borrowed references.
-        Py_INCREF(ret.f2);
     } else {
         // offset is dummy in this case, just use the old value.
         ret.f1 = offset;
@@ -370,16 +378,24 @@ tuple_T3CIO CPyDict_NextValue(PyObject *dict_or_iter, CPyTagged offset) {
     PyObject *dummy;
 
     if (PyDict_CheckExact(dict_or_iter)) {
-        ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &dummy, &ret.f2);
+        PyObject *value;
+        // PyDict_Next() returns a borrowed reference. On free-threaded builds,
+        // hold the dict lock until we have converted it to a strong reference.
+        Py_BEGIN_CRITICAL_SECTION(dict_or_iter);
+        ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &dummy, &value);
+        if (ret.f0) {
+            ret.f2 = Py_NewRef(value);
+        }
+        Py_END_CRITICAL_SECTION();
+
         if (ret.f0) {
             ret.f1 = CPyTagged_FromSsize_t(py_offset);
         } else {
             // Set value to None, so mypyc can manage refcounts.
             ret.f1 = 0;
             ret.f2 = Py_None;
+            Py_INCREF(ret.f2);
         }
-        // PyDict_Next() returns borrowed references.
-        Py_INCREF(ret.f2);
     } else {
         // offset is dummy in this case, just use the old value.
         ret.f1 = offset;
@@ -393,14 +409,27 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) {
     Py_ssize_t py_offset = CPyTagged_AsSsize_t(offset);
 
     if (PyDict_CheckExact(dict_or_iter)) {
-        ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &ret.f2, &ret.f3);
+        PyObject *key;
+        PyObject *value;
+        // PyDict_Next() returns borrowed references. On free-threaded builds,
+        // hold the dict lock until we have converted them to strong references.
+        Py_BEGIN_CRITICAL_SECTION(dict_or_iter);
+        ret.f0 = PyDict_Next(dict_or_iter, &py_offset, &key, &value);
+        if (ret.f0) {
+            ret.f2 = Py_NewRef(key);
+            ret.f3 = Py_NewRef(value);
+        }
+        Py_END_CRITICAL_SECTION();
+
         if (ret.f0) {
             ret.f1 = CPyTagged_FromSsize_t(py_offset);
         } else {
             // Set key and value to None, so mypyc can manage refcounts.
             ret.f1 = 0;
             ret.f2 = Py_None;
             ret.f3 = Py_None;
+            Py_INCREF(ret.f2);
+            Py_INCREF(ret.f3);
         }
     } else {
         ret.f1 = offset;
@@ -413,6 +442,8 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) {
             ret.f0 = 0;
             ret.f2 = Py_None;
             ret.f3 = Py_None;
+            Py_INCREF(ret.f2);
+            Py_INCREF(ret.f3);
         } else {
             ret.f0 = 1;
             ret.f2 = PyTuple_GET_ITEM(item, 0);
@@ -423,9 +454,6 @@ tuple_T4CIOO CPyDict_NextItem(PyObject *dict_or_iter, CPyTagged offset) {
             return ret;
         }
     }
-    // PyDict_Next() returns borrowed references.
-    Py_INCREF(ret.f2);
-    Py_INCREF(ret.f3);
     return ret;
 }