bpo-45609: Specialize STORE_SUBSCR (GH-29242)

* Specialize STORE_SUBSCR for list[int], and dict[object]

* Adds _PyDict_SetItem_Take2 which consumes references to the key and values.

3 files changed, 147 insertions, 21 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 2b7b31c..9d3ff74 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -2279,6 +2279,8 @@ check_eval_breaker:
         }
 
         TARGET(STORE_SUBSCR) {
+            PREDICTED(STORE_SUBSCR);
+            STAT_INC(STORE_SUBSCR, unquickened);
             PyObject *sub = TOP();
             PyObject *container = SECOND();
             PyObject *v = THIRD();
@@ -2294,6 +2296,64 @@ check_eval_breaker:
             DISPATCH();
         }
 
+        TARGET(STORE_SUBSCR_ADAPTIVE) {
+            if (oparg == 0) {
+                PyObject *sub = TOP();
+                PyObject *container = SECOND();
+                next_instr--;
+                if (_Py_Specialize_StoreSubscr(container, sub, next_instr) < 0) {
+                    goto error;
+                }
+                DISPATCH();
+            }
+            else {
+                STAT_INC(STORE_SUBSCR, deferred);
+                // oparg is the adaptive cache counter
+                UPDATE_PREV_INSTR_OPARG(next_instr, oparg - 1);
+                STAT_DEC(STORE_SUBSCR, unquickened);
+                JUMP_TO_INSTRUCTION(STORE_SUBSCR);
+            }
+        }
+
+        TARGET(STORE_SUBSCR_LIST_INT) {
+            PyObject *sub = TOP();
+            PyObject *list = SECOND();
+            PyObject *value = THIRD();
+            DEOPT_IF(!PyLong_CheckExact(sub), STORE_SUBSCR);
+            DEOPT_IF(!PyList_CheckExact(list), STORE_SUBSCR);
+
+            // Ensure nonnegative, zero-or-one-digit ints.
+            DEOPT_IF(((size_t)Py_SIZE(sub)) > 1, STORE_SUBSCR);
+            Py_ssize_t index = ((PyLongObject*)sub)->ob_digit[0];
+            // Ensure index < len(list)
+            DEOPT_IF(index >= PyList_GET_SIZE(list), STORE_SUBSCR);
+            STAT_INC(STORE_SUBSCR, hit);
+
+            PyObject *old_value = PyList_GET_ITEM(list, index);
+            PyList_SET_ITEM(list, index, value);
+            STACK_SHRINK(3);
+            assert(old_value != NULL);
+            Py_DECREF(old_value);
+            Py_DECREF(sub);
+            Py_DECREF(list);
+            DISPATCH();
+        }
+
+        TARGET(STORE_SUBSCR_DICT) {
+            PyObject *sub = TOP();
+            PyObject *dict = SECOND();
+            PyObject *value = THIRD();
+            DEOPT_IF(!PyDict_CheckExact(dict), STORE_SUBSCR);
+            STACK_SHRINK(3);
+            STAT_INC(STORE_SUBSCR, hit);
+            int err = _PyDict_SetItem_Take2((PyDictObject *)dict, sub, value);
+            Py_DECREF(dict);
+            if (err != 0) {
+                goto error;
+            }
+            DISPATCH();
+        }
+
         TARGET(DELETE_SUBSCR) {
             PyObject *sub = TOP();
             PyObject *container = SECOND();
@@ -3374,15 +3434,13 @@ check_eval_breaker:
             PyObject *value = TOP();
             PyObject *key = SECOND();
             PyObject *map;
-            int err;
             STACK_SHRINK(2);
             map = PEEK(oparg);                      /* dict */
             assert(PyDict_CheckExact(map));
-            err = PyDict_SetItem(map, key, value);  /* map[key] = value */
-            Py_DECREF(value);
-            Py_DECREF(key);
-            if (err != 0)
+            /* map[key] = value */
+            if (_PyDict_SetItem_Take2((PyDictObject *)map, key, value) != 0) {
                 goto error;
+            }
             PREDICT(JUMP_ABSOLUTE);
             DISPATCH();
         }
@@ -4909,6 +4967,22 @@ opname ## _miss: \
         JUMP_TO_INSTRUCTION(opname); \
     }
 
+#define MISS_WITH_OPARG_COUNTER(opname) \
+opname ## _miss: \
+    { \
+        STAT_INC(opname, miss); \
+        uint8_t oparg = _Py_OPARG(next_instr[-1])-1; \
+        UPDATE_PREV_INSTR_OPARG(next_instr, oparg); \
+        assert(_Py_OPARG(next_instr[-1]) == oparg); \
+        if (oparg == 0) /* too many cache misses */ { \
+            oparg = ADAPTIVE_CACHE_BACKOFF; \
+            next_instr[-1] = _Py_MAKECODEUNIT(opname ## _ADAPTIVE, oparg); \
+            STAT_INC(opname, deopt); \
+        } \
+        STAT_DEC(opname, unquickened); \
+        JUMP_TO_INSTRUCTION(opname); \
+    }
+
 MISS_WITH_CACHE(LOAD_ATTR)
 MISS_WITH_CACHE(STORE_ATTR)
 MISS_WITH_CACHE(LOAD_GLOBAL)
@@ -4916,6 +4990,7 @@ MISS_WITH_CACHE(LOAD_METHOD)
 MISS_WITH_CACHE(CALL_FUNCTION)
 MISS_WITH_CACHE(BINARY_OP)
 MISS_WITH_CACHE(BINARY_SUBSCR)
+MISS_WITH_OPARG_COUNTER(STORE_SUBSCR)
 
 binary_subscr_dict_error:
         {
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 4703255..903b967 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -26,17 +26,20 @@ static void *opcode_targets[256] = {
     &&TARGET_BINARY_SUBSCR_TUPLE_INT,
     &&TARGET_BINARY_SUBSCR,
     &&TARGET_BINARY_SUBSCR_DICT,
-    &&TARGET_CALL_FUNCTION_ADAPTIVE,
-    &&TARGET_CALL_FUNCTION_BUILTIN_O,
-    &&TARGET_CALL_FUNCTION_BUILTIN_FAST,
+    &&TARGET_STORE_SUBSCR_ADAPTIVE,
+    &&TARGET_STORE_SUBSCR_LIST_INT,
+    &&TARGET_STORE_SUBSCR_DICT,
     &&TARGET_GET_LEN,
     &&TARGET_MATCH_MAPPING,
     &&TARGET_MATCH_SEQUENCE,
     &&TARGET_MATCH_KEYS,
-    &&TARGET_CALL_FUNCTION_LEN,
+    &&TARGET_CALL_FUNCTION_ADAPTIVE,
     &&TARGET_PUSH_EXC_INFO,
-    &&TARGET_CALL_FUNCTION_ISINSTANCE,
+    &&TARGET_CALL_FUNCTION_BUILTIN_O,
     &&TARGET_POP_EXCEPT_AND_RERAISE,
+    &&TARGET_CALL_FUNCTION_BUILTIN_FAST,
+    &&TARGET_CALL_FUNCTION_LEN,
+    &&TARGET_CALL_FUNCTION_ISINSTANCE,
     &&TARGET_CALL_FUNCTION_PY_SIMPLE,
     &&TARGET_JUMP_ABSOLUTE_QUICK,
     &&TARGET_LOAD_ATTR_ADAPTIVE,
@@ -45,28 +48,25 @@ static void *opcode_targets[256] = {
     &&TARGET_LOAD_ATTR_SLOT,
     &&TARGET_LOAD_ATTR_MODULE,
     &&TARGET_LOAD_GLOBAL_ADAPTIVE,
-    &&TARGET_LOAD_GLOBAL_MODULE,
-    &&TARGET_LOAD_GLOBAL_BUILTIN,
-    &&TARGET_LOAD_METHOD_ADAPTIVE,
     &&TARGET_WITH_EXCEPT_START,
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
     &&TARGET_BEFORE_ASYNC_WITH,
     &&TARGET_BEFORE_WITH,
     &&TARGET_END_ASYNC_FOR,
+    &&TARGET_LOAD_GLOBAL_MODULE,
+    &&TARGET_LOAD_GLOBAL_BUILTIN,
+    &&TARGET_LOAD_METHOD_ADAPTIVE,
     &&TARGET_LOAD_METHOD_CACHED,
     &&TARGET_LOAD_METHOD_CLASS,
+    &&TARGET_STORE_SUBSCR,
+    &&TARGET_DELETE_SUBSCR,
     &&TARGET_LOAD_METHOD_MODULE,
     &&TARGET_LOAD_METHOD_NO_DICT,
     &&TARGET_STORE_ATTR_ADAPTIVE,
-    &&TARGET_STORE_SUBSCR,
-    &&TARGET_DELETE_SUBSCR,
     &&TARGET_STORE_ATTR_INSTANCE_VALUE,
     &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_STORE_ATTR_WITH_HINT,
-    &&TARGET_LOAD_FAST__LOAD_FAST,
-    &&TARGET_STORE_FAST__LOAD_FAST,
-    &&TARGET_LOAD_FAST__LOAD_CONST,
     &&TARGET_GET_ITER,
     &&TARGET_GET_YIELD_FROM_ITER,
     &&TARGET_PRINT_EXPR,
@@ -74,13 +74,13 @@ static void *opcode_targets[256] = {
     &&TARGET_YIELD_FROM,
     &&TARGET_GET_AWAITABLE,
     &&TARGET_LOAD_ASSERTION_ERROR,
+    &&TARGET_LOAD_FAST__LOAD_FAST,
+    &&TARGET_STORE_FAST__LOAD_FAST,
+    &&TARGET_LOAD_FAST__LOAD_CONST,
     &&TARGET_LOAD_CONST__LOAD_FAST,
     &&TARGET_STORE_FAST__STORE_FAST,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
-    &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_LIST_TO_TUPLE,
     &&TARGET_RETURN_VALUE,
     &&TARGET_IMPORT_STAR,
diff --git a/Python/specialize.c b/Python/specialize.c
index 06b0764..130da00 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -125,6 +125,7 @@ _Py_GetSpecializationStats(void) {
     err += add_stat_dict(stats, LOAD_GLOBAL, "load_global");
     err += add_stat_dict(stats, LOAD_METHOD, "load_method");
     err += add_stat_dict(stats, BINARY_SUBSCR, "binary_subscr");
+    err += add_stat_dict(stats, STORE_SUBSCR, "store_subscr");
     err += add_stat_dict(stats, STORE_ATTR, "store_attr");
     err += add_stat_dict(stats, CALL_FUNCTION, "call_function");
     err += add_stat_dict(stats, BINARY_OP, "binary_op");
@@ -182,6 +183,7 @@ _Py_PrintSpecializationStats(void)
     print_stats(out, &_specialization_stats[LOAD_GLOBAL], "load_global");
     print_stats(out, &_specialization_stats[LOAD_METHOD], "load_method");
     print_stats(out, &_specialization_stats[BINARY_SUBSCR], "binary_subscr");
+    print_stats(out, &_specialization_stats[STORE_SUBSCR], "store_subscr");
     print_stats(out, &_specialization_stats[STORE_ATTR], "store_attr");
     print_stats(out, &_specialization_stats[CALL_FUNCTION], "call_function");
     print_stats(out, &_specialization_stats[BINARY_OP], "binary_op");
@@ -233,6 +235,7 @@ static uint8_t adaptive_opcodes[256] = {
     [LOAD_GLOBAL] = LOAD_GLOBAL_ADAPTIVE,
     [LOAD_METHOD] = LOAD_METHOD_ADAPTIVE,
     [BINARY_SUBSCR] = BINARY_SUBSCR_ADAPTIVE,
+    [STORE_SUBSCR] = STORE_SUBSCR_ADAPTIVE,
     [CALL_FUNCTION] = CALL_FUNCTION_ADAPTIVE,
     [STORE_ATTR] = STORE_ATTR_ADAPTIVE,
     [BINARY_OP] = BINARY_OP_ADAPTIVE,
@@ -244,6 +247,7 @@ static uint8_t cache_requirements[256] = {
     [LOAD_GLOBAL] = 2, /* _PyAdaptiveEntry and _PyLoadGlobalCache */
     [LOAD_METHOD] = 3, /* _PyAdaptiveEntry, _PyAttrCache and _PyObjectCache */
     [BINARY_SUBSCR] = 2, /* _PyAdaptiveEntry, _PyObjectCache */
+    [STORE_SUBSCR] = 0,
     [CALL_FUNCTION] = 2, /* _PyAdaptiveEntry and _PyObjectCache/_PyCallCache */
     [STORE_ATTR] = 2, /* _PyAdaptiveEntry and _PyAttrCache */
     [BINARY_OP] = 1,  // _PyAdaptiveEntry
@@ -1228,6 +1232,53 @@ success:
     return 0;
 }
 
+int
+_Py_Specialize_StoreSubscr(PyObject *container, PyObject *sub, _Py_CODEUNIT *instr)
+{
+    PyTypeObject *container_type = Py_TYPE(container);
+    if (container_type == &PyList_Type) {
+        if (PyLong_CheckExact(sub)) {
+            if ((Py_SIZE(sub) == 0 || Py_SIZE(sub) == 1)
+                && ((PyLongObject *)sub)->ob_digit[0] < PyList_GET_SIZE(container))
+            {
+                *instr = _Py_MAKECODEUNIT(STORE_SUBSCR_LIST_INT,
+                                          initial_counter_value());
+                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_LIST_SLICE);
+            goto fail;
+        }
+        else {
+            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+            goto fail;
+        }
+    }
+    else if (container_type == &PyDict_Type) {
+        *instr = _Py_MAKECODEUNIT(STORE_SUBSCR_DICT,
+                                  initial_counter_value());
+         goto success;
+    }
+    else {
+        SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+        goto fail;
+    }
+fail:
+    STAT_INC(STORE_SUBSCR, specialization_failure);
+    assert(!PyErr_Occurred());
+    *instr = _Py_MAKECODEUNIT(_Py_OPCODE(*instr), ADAPTIVE_CACHE_BACKOFF);
+    return 0;
+success:
+    STAT_INC(STORE_SUBSCR, specialization_success);
+    assert(!PyErr_Occurred());
+    return 0;
+}
+
 static int
 specialize_class_call(
     PyObject *callable, _Py_CODEUNIT *instr,