GH-107596: Specialize str[int] (GH-107597)

5 files changed, 94 insertions, 28 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 90e26d3..d6bfb62 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -509,6 +509,7 @@ dummy_func(
             BINARY_SUBSCR_DICT,
             BINARY_SUBSCR_GETITEM,
             BINARY_SUBSCR_LIST_INT,
+            BINARY_SUBSCR_STR_INT,
             BINARY_SUBSCR_TUPLE_INT,
         };
 
@@ -574,6 +575,21 @@ dummy_func(
             Py_DECREF(list);
         }
 
+        inst(BINARY_SUBSCR_STR_INT, (unused/1, str, sub -- res)) {
+            DEOPT_IF(!PyLong_CheckExact(sub), BINARY_SUBSCR);
+            DEOPT_IF(!PyUnicode_CheckExact(str), BINARY_SUBSCR);
+            DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub), BINARY_SUBSCR);
+            Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
+            DEOPT_IF(PyUnicode_GET_LENGTH(str) <= index, BINARY_SUBSCR);
+            // Specialize for reading an ASCII character from any string:
+            Py_UCS4 c = PyUnicode_READ_CHAR(str, index);
+            DEOPT_IF(Py_ARRAY_LENGTH(_Py_SINGLETON(strings).ascii) <= c, BINARY_SUBSCR);
+            STAT_INC(BINARY_SUBSCR, hit);
+            res = (PyObject*)&_Py_SINGLETON(strings).ascii[c];
+            _Py_DECREF_SPECIALIZED(sub, (destructor)PyObject_Free);
+            Py_DECREF(str);
+        }
+
         inst(BINARY_SUBSCR_TUPLE_INT, (unused/1, tuple, sub -- res)) {
             DEOPT_IF(!PyLong_CheckExact(sub), BINARY_SUBSCR);
             DEOPT_IF(!PyTuple_CheckExact(tuple), BINARY_SUBSCR);
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 9363b49..27be8a3 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -462,6 +462,29 @@
             break;
         }
 
+        case BINARY_SUBSCR_STR_INT: {
+            PyObject *sub;
+            PyObject *str;
+            PyObject *res;
+            sub = stack_pointer[-1];
+            str = stack_pointer[-2];
+            DEOPT_IF(!PyLong_CheckExact(sub), BINARY_SUBSCR);
+            DEOPT_IF(!PyUnicode_CheckExact(str), BINARY_SUBSCR);
+            DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub), BINARY_SUBSCR);
+            Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
+            DEOPT_IF(PyUnicode_GET_LENGTH(str) <= index, BINARY_SUBSCR);
+            // Specialize for reading an ASCII character from any string:
+            Py_UCS4 c = PyUnicode_READ_CHAR(str, index);
+            DEOPT_IF(Py_ARRAY_LENGTH(_Py_SINGLETON(strings).ascii) <= c, BINARY_SUBSCR);
+            STAT_INC(BINARY_SUBSCR, hit);
+            res = (PyObject*)&_Py_SINGLETON(strings).ascii[c];
+            _Py_DECREF_SPECIALIZED(sub, (destructor)PyObject_Free);
+            Py_DECREF(str);
+            STACK_SHRINK(1);
+            stack_pointer[-1] = res;
+            break;
+        }
+
         case BINARY_SUBSCR_TUPLE_INT: {
             PyObject *sub;
             PyObject *tuple;
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 7250240..d7db8b0 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -685,6 +685,30 @@
             DISPATCH();
         }
 
+        TARGET(BINARY_SUBSCR_STR_INT) {
+            PyObject *sub;
+            PyObject *str;
+            PyObject *res;
+            sub = stack_pointer[-1];
+            str = stack_pointer[-2];
+            DEOPT_IF(!PyLong_CheckExact(sub), BINARY_SUBSCR);
+            DEOPT_IF(!PyUnicode_CheckExact(str), BINARY_SUBSCR);
+            DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub), BINARY_SUBSCR);
+            Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
+            DEOPT_IF(PyUnicode_GET_LENGTH(str) <= index, BINARY_SUBSCR);
+            // Specialize for reading an ASCII character from any string:
+            Py_UCS4 c = PyUnicode_READ_CHAR(str, index);
+            DEOPT_IF(Py_ARRAY_LENGTH(_Py_SINGLETON(strings).ascii) <= c, BINARY_SUBSCR);
+            STAT_INC(BINARY_SUBSCR, hit);
+            res = (PyObject*)&_Py_SINGLETON(strings).ascii[c];
+            _Py_DECREF_SPECIALIZED(sub, (destructor)PyObject_Free);
+            Py_DECREF(str);
+            STACK_SHRINK(1);
+            stack_pointer[-1] = res;
+            next_instr += 1;
+            DISPATCH();
+        }
+
         TARGET(BINARY_SUBSCR_TUPLE_INT) {
             PyObject *sub;
             PyObject *tuple;
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index d84d253..210c37b 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -42,12 +42,12 @@ static void *opcode_targets[256] = {
     &&TARGET_FORMAT_SIMPLE,
     &&TARGET_FORMAT_WITH_SPEC,
     &&TARGET_BINARY_SUBSCR_LIST_INT,
+    &&TARGET_BINARY_SUBSCR_STR_INT,
     &&TARGET_BINARY_SUBSCR_TUPLE_INT,
     &&TARGET_STORE_SUBSCR_DICT,
     &&TARGET_STORE_SUBSCR_LIST_INT,
     &&TARGET_SEND_GEN,
     &&TARGET_UNPACK_SEQUENCE_TWO_TUPLE,
-    &&TARGET_UNPACK_SEQUENCE_TUPLE,
     &&TARGET_WITH_EXCEPT_START,
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
@@ -55,39 +55,39 @@ static void *opcode_targets[256] = {
     &&TARGET_BEFORE_WITH,
     &&TARGET_END_ASYNC_FOR,
     &&TARGET_CLEANUP_THROW,
+    &&TARGET_UNPACK_SEQUENCE_TUPLE,
     &&TARGET_UNPACK_SEQUENCE_LIST,
     &&TARGET_STORE_ATTR_INSTANCE_VALUE,
     &&TARGET_STORE_ATTR_SLOT,
-    &&TARGET_STORE_ATTR_WITH_HINT,
     &&TARGET_STORE_SUBSCR,
     &&TARGET_DELETE_SUBSCR,
+    &&TARGET_STORE_ATTR_WITH_HINT,
     &&TARGET_LOAD_GLOBAL_MODULE,
     &&TARGET_LOAD_GLOBAL_BUILTIN,
     &&TARGET_LOAD_SUPER_ATTR_ATTR,
     &&TARGET_LOAD_SUPER_ATTR_METHOD,
     &&TARGET_LOAD_ATTR_INSTANCE_VALUE,
-    &&TARGET_LOAD_ATTR_MODULE,
     &&TARGET_GET_ITER,
     &&TARGET_GET_YIELD_FROM_ITER,
-    &&TARGET_LOAD_ATTR_WITH_HINT,
+    &&TARGET_LOAD_ATTR_MODULE,
     &&TARGET_LOAD_BUILD_CLASS,
+    &&TARGET_LOAD_ATTR_WITH_HINT,
     &&TARGET_LOAD_ATTR_SLOT,
-    &&TARGET_LOAD_ATTR_CLASS,
     &&TARGET_LOAD_ASSERTION_ERROR,
     &&TARGET_RETURN_GENERATOR,
+    &&TARGET_LOAD_ATTR_CLASS,
     &&TARGET_LOAD_ATTR_PROPERTY,
     &&TARGET_LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN,
     &&TARGET_LOAD_ATTR_METHOD_WITH_VALUES,
     &&TARGET_LOAD_ATTR_METHOD_NO_DICT,
     &&TARGET_LOAD_ATTR_METHOD_LAZY_DICT,
     &&TARGET_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES,
-    &&TARGET_LOAD_ATTR_NONDESCRIPTOR_NO_DICT,
     &&TARGET_RETURN_VALUE,
-    &&TARGET_COMPARE_OP_FLOAT,
+    &&TARGET_LOAD_ATTR_NONDESCRIPTOR_NO_DICT,
     &&TARGET_SETUP_ANNOTATIONS,
-    &&TARGET_COMPARE_OP_INT,
+    &&TARGET_COMPARE_OP_FLOAT,
     &&TARGET_LOAD_LOCALS,
-    &&TARGET_COMPARE_OP_STR,
+    &&TARGET_COMPARE_OP_INT,
     &&TARGET_POP_EXCEPT,
     &&TARGET_STORE_NAME,
     &&TARGET_DELETE_NAME,
@@ -110,9 +110,9 @@ static void *opcode_targets[256] = {
     &&TARGET_IMPORT_NAME,
     &&TARGET_IMPORT_FROM,
     &&TARGET_JUMP_FORWARD,
+    &&TARGET_COMPARE_OP_STR,
     &&TARGET_FOR_ITER_LIST,
     &&TARGET_FOR_ITER_TUPLE,
-    &&TARGET_FOR_ITER_RANGE,
     &&TARGET_POP_JUMP_IF_FALSE,
     &&TARGET_POP_JUMP_IF_TRUE,
     &&TARGET_LOAD_GLOBAL,
@@ -131,11 +131,11 @@ static void *opcode_targets[256] = {
     &&TARGET_POP_JUMP_IF_NONE,
     &&TARGET_RAISE_VARARGS,
     &&TARGET_GET_AWAITABLE,
-    &&TARGET_FOR_ITER_GEN,
+    &&TARGET_FOR_ITER_RANGE,
     &&TARGET_BUILD_SLICE,
     &&TARGET_JUMP_BACKWARD_NO_INTERRUPT,
     &&TARGET_MAKE_CELL,
-    &&TARGET_CALL_BOUND_METHOD_EXACT_ARGS,
+    &&TARGET_FOR_ITER_GEN,
     &&TARGET_LOAD_DEREF,
     &&TARGET_STORE_DEREF,
     &&TARGET_DELETE_DEREF,
@@ -147,26 +147,26 @@ static void *opcode_targets[256] = {
     &&TARGET_LIST_APPEND,
     &&TARGET_SET_ADD,
     &&TARGET_MAP_ADD,
-    &&TARGET_CALL_PY_EXACT_ARGS,
+    &&TARGET_CALL_BOUND_METHOD_EXACT_ARGS,
     &&TARGET_COPY_FREE_VARS,
     &&TARGET_YIELD_VALUE,
     &&TARGET_RESUME,
     &&TARGET_MATCH_CLASS,
+    &&TARGET_CALL_PY_EXACT_ARGS,
     &&TARGET_CALL_PY_WITH_DEFAULTS,
     &&TARGET_CALL_NO_KW_TYPE_1,
-    &&TARGET_CALL_NO_KW_STR_1,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
     &&TARGET_CONVERT_VALUE,
+    &&TARGET_CALL_NO_KW_STR_1,
     &&TARGET_CALL_NO_KW_TUPLE_1,
     &&TARGET_CALL_BUILTIN_CLASS,
-    &&TARGET_CALL_NO_KW_BUILTIN_O,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_CALL_NO_KW_BUILTIN_O,
     &&TARGET_CALL_NO_KW_BUILTIN_FAST,
-    &&TARGET_CALL_BUILTIN_FAST_WITH_KEYWORDS,
     &&TARGET_LOAD_FAST_LOAD_FAST,
     &&TARGET_STORE_FAST_LOAD_FAST,
     &&TARGET_STORE_FAST_STORE_FAST,
@@ -177,6 +177,7 @@ static void *opcode_targets[256] = {
     &&TARGET_LOAD_FROM_DICT_OR_GLOBALS,
     &&TARGET_LOAD_FROM_DICT_OR_DEREF,
     &&TARGET_SET_FUNCTION_ATTRIBUTE,
+    &&TARGET_CALL_BUILTIN_FAST_WITH_KEYWORDS,
     &&TARGET_CALL_NO_KW_LEN,
     &&TARGET_CALL_NO_KW_ISINSTANCE,
     &&TARGET_CALL_NO_KW_LIST_APPEND,
@@ -228,7 +229,6 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_ENTER_EXECUTOR,
     &&_unknown_opcode,
     &&_unknown_opcode,
diff --git a/Python/specialize.c b/Python/specialize.c
index de329ef..855252e 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -363,7 +363,6 @@ _PyCode_Quicken(PyCodeObject *code)
 #define SPEC_FAIL_SUBSCR_ARRAY_SLICE 10
 #define SPEC_FAIL_SUBSCR_LIST_SLICE 11
 #define SPEC_FAIL_SUBSCR_TUPLE_SLICE 12
-#define SPEC_FAIL_SUBSCR_STRING_INT 13
 #define SPEC_FAIL_SUBSCR_STRING_SLICE 14
 #define SPEC_FAIL_SUBSCR_BUFFER_INT 15
 #define SPEC_FAIL_SUBSCR_BUFFER_SLICE 16
@@ -1260,16 +1259,7 @@ success:
 static int
 binary_subscr_fail_kind(PyTypeObject *container_type, PyObject *sub)
 {
-    if (container_type == &PyUnicode_Type) {
-        if (PyLong_CheckExact(sub)) {
-            return SPEC_FAIL_SUBSCR_STRING_INT;
-        }
-        if (PySlice_Check(sub)) {
-            return SPEC_FAIL_SUBSCR_STRING_SLICE;
-        }
-        return SPEC_FAIL_OTHER;
-    }
-    else if (strcmp(container_type->tp_name, "array.array") == 0) {
+    if (strcmp(container_type->tp_name, "array.array") == 0) {
         if (PyLong_CheckExact(sub)) {
             return SPEC_FAIL_SUBSCR_ARRAY_INT;
         }
@@ -1376,6 +1366,19 @@ _Py_Specialize_BinarySubscr(
             PySlice_Check(sub) ? SPEC_FAIL_SUBSCR_TUPLE_SLICE : SPEC_FAIL_OTHER);
         goto fail;
     }
+    if (container_type == &PyUnicode_Type) {
+        if (PyLong_CheckExact(sub)) {
+            if (_PyLong_IsNonNegativeCompact((PyLongObject *)sub)) {
+                instr->op.code = BINARY_SUBSCR_STR_INT;
+                goto success;
+            }
+            SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
+            goto fail;
+        }
+        SPECIALIZATION_FAIL(BINARY_SUBSCR,
+            PySlice_Check(sub) ? SPEC_FAIL_SUBSCR_STRING_SLICE : SPEC_FAIL_OTHER);
+        goto fail;
+    }
     if (container_type == &PyDict_Type) {
         instr->op.code = BINARY_SUBSCR_DICT;
         goto success;