[Python-checkins] bpo-45609: More specialization stats for STORE_SUBSCR (GH-30193)

[markshannon]

https://github.com/python/cpython/commit/7537f6008704b20e2d04a7ef1c0cfa34121cc5eb
commit: 7537f6008704b20e2d04a7ef1c0cfa34121cc5eb
branch: main
author: Dennis Sweeney <36520290+sweeneyde at users.noreply.github.com>
committer: markshannon <mark at hotpy.org>
date: 2022-01-04T18:05:09Z
summary:

bpo-45609: More specialization stats for STORE_SUBSCR (GH-30193)

files:
M Python/specialize.c
M Tools/scripts/summarize_stats.py

diff --git a/Python/specialize.c b/Python/specialize.c
index 8991fa94f8e36..2da9e0f29b7a4 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -486,6 +486,13 @@ initial_counter_value(void) {
 #define SPEC_FAIL_BUFFER_SLICE 16
 #define SPEC_FAIL_SEQUENCE_INT 17
 
+/* Store subscr */
+#define SPEC_FAIL_BYTEARRAY_INT 18
+#define SPEC_FAIL_BYTEARRAY_SLICE 19
+#define SPEC_FAIL_PY_SIMPLE 20
+#define SPEC_FAIL_PY_OTHER 21
+#define SPEC_FAIL_DICT_SUBCLASS_NO_OVERRIDE 22
+
 /* Binary add */
 
 #define SPEC_FAIL_NON_FUNCTION_SCOPE 11
@@ -1253,15 +1260,73 @@ _Py_Specialize_StoreSubscr(PyObject *container, PyObject *sub, _Py_CODEUNIT *ins
             goto fail;
         }
     }
-    else if (container_type == &PyDict_Type) {
+    if (container_type == &PyDict_Type) {
         *instr = _Py_MAKECODEUNIT(STORE_SUBSCR_DICT,
                                   initial_counter_value());
          goto success;
     }
-    else {
-        SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+#ifdef Py_STATS
+    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_DICT_SUBCLASS_NO_OVERRIDE);
         goto fail;
     }
+    if (PyObject_CheckBuffer(container)) {
+        if (PyLong_CheckExact(sub) && (((size_t)Py_SIZE(sub)) > 1)) {
+            SPECIALIZATION_FAIL(STORE_SUBSCR, 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_ARRAY_INT);
+            }
+            else if (PySlice_Check(sub)) {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_ARRAY_SLICE);
+            }
+            else {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+            }
+        }
+        else if (PyByteArray_CheckExact(container)) {
+            if (PyLong_CheckExact(sub)) {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_BYTEARRAY_INT);
+            }
+            else if (PySlice_Check(sub)) {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_BYTEARRAY_SLICE);
+            }
+            else {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+            }
+        }
+        else {
+            if (PyLong_CheckExact(sub)) {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_BUFFER_INT);
+            }
+            else if (PySlice_Check(sub)) {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_BUFFER_SLICE);
+            }
+            else {
+                SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
+            }
+        }
+        goto fail;
+    }
+    _Py_IDENTIFIER(__setitem__);
+    PyObject *descriptor = _PyType_LookupId(container_type, &PyId___setitem__);
+    if (descriptor && Py_TYPE(descriptor) == &PyFunction_Type) {
+        PyFunctionObject *func = (PyFunctionObject *)descriptor;
+        PyCodeObject *code = (PyCodeObject *)func->func_code;
+        int kind = function_kind(code);
+        if (kind == SIMPLE_FUNCTION) {
+            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_PY_SIMPLE);
+        }
+        else {
+            SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_PY_OTHER);
+        }
+        goto fail;
+    }
+#endif
+    SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
 fail:
     STAT_INC(STORE_SUBSCR, failure);
     assert(!PyErr_Occurred());
diff --git a/Tools/scripts/summarize_stats.py b/Tools/scripts/summarize_stats.py
index a5a8e93c17392..3a77125035a32 100644
--- a/Tools/scripts/summarize_stats.py
+++ b/Tools/scripts/summarize_stats.py
@@ -28,7 +28,7 @@
 def print_specialization_stats(name, family_stats):
     if "specialization.deferred" not in family_stats:
         return
-    total = sum(family_stats[kind] for kind in TOTAL)
+    total = sum(family_stats.get(kind, 0) for kind in TOTAL)
     if total == 0:
         return
     print(name+":")
@@ -44,7 +44,7 @@ def print_specialization_stats(name, family_stats):
     for key in ("specialization.success",  "specialization.failure"):
         label = key[len("specialization."):]
         print(f"  {label}:{family_stats.get(key, 0):>12}")
-    total_failures = family_stats["specialization.failure"]
+    total_failures = family_stats.get("specialization.failure", 0)
     failure_kinds = [ 0 ] * 30
     for key in family_stats:
         if not key.startswith("specialization.failure_kind"):