gh-92777: Add LOAD_METHOD_LAZY_DICT (GH-92778)

7 files changed, 85 insertions, 49 deletions

diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
index 79926b6..c693fe3 100644
--- a/Include/internal/pycore_opcode.h
+++ b/Include/internal/pycore_opcode.h
@@ -167,6 +167,7 @@ const uint8_t _PyOpcode_Deopt[256] = {
     [LOAD_METHOD] = LOAD_METHOD,
     [LOAD_METHOD_ADAPTIVE] = LOAD_METHOD,
     [LOAD_METHOD_CLASS] = LOAD_METHOD,
+    [LOAD_METHOD_LAZY_DICT] = LOAD_METHOD,
     [LOAD_METHOD_MODULE] = LOAD_METHOD,
     [LOAD_METHOD_NO_DICT] = LOAD_METHOD,
     [LOAD_METHOD_WITH_DICT] = LOAD_METHOD,
@@ -348,6 +349,7 @@ const uint8_t _PyOpcode_Original[256] = {
     [LOAD_METHOD] = LOAD_METHOD,
     [LOAD_METHOD_ADAPTIVE] = LOAD_METHOD,
     [LOAD_METHOD_CLASS] = LOAD_METHOD,
+    [LOAD_METHOD_LAZY_DICT] = LOAD_METHOD,
     [LOAD_METHOD_MODULE] = LOAD_METHOD,
     [LOAD_METHOD_NO_DICT] = LOAD_METHOD,
     [LOAD_METHOD_WITH_DICT] = LOAD_METHOD,
@@ -499,12 +501,12 @@ static const char *const _PyOpcode_OpName[256] = {
     [LOAD_GLOBAL_MODULE] = "LOAD_GLOBAL_MODULE",
     [LOAD_METHOD_ADAPTIVE] = "LOAD_METHOD_ADAPTIVE",
     [LOAD_METHOD_CLASS] = "LOAD_METHOD_CLASS",
-    [LOAD_METHOD_MODULE] = "LOAD_METHOD_MODULE",
+    [LOAD_METHOD_LAZY_DICT] = "LOAD_METHOD_LAZY_DICT",
     [LIST_TO_TUPLE] = "LIST_TO_TUPLE",
     [RETURN_VALUE] = "RETURN_VALUE",
     [IMPORT_STAR] = "IMPORT_STAR",
     [SETUP_ANNOTATIONS] = "SETUP_ANNOTATIONS",
-    [LOAD_METHOD_NO_DICT] = "LOAD_METHOD_NO_DICT",
+    [LOAD_METHOD_MODULE] = "LOAD_METHOD_MODULE",
     [ASYNC_GEN_WRAP] = "ASYNC_GEN_WRAP",
     [PREP_RERAISE_STAR] = "PREP_RERAISE_STAR",
     [POP_EXCEPT] = "POP_EXCEPT",
@@ -531,7 +533,7 @@ static const char *const _PyOpcode_OpName[256] = {
     [JUMP_FORWARD] = "JUMP_FORWARD",
     [JUMP_IF_FALSE_OR_POP] = "JUMP_IF_FALSE_OR_POP",
     [JUMP_IF_TRUE_OR_POP] = "JUMP_IF_TRUE_OR_POP",
-    [LOAD_METHOD_WITH_DICT] = "LOAD_METHOD_WITH_DICT",
+    [LOAD_METHOD_NO_DICT] = "LOAD_METHOD_NO_DICT",
     [POP_JUMP_FORWARD_IF_FALSE] = "POP_JUMP_FORWARD_IF_FALSE",
     [POP_JUMP_FORWARD_IF_TRUE] = "POP_JUMP_FORWARD_IF_TRUE",
     [LOAD_GLOBAL] = "LOAD_GLOBAL",
@@ -539,13 +541,13 @@ static const char *const _PyOpcode_OpName[256] = {
     [CONTAINS_OP] = "CONTAINS_OP",
     [RERAISE] = "RERAISE",
     [COPY] = "COPY",
-    [LOAD_METHOD_WITH_VALUES] = "LOAD_METHOD_WITH_VALUES",
+    [LOAD_METHOD_WITH_DICT] = "LOAD_METHOD_WITH_DICT",
     [BINARY_OP] = "BINARY_OP",
     [SEND] = "SEND",
     [LOAD_FAST] = "LOAD_FAST",
     [STORE_FAST] = "STORE_FAST",
     [DELETE_FAST] = "DELETE_FAST",
-    [RESUME_QUICK] = "RESUME_QUICK",
+    [LOAD_METHOD_WITH_VALUES] = "LOAD_METHOD_WITH_VALUES",
     [POP_JUMP_FORWARD_IF_NOT_NONE] = "POP_JUMP_FORWARD_IF_NOT_NONE",
     [POP_JUMP_FORWARD_IF_NONE] = "POP_JUMP_FORWARD_IF_NONE",
     [RAISE_VARARGS] = "RAISE_VARARGS",
@@ -559,9 +561,9 @@ static const char *const _PyOpcode_OpName[256] = {
     [STORE_DEREF] = "STORE_DEREF",
     [DELETE_DEREF] = "DELETE_DEREF",
     [JUMP_BACKWARD] = "JUMP_BACKWARD",
-    [STORE_ATTR_ADAPTIVE] = "STORE_ATTR_ADAPTIVE",
+    [RESUME_QUICK] = "RESUME_QUICK",
     [CALL_FUNCTION_EX] = "CALL_FUNCTION_EX",
-    [STORE_ATTR_INSTANCE_VALUE] = "STORE_ATTR_INSTANCE_VALUE",
+    [STORE_ATTR_ADAPTIVE] = "STORE_ATTR_ADAPTIVE",
     [EXTENDED_ARG] = "EXTENDED_ARG",
     [LIST_APPEND] = "LIST_APPEND",
     [SET_ADD] = "SET_ADD",
@@ -571,32 +573,32 @@ static const char *const _PyOpcode_OpName[256] = {
     [YIELD_VALUE] = "YIELD_VALUE",
     [RESUME] = "RESUME",
     [MATCH_CLASS] = "MATCH_CLASS",
+    [STORE_ATTR_INSTANCE_VALUE] = "STORE_ATTR_INSTANCE_VALUE",
     [STORE_ATTR_SLOT] = "STORE_ATTR_SLOT",
-    [STORE_ATTR_WITH_HINT] = "STORE_ATTR_WITH_HINT",
     [FORMAT_VALUE] = "FORMAT_VALUE",
     [BUILD_CONST_KEY_MAP] = "BUILD_CONST_KEY_MAP",
     [BUILD_STRING] = "BUILD_STRING",
+    [STORE_ATTR_WITH_HINT] = "STORE_ATTR_WITH_HINT",
     [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
-    [STORE_FAST__STORE_FAST] = "STORE_FAST__STORE_FAST",
     [LOAD_METHOD] = "LOAD_METHOD",
-    [STORE_SUBSCR_ADAPTIVE] = "STORE_SUBSCR_ADAPTIVE",
+    [STORE_FAST__STORE_FAST] = "STORE_FAST__STORE_FAST",
     [LIST_EXTEND] = "LIST_EXTEND",
     [SET_UPDATE] = "SET_UPDATE",
     [DICT_MERGE] = "DICT_MERGE",
     [DICT_UPDATE] = "DICT_UPDATE",
+    [STORE_SUBSCR_ADAPTIVE] = "STORE_SUBSCR_ADAPTIVE",
     [STORE_SUBSCR_DICT] = "STORE_SUBSCR_DICT",
     [STORE_SUBSCR_LIST_INT] = "STORE_SUBSCR_LIST_INT",
     [UNPACK_SEQUENCE_ADAPTIVE] = "UNPACK_SEQUENCE_ADAPTIVE",
     [UNPACK_SEQUENCE_LIST] = "UNPACK_SEQUENCE_LIST",
-    [UNPACK_SEQUENCE_TUPLE] = "UNPACK_SEQUENCE_TUPLE",
     [CALL] = "CALL",
     [KW_NAMES] = "KW_NAMES",
     [POP_JUMP_BACKWARD_IF_NOT_NONE] = "POP_JUMP_BACKWARD_IF_NOT_NONE",
     [POP_JUMP_BACKWARD_IF_NONE] = "POP_JUMP_BACKWARD_IF_NONE",
     [POP_JUMP_BACKWARD_IF_FALSE] = "POP_JUMP_BACKWARD_IF_FALSE",
     [POP_JUMP_BACKWARD_IF_TRUE] = "POP_JUMP_BACKWARD_IF_TRUE",
+    [UNPACK_SEQUENCE_TUPLE] = "UNPACK_SEQUENCE_TUPLE",
     [UNPACK_SEQUENCE_TWO_TUPLE] = "UNPACK_SEQUENCE_TWO_TUPLE",
-    [178] = "<178>",
     [179] = "<179>",
     [180] = "<180>",
     [181] = "<181>",
@@ -678,7 +680,6 @@ static const char *const _PyOpcode_OpName[256] = {
 #endif
 
 #define EXTRA_CASES \
-    case 178: \
     case 179: \
     case 180: \
     case 181: \
diff --git a/Include/opcode.h b/Include/opcode.h
index e04b5a6..f76ca94 100644
--- a/Include/opcode.h
+++ b/Include/opcode.h
@@ -169,24 +169,25 @@ extern "C" {
 #define LOAD_GLOBAL_MODULE                      78
 #define LOAD_METHOD_ADAPTIVE                    79
 #define LOAD_METHOD_CLASS                       80
-#define LOAD_METHOD_MODULE                      81
-#define LOAD_METHOD_NO_DICT                     86
-#define LOAD_METHOD_WITH_DICT                  113
-#define LOAD_METHOD_WITH_VALUES                121
-#define RESUME_QUICK                           127
-#define STORE_ATTR_ADAPTIVE                    141
-#define STORE_ATTR_INSTANCE_VALUE              143
-#define STORE_ATTR_SLOT                        153
-#define STORE_ATTR_WITH_HINT                   154
-#define STORE_FAST__LOAD_FAST                  158
-#define STORE_FAST__STORE_FAST                 159
-#define STORE_SUBSCR_ADAPTIVE                  161
-#define STORE_SUBSCR_DICT                      166
-#define STORE_SUBSCR_LIST_INT                  167
-#define UNPACK_SEQUENCE_ADAPTIVE               168
-#define UNPACK_SEQUENCE_LIST                   169
-#define UNPACK_SEQUENCE_TUPLE                  170
-#define UNPACK_SEQUENCE_TWO_TUPLE              177
+#define LOAD_METHOD_LAZY_DICT                   81
+#define LOAD_METHOD_MODULE                      86
+#define LOAD_METHOD_NO_DICT                    113
+#define LOAD_METHOD_WITH_DICT                  121
+#define LOAD_METHOD_WITH_VALUES                127
+#define RESUME_QUICK                           141
+#define STORE_ATTR_ADAPTIVE                    143
+#define STORE_ATTR_INSTANCE_VALUE              153
+#define STORE_ATTR_SLOT                        154
+#define STORE_ATTR_WITH_HINT                   158
+#define STORE_FAST__LOAD_FAST                  159
+#define STORE_FAST__STORE_FAST                 161
+#define STORE_SUBSCR_ADAPTIVE                  166
+#define STORE_SUBSCR_DICT                      167
+#define STORE_SUBSCR_LIST_INT                  168
+#define UNPACK_SEQUENCE_ADAPTIVE               169
+#define UNPACK_SEQUENCE_LIST                   170
+#define UNPACK_SEQUENCE_TUPLE                  177
+#define UNPACK_SEQUENCE_TWO_TUPLE              178
 #define DO_TRACING                             255
 
 #define HAS_CONST(op) (false\
diff --git a/Lib/opcode.py b/Lib/opcode.py
index 410853a..573e461 100644
--- a/Lib/opcode.py
+++ b/Lib/opcode.py
@@ -303,6 +303,7 @@ _specializations = {
     "LOAD_METHOD": [
         "LOAD_METHOD_ADAPTIVE",
         "LOAD_METHOD_CLASS",
+        "LOAD_METHOD_LAZY_DICT",
         "LOAD_METHOD_MODULE",
         "LOAD_METHOD_NO_DICT",
         "LOAD_METHOD_WITH_DICT",
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-05-13-12-36-10.gh-issue-92777.Odo4vP.rst b/Misc/NEWS.d/next/Core and Builtins/2022-05-13-12-36-10.gh-issue-92777.Odo4vP.rst
new file mode 100644
index 0000000..1941659
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-05-13-12-36-10.gh-issue-92777.Odo4vP.rst
@@ -0,0 +1 @@
+Specialize ``LOAD_METHOD`` for objects with lazy dictionaries. Patch by Ken Jin.
diff --git a/Python/ceval.c b/Python/ceval.c
index e82f734..56b1017 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -4673,6 +4673,29 @@ handle_eval_breaker:
             NOTRACE_DISPATCH();
         }
 
+        TARGET(LOAD_METHOD_LAZY_DICT) {
+            assert(cframe.use_tracing == 0);
+            PyObject *self = TOP();
+            PyTypeObject *self_cls = Py_TYPE(self);
+            _PyLoadMethodCache *cache = (_PyLoadMethodCache *)next_instr;
+            uint32_t type_version = read_u32(cache->type_version);
+            DEOPT_IF(self_cls->tp_version_tag != type_version, LOAD_METHOD);
+            int dictoffset = cache->dict_offset;
+            PyObject *dict = *(PyObject **)((char *)self + dictoffset);
+            assert(dictoffset == self_cls->tp_dictoffset && dictoffset > 0);
+            /* This object has a __dict__, just not yet created */
+            DEOPT_IF(dict != NULL, LOAD_METHOD);
+            STAT_INC(LOAD_METHOD, hit);
+            PyObject *res = read_obj(cache->descr);
+            assert(res != NULL);
+            assert(_PyType_HasFeature(Py_TYPE(res), Py_TPFLAGS_METHOD_DESCRIPTOR));
+            Py_INCREF(res);
+            SET_TOP(res);
+            PUSH(self);
+            JUMPBY(INLINE_CACHE_ENTRIES_LOAD_METHOD);
+            NOTRACE_DISPATCH();
+        }
+
         TARGET(LOAD_METHOD_MODULE) {
             /* LOAD_METHOD, for module methods */
             assert(cframe.use_tracing == 0);
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 0d8faf2..71b7a76 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -80,12 +80,12 @@ static void *opcode_targets[256] = {
     &&TARGET_LOAD_GLOBAL_MODULE,
     &&TARGET_LOAD_METHOD_ADAPTIVE,
     &&TARGET_LOAD_METHOD_CLASS,
-    &&TARGET_LOAD_METHOD_MODULE,
+    &&TARGET_LOAD_METHOD_LAZY_DICT,
     &&TARGET_LIST_TO_TUPLE,
     &&TARGET_RETURN_VALUE,
     &&TARGET_IMPORT_STAR,
     &&TARGET_SETUP_ANNOTATIONS,
-    &&TARGET_LOAD_METHOD_NO_DICT,
+    &&TARGET_LOAD_METHOD_MODULE,
     &&TARGET_ASYNC_GEN_WRAP,
     &&TARGET_PREP_RERAISE_STAR,
     &&TARGET_POP_EXCEPT,
@@ -112,7 +112,7 @@ static void *opcode_targets[256] = {
     &&TARGET_JUMP_FORWARD,
     &&TARGET_JUMP_IF_FALSE_OR_POP,
     &&TARGET_JUMP_IF_TRUE_OR_POP,
-    &&TARGET_LOAD_METHOD_WITH_DICT,
+    &&TARGET_LOAD_METHOD_NO_DICT,
     &&TARGET_POP_JUMP_FORWARD_IF_FALSE,
     &&TARGET_POP_JUMP_FORWARD_IF_TRUE,
     &&TARGET_LOAD_GLOBAL,
@@ -120,13 +120,13 @@ static void *opcode_targets[256] = {
     &&TARGET_CONTAINS_OP,
     &&TARGET_RERAISE,
     &&TARGET_COPY,
-    &&TARGET_LOAD_METHOD_WITH_VALUES,
+    &&TARGET_LOAD_METHOD_WITH_DICT,
     &&TARGET_BINARY_OP,
     &&TARGET_SEND,
     &&TARGET_LOAD_FAST,
     &&TARGET_STORE_FAST,
     &&TARGET_DELETE_FAST,
-    &&TARGET_RESUME_QUICK,
+    &&TARGET_LOAD_METHOD_WITH_VALUES,
     &&TARGET_POP_JUMP_FORWARD_IF_NOT_NONE,
     &&TARGET_POP_JUMP_FORWARD_IF_NONE,
     &&TARGET_RAISE_VARARGS,
@@ -140,9 +140,9 @@ static void *opcode_targets[256] = {
     &&TARGET_STORE_DEREF,
     &&TARGET_DELETE_DEREF,
     &&TARGET_JUMP_BACKWARD,
-    &&TARGET_STORE_ATTR_ADAPTIVE,
+    &&TARGET_RESUME_QUICK,
     &&TARGET_CALL_FUNCTION_EX,
-    &&TARGET_STORE_ATTR_INSTANCE_VALUE,
+    &&TARGET_STORE_ATTR_ADAPTIVE,
     &&TARGET_EXTENDED_ARG,
     &&TARGET_LIST_APPEND,
     &&TARGET_SET_ADD,
@@ -152,30 +152,31 @@ static void *opcode_targets[256] = {
     &&TARGET_YIELD_VALUE,
     &&TARGET_RESUME,
     &&TARGET_MATCH_CLASS,
+    &&TARGET_STORE_ATTR_INSTANCE_VALUE,
     &&TARGET_STORE_ATTR_SLOT,
-    &&TARGET_STORE_ATTR_WITH_HINT,
     &&TARGET_FORMAT_VALUE,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
+    &&TARGET_STORE_ATTR_WITH_HINT,
     &&TARGET_STORE_FAST__LOAD_FAST,
-    &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_LOAD_METHOD,
-    &&TARGET_STORE_SUBSCR_ADAPTIVE,
+    &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_STORE_SUBSCR_ADAPTIVE,
     &&TARGET_STORE_SUBSCR_DICT,
     &&TARGET_STORE_SUBSCR_LIST_INT,
     &&TARGET_UNPACK_SEQUENCE_ADAPTIVE,
     &&TARGET_UNPACK_SEQUENCE_LIST,
-    &&TARGET_UNPACK_SEQUENCE_TUPLE,
     &&TARGET_CALL,
     &&TARGET_KW_NAMES,
     &&TARGET_POP_JUMP_BACKWARD_IF_NOT_NONE,
     &&TARGET_POP_JUMP_BACKWARD_IF_NONE,
     &&TARGET_POP_JUMP_BACKWARD_IF_FALSE,
     &&TARGET_POP_JUMP_BACKWARD_IF_TRUE,
+    &&TARGET_UNPACK_SEQUENCE_TUPLE,
     &&TARGET_UNPACK_SEQUENCE_TWO_TUPLE,
     &&_unknown_opcode,
     &&_unknown_opcode,
@@ -253,6 +254,5 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_DO_TRACING
 };
diff --git a/Python/specialize.c b/Python/specialize.c
index 9579369..7148b5d 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -903,7 +903,8 @@ typedef enum {
     MANAGED_VALUES = 1,
     MANAGED_DICT = 2,
     OFFSET_DICT = 3,
-    NO_DICT = 4
+    NO_DICT = 4,
+    LAZY_DICT = 5,    
 } ObjectDictKind;
 
 // Please collect stats carefully before and after modifying. A subtle change
@@ -972,14 +973,17 @@ _Py_Specialize_LoadMethod(PyObject *owner, _Py_CODEUNIT *instr, PyObject *name)
         else {
             PyObject *dict = *(PyObject **) ((char *)owner + dictoffset);
             if (dict == NULL) {
-                SPECIALIZATION_FAIL(LOAD_METHOD, SPEC_FAIL_NO_DICT);
-                goto fail;
+                // This object will have a dict if user access __dict__
+                dictkind = LAZY_DICT;
+                keys = NULL;
+            }
+            else {
+                keys = ((PyDictObject *)dict)->ma_keys;
+                dictkind = OFFSET_DICT;
             }
-            keys = ((PyDictObject *)dict)->ma_keys;
-            dictkind = OFFSET_DICT;
         }
     }
-    if (dictkind != NO_DICT) {
+    if (dictkind == MANAGED_VALUES || dictkind == MANAGED_DICT || dictkind == OFFSET_DICT) {
         Py_ssize_t index = _PyDictKeys_StringLookup(keys, name);
         if (index != DKIX_EMPTY) {
             SPECIALIZATION_FAIL(LOAD_METHOD, SPEC_FAIL_LOAD_METHOD_IS_ATTR);
@@ -1008,6 +1012,11 @@ _Py_Specialize_LoadMethod(PyObject *owner, _Py_CODEUNIT *instr, PyObject *name)
             cache->dict_offset = (uint16_t)owner_cls->tp_dictoffset;
             _Py_SET_OPCODE(*instr, LOAD_METHOD_WITH_DICT);
             break;
+        case LAZY_DICT:
+            assert(owner_cls->tp_dictoffset > 0 && owner_cls->tp_dictoffset <= INT16_MAX);
+            cache->dict_offset = (uint16_t)owner_cls->tp_dictoffset;
+            _Py_SET_OPCODE(*instr, LOAD_METHOD_LAZY_DICT);
+            break;
     }
     /* `descr` is borrowed. This is safe for methods (even inherited ones from
     *  super classes!) as long as tp_version_tag is validated for two main reasons: