gh-115999: Add free-threaded specialization for `STORE_SUBSCR` (#127169)

The specialization only depends on the type, so no special thread-safety
considerations there.

STORE_SUBSCR_LIST_INT needs to lock the list before modifying it.

`_PyDict_SetItem_Take2` already internally locks the dictionary using a
critical section.

5 files changed, 109 insertions, 70 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 88e96af..a14b32b 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -910,7 +910,7 @@ dummy_func(
         };
 
         specializing op(_SPECIALIZE_STORE_SUBSCR, (counter/1, container, sub -- container, sub)) {
-            #if ENABLE_SPECIALIZATION
+            #if ENABLE_SPECIALIZATION_FT
             if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                 next_instr = this_instr;
                 _Py_Specialize_StoreSubscr(container, sub, next_instr);
@@ -918,7 +918,7 @@ dummy_func(
             }
             OPCODE_DEFERRED_INC(STORE_SUBSCR);
             ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-            #endif  /* ENABLE_SPECIALIZATION */
+            #endif  /* ENABLE_SPECIALIZATION_FT */
         }
 
         op(_STORE_SUBSCR, (v, container, sub -- )) {
@@ -940,13 +940,18 @@ dummy_func(
             // Ensure nonnegative, zero-or-one-digit ints.
             DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub));
             Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
+            DEOPT_IF(!LOCK_OBJECT(list));
             // Ensure index < len(list)
-            DEOPT_IF(index >= PyList_GET_SIZE(list));
+            if (index >= PyList_GET_SIZE(list)) {
+                UNLOCK_OBJECT(list);
+                DEOPT_IF(true);
+            }
             STAT_INC(STORE_SUBSCR, hit);
 
             PyObject *old_value = PyList_GET_ITEM(list, index);
             PyList_SET_ITEM(list, index, PyStackRef_AsPyObjectSteal(value));
             assert(old_value != NULL);
+            UNLOCK_OBJECT(list);  // unlock before decrefs!
             Py_DECREF(old_value);
             PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
             DEAD(sub_st);
diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
index 603b71e..9250b86 100644
--- a/Python/ceval_macros.h
+++ b/Python/ceval_macros.h
@@ -284,6 +284,29 @@ GETITEM(PyObject *v, Py_ssize_t i) {
     }
 
 
+// Try to lock an object in the free threading build, if it's not already
+// locked. Use with a DEOPT_IF() to deopt if the object is already locked.
+// These are no-ops in the default GIL build. The general pattern is:
+//
+// DEOPT_IF(!LOCK_OBJECT(op));
+// if (/* condition fails */) {
+//     UNLOCK_OBJECT(op);
+//     DEOPT_IF(true);
+//  }
+//  ...
+//  UNLOCK_OBJECT(op);
+//
+// NOTE: The object must be unlocked on every exit code path and you should
+// avoid any potentially escaping calls (like PyStackRef_CLOSE) while the
+// object is locked.
+#ifdef Py_GIL_DISABLED
+#  define LOCK_OBJECT(op) PyMutex_LockFast(&(_PyObject_CAST(op))->ob_mutex._bits)
+#  define UNLOCK_OBJECT(op) PyMutex_Unlock(&(_PyObject_CAST(op))->ob_mutex)
+#else
+#  define LOCK_OBJECT(op) (1)
+#  define UNLOCK_OBJECT(op) ((void)0)
+#endif
+
 #define GLOBALS() frame->f_globals
 #define BUILTINS() frame->f_builtins
 #define LOCALS() frame->f_locals
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 5af970e..d46412a 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -1235,15 +1235,23 @@
                 JUMP_TO_JUMP_TARGET();
             }
             Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
-            // Ensure index < len(list)
-            if (index >= PyList_GET_SIZE(list)) {
+            if (!LOCK_OBJECT(list)) {
                 UOP_STAT_INC(uopcode, miss);
                 JUMP_TO_JUMP_TARGET();
             }
+            // Ensure index < len(list)
+            if (index >= PyList_GET_SIZE(list)) {
+                UNLOCK_OBJECT(list);
+                if (true) {
+                    UOP_STAT_INC(uopcode, miss);
+                    JUMP_TO_JUMP_TARGET();
+                }
+            }
             STAT_INC(STORE_SUBSCR, hit);
             PyObject *old_value = PyList_GET_ITEM(list, index);
             PyList_SET_ITEM(list, index, PyStackRef_AsPyObjectSteal(value));
             assert(old_value != NULL);
+            UNLOCK_OBJECT(list);  // unlock before decrefs!
             Py_DECREF(old_value);
             PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
             PyStackRef_CLOSE(list_st);
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 36ec727..c9a5132 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -7629,7 +7629,7 @@
                 container = stack_pointer[-2];
                 uint16_t counter = read_u16(&this_instr[1].cache);
                 (void)counter;
-                #if ENABLE_SPECIALIZATION
+                #if ENABLE_SPECIALIZATION_FT
                 if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                     next_instr = this_instr;
                     _PyFrame_SetStackPointer(frame, stack_pointer);
@@ -7639,7 +7639,7 @@
                 }
                 OPCODE_DEFERRED_INC(STORE_SUBSCR);
                 ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-                #endif  /* ENABLE_SPECIALIZATION */
+                #endif  /* ENABLE_SPECIALIZATION_FT */
             }
             // _STORE_SUBSCR
             {
@@ -7704,12 +7704,17 @@
             // Ensure nonnegative, zero-or-one-digit ints.
             DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub), STORE_SUBSCR);
             Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
+            DEOPT_IF(!LOCK_OBJECT(list), STORE_SUBSCR);
             // Ensure index < len(list)
-            DEOPT_IF(index >= PyList_GET_SIZE(list), STORE_SUBSCR);
+            if (index >= PyList_GET_SIZE(list)) {
+                UNLOCK_OBJECT(list);
+                DEOPT_IF(true, STORE_SUBSCR);
+            }
             STAT_INC(STORE_SUBSCR, hit);
             PyObject *old_value = PyList_GET_ITEM(list, index);
             PyList_SET_ITEM(list, index, PyStackRef_AsPyObjectSteal(value));
             assert(old_value != NULL);
+            UNLOCK_OBJECT(list);  // unlock before decrefs!
             Py_DECREF(old_value);
             PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
             PyStackRef_CLOSE(list_st);
diff --git a/Python/specialize.c b/Python/specialize.c
index c1f4b06..172dae7 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -1814,86 +1814,54 @@ success:
     cache->counter = adaptive_counter_cooldown();
 }
 
-void
-_Py_Specialize_StoreSubscr(_PyStackRef container_st, _PyStackRef sub_st, _Py_CODEUNIT *instr)
-{
-    PyObject *container = PyStackRef_AsPyObjectBorrow(container_st);
-    PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
 
-    assert(ENABLE_SPECIALIZATION);
-    _PyStoreSubscrCache *cache = (_PyStoreSubscrCache *)(instr + 1);
-    PyTypeObject *container_type = Py_TYPE(container);
-    if (container_type == &PyList_Type) {
-        if (PyLong_CheckExact(sub)) {
-            if (_PyLong_IsNonNegativeCompact((PyLongObject *)sub)
-                && ((PyLongObject *)sub)->long_value.ob_digit[0] < (size_t)PyList_GET_SIZE(container))
-            {
-                instr->op.code = STORE_SUBSCR_LIST_INT;
-                goto success;
-            }
-            else {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
-                goto fail;
-            }
-        }
-        else if (PySlice_Check(sub)) {
-            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_LIST_SLICE);
-            goto fail;
-        }
-        else {
-            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
-            goto fail;
-        }
-    }
-    if (container_type == &PyDict_Type) {
-        instr->op.code = STORE_SUBSCR_DICT;
-        goto success;
-    }
 #ifdef Py_STATS
+static int
+store_subscr_fail_kind(PyObject *container_type)
+{
     PyMappingMethods *as_mapping = container_type->tp_as_mapping;
     if (as_mapping && (as_mapping->mp_ass_subscript
                        == PyDict_Type.tp_as_mapping->mp_ass_subscript)) {
-        SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_DICT_SUBCLASS_NO_OVERRIDE);
-        goto fail;
+        return SPEC_FAIL_SUBSCR_DICT_SUBCLASS_NO_OVERRIDE;
     }
     if (PyObject_CheckBuffer(container)) {
         if (PyLong_CheckExact(sub) && (!_PyLong_IsNonNegativeCompact((PyLongObject *)sub))) {
-            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
+            return SPEC_FAIL_OUT_OF_RANGE;
         }
         else if (strcmp(container_type->tp_name, "array.array") == 0) {
             if (PyLong_CheckExact(sub)) {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_ARRAY_INT);
+                return SPEC_FAIL_SUBSCR_ARRAY_INT;
             }
             else if (PySlice_Check(sub)) {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_ARRAY_SLICE);
+                return SPEC_FAIL_SUBSCR_ARRAY_SLICE;
             }
             else {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+                return SPEC_FAIL_OTHER;
             }
         }
         else if (PyByteArray_CheckExact(container)) {
             if (PyLong_CheckExact(sub)) {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BYTEARRAY_INT);
+                return SPEC_FAIL_SUBSCR_BYTEARRAY_INT;
             }
             else if (PySlice_Check(sub)) {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BYTEARRAY_SLICE);
+                return SPEC_FAIL_SUBSCR_BYTEARRAY_SLICE;
             }
             else {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+                return SPEC_FAIL_OTHER;
             }
         }
         else {
             if (PyLong_CheckExact(sub)) {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BUFFER_INT);
+                return SPEC_FAIL_SUBSCR_BUFFER_INT;
             }
             else if (PySlice_Check(sub)) {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BUFFER_SLICE);
+                return SPEC_FAIL_SUBSCR_BUFFER_SLICE;
             }
             else {
-                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+                return SPEC_FAIL_OTHER;
             }
         }
-        goto fail;
+        return SPEC_FAIL_OTHER;
     }
     PyObject *descriptor = _PyType_Lookup(container_type, &_Py_ID(__setitem__));
     if (descriptor && Py_TYPE(descriptor) == &PyFunction_Type) {
@@ -1901,25 +1869,55 @@ _Py_Specialize_StoreSubscr(_PyStackRef container_st, _PyStackRef sub_st, _Py_COD
         PyCodeObject *code = (PyCodeObject *)func->func_code;
         int kind = function_kind(code);
         if (kind == SIMPLE_FUNCTION) {
-            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_PY_SIMPLE);
+            return SPEC_FAIL_SUBSCR_PY_SIMPLE;
         }
         else {
-            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_PY_OTHER);
+            return SPEC_FAIL_SUBSCR_PY_OTHER;
         }
-        goto fail;
     }
-#endif   // Py_STATS
-    SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
-fail:
-    STAT_INC(STORE_SUBSCR, failure);
-    assert(!PyErr_Occurred());
-    instr->op.code = STORE_SUBSCR;
-    cache->counter = adaptive_counter_backoff(cache->counter);
-    return;
-success:
-    STAT_INC(STORE_SUBSCR, success);
-    assert(!PyErr_Occurred());
-    cache->counter = adaptive_counter_cooldown();
+    return SPEC_FAIL_OTHER;
+}
+#endif
+
+void
+_Py_Specialize_StoreSubscr(_PyStackRef container_st, _PyStackRef sub_st, _Py_CODEUNIT *instr)
+{
+    PyObject *container = PyStackRef_AsPyObjectBorrow(container_st);
+    PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
+
+    assert(ENABLE_SPECIALIZATION_FT);
+    PyTypeObject *container_type = Py_TYPE(container);
+    if (container_type == &PyList_Type) {
+        if (PyLong_CheckExact(sub)) {
+            if (_PyLong_IsNonNegativeCompact((PyLongObject *)sub)
+                && ((PyLongObject *)sub)->long_value.ob_digit[0] < (size_t)PyList_GET_SIZE(container))
+            {
+                specialize(instr, STORE_SUBSCR_LIST_INT);
+                return;
+            }
+            else {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
+                unspecialize(instr);
+                return;
+            }
+        }
+        else if (PySlice_Check(sub)) {
+            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_LIST_SLICE);
+            unspecialize(instr);
+            return;
+        }
+        else {
+            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+            unspecialize(instr);
+            return;
+        }
+    }
+    if (container_type == &PyDict_Type) {
+        specialize(instr, STORE_SUBSCR_DICT);
+        return;
+    }
+    SPECIALIZATION_FAIL(STORE_SUBSCR, store_subscr_fail_kind(container_type));
+    unspecialize(instr);
 }
 
 /* Returns a borrowed reference.