1 files changed, 61 insertions, 25 deletions

diff --git a/Python/specialize.c b/Python/specialize.c
index dd15de7..06b0764 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -243,7 +243,7 @@ static uint8_t cache_requirements[256] = {
     [LOAD_ATTR] = 2, /* _PyAdaptiveEntry and _PyAttrCache */
     [LOAD_GLOBAL] = 2, /* _PyAdaptiveEntry and _PyLoadGlobalCache */
     [LOAD_METHOD] = 3, /* _PyAdaptiveEntry, _PyAttrCache and _PyObjectCache */
-    [BINARY_SUBSCR] = 0,
+    [BINARY_SUBSCR] = 2, /* _PyAdaptiveEntry, _PyObjectCache */
     [CALL_FUNCTION] = 2, /* _PyAdaptiveEntry and _PyObjectCache/_PyCallCache */
     [STORE_ATTR] = 2, /* _PyAdaptiveEntry and _PyAttrCache */
     [BINARY_OP] = 1,  // _PyAdaptiveEntry
@@ -1100,7 +1100,7 @@ success:
 
 #if COLLECT_SPECIALIZATION_STATS_DETAILED
 static int
-binary_subscr_faiL_kind(PyTypeObject *container_type, PyObject *sub)
+binary_subscr_fail_kind(PyTypeObject *container_type, PyObject *sub)
 {
     if (container_type == &PyUnicode_Type) {
         if (PyLong_CheckExact(sub)) {
@@ -1138,14 +1138,37 @@ binary_subscr_faiL_kind(PyTypeObject *container_type, PyObject *sub)
 }
 #endif
 
+_Py_IDENTIFIER(__getitem__);
+
+#define SIMPLE_FUNCTION 0
+
+static int
+function_kind(PyCodeObject *code) {
+    int flags = code->co_flags;
+    if (flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
+        return SPEC_FAIL_GENERATOR;
+    }
+    if ((flags & (CO_VARKEYWORDS | CO_VARARGS)) || code->co_kwonlyargcount) {
+        return SPEC_FAIL_COMPLEX_PARAMETERS;
+    }
+    if ((flags & CO_OPTIMIZED) == 0) {
+        return SPEC_FAIL_CO_NOT_OPTIMIZED;
+    }
+    if (code->co_nfreevars) {
+        return SPEC_FAIL_FREE_VARS;
+    }
+    return SIMPLE_FUNCTION;
+}
+
 int
 _Py_Specialize_BinarySubscr(
-     PyObject *container, PyObject *sub, _Py_CODEUNIT *instr)
+     PyObject *container, PyObject *sub, _Py_CODEUNIT *instr, SpecializedCacheEntry *cache)
 {
+    _PyAdaptiveEntry *cache0 = &cache->adaptive;
     PyTypeObject *container_type = Py_TYPE(container);
     if (container_type == &PyList_Type) {
         if (PyLong_CheckExact(sub)) {
-            *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_LIST_INT, initial_counter_value());
+            *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_LIST_INT, _Py_OPARG(*instr));
             goto success;
         }
         SPECIALIZATION_FAIL(BINARY_SUBSCR,
@@ -1154,7 +1177,7 @@ _Py_Specialize_BinarySubscr(
     }
     if (container_type == &PyTuple_Type) {
         if (PyLong_CheckExact(sub)) {
-            *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_TUPLE_INT, initial_counter_value());
+            *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_TUPLE_INT, _Py_OPARG(*instr));
             goto success;
         }
         SPECIALIZATION_FAIL(BINARY_SUBSCR,
@@ -1162,20 +1185,46 @@ _Py_Specialize_BinarySubscr(
         goto fail;
     }
     if (container_type == &PyDict_Type) {
-        *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_DICT, initial_counter_value());
+        *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_DICT, _Py_OPARG(*instr));
+        goto success;
+    }
+    PyTypeObject *cls = Py_TYPE(container);
+    PyObject *descriptor = _PyType_LookupId(cls, &PyId___getitem__);
+    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(BINARY_SUBSCR, kind);
+            goto fail;
+        }
+        if (code->co_argcount != 2) {
+            SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_WRONG_NUMBER_ARGUMENTS);
+            goto fail;
+        }
+        assert(cls->tp_version_tag != 0);
+        cache0->version = cls->tp_version_tag;
+        int version = _PyFunction_GetVersionForCurrentState(func);
+        if (version == 0) {
+            SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_OUT_OF_VERSIONS);
+            goto fail;
+        }
+        cache0->index = version;
+        cache[-1].obj.obj = descriptor;
+        *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_GETITEM, _Py_OPARG(*instr));
         goto success;
     }
     SPECIALIZATION_FAIL(BINARY_SUBSCR,
-                        binary_subscr_faiL_kind(container_type, sub));
-    goto fail;
+                        binary_subscr_fail_kind(container_type, sub));
 fail:
     STAT_INC(BINARY_SUBSCR, specialization_failure);
     assert(!PyErr_Occurred());
-    *instr = _Py_MAKECODEUNIT(_Py_OPCODE(*instr), ADAPTIVE_CACHE_BACKOFF);
+    cache_backoff(cache0);
     return 0;
 success:
     STAT_INC(BINARY_SUBSCR, specialization_success);
     assert(!PyErr_Occurred());
+    cache0->counter = initial_counter_value();
     return 0;
 }
 
@@ -1194,23 +1243,10 @@ specialize_py_call(
     int nargs, SpecializedCacheEntry *cache)
 {
     _PyCallCache *cache1 = &cache[-1].call;
-    /* Exclude generator or coroutines for now */
     PyCodeObject *code = (PyCodeObject *)func->func_code;
-    int flags = code->co_flags;
-    if (flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
-        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_GENERATOR);
-        return -1;
-    }
-    if ((flags & (CO_VARKEYWORDS | CO_VARARGS)) || code->co_kwonlyargcount) {
-        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_COMPLEX_PARAMETERS);
-        return -1;
-    }
-    if ((flags & CO_OPTIMIZED) == 0) {
-        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_CO_NOT_OPTIMIZED);
-        return -1;
-    }
-    if (code->co_nfreevars) {
-        SPECIALIZATION_FAIL(CALL_FUNCTION, SPEC_FAIL_FREE_VARS);
+    int kind = function_kind(code);
+    if (kind != SIMPLE_FUNCTION) {
+        SPECIALIZATION_FAIL(CALL_FUNCTION, kind);
         return -1;
     }
     int argcount = code->co_argcount;