gh-122712: Test CALL_ALLOC_AND_ENTER_INIT handles reassignment of __code__

Include/internal/pycore_code.h

@@ -346,6 +346,27 @@ extern void _Py_Specialize_Send(_PyStackRef receiver, _Py_CODEUNIT *instr);
 extern void _Py_Specialize_ToBool(_PyStackRef value, _Py_CODEUNIT *instr);
 extern void _Py_Specialize_ContainsOp(_PyStackRef value, _Py_CODEUNIT *instr);
 
+/* Helpers shared by the specializer and instructions */
+
+typedef enum {
+    _Py_Specialize_Ok = 0,
+    _Py_Specialize_ErrComplexParams = 1,
+    _Py_Specialize_ErrCodeUnoptimized = 2,
+} _Py_Specialize_CheckCodeResult;
+
+static inline _Py_Specialize_CheckCodeResult
+_Py_Specialize_CheckCode(PyCodeObject *code)
+{
+    int flags = code->co_flags;
+    if ((flags & (CO_VARKEYWORDS | CO_VARARGS)) || code->co_kwonlyargcount) {
+        return _Py_Specialize_ErrComplexParams;
+    }
+    if ((flags & CO_OPTIMIZED) == 0) {
+        return _Py_Specialize_ErrCodeUnoptimized;
+    }
+    return _Py_Specialize_Ok;
+}
+
 #ifdef Py_STATS
 
 #include "pycore_bitutils.h"  // _Py_bit_length

Lib/test/test_opcache.py

@@ -28,6 +28,13 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+class TestBase(unittest.TestCase):
+    def assert_specialized(self, f, opname):
+        instructions = dis.get_instructions(f, adaptive=True)
+        opnames = {instruction.opname for instruction in instructions}
+        self.assertIn(opname, opnames)
+
+
 class TestLoadSuperAttrCache(unittest.TestCase):
     def test_descriptor_not_double_executed_on_spec_fail(self):
         calls = []
@@ -479,7 +486,7 @@ def f():
             self.assertFalse(f())
 
 
-class TestCallCache(unittest.TestCase):
+class TestCallCache(TestBase):
     def test_too_many_defaults_0(self):
         def f():
             pass
@@ -507,22 +514,40 @@ def f(x, y):
             f(None)
             f()
 
+    @disabling_optimizer
+    @requires_specialization
+    def test_assign_init_code(self):
+        class MyClass:
+            def __init__(self):
+                pass
+
+        def instantiate():
+            return MyClass()
+
+        # Trigger specialization
+        for _ in range(1025):
+            instantiate()
+        self.assert_specialized(instantiate, "CALL_ALLOC_AND_ENTER_INIT")
+
+        def count_args(self, *args):
+            self.num_args = len(args)
+
+        # Set MyClass.__init__.__code__ to a code object that is incompatible
+        # (uses varargs) with the current specialization
+        MyClass.__init__.__code__ = count_args.__code__
+        instantiate()
+
 
 @threading_helper.requires_working_threading()
 @requires_specialization
-class TestRacesDoNotCrash(unittest.TestCase):
+class TestRacesDoNotCrash(TestBase):
     # Careful with these. Bigger numbers have a higher chance of catching bugs,
     # but you can also burn through a *ton* of type/dict/function versions:
     ITEMS = 1000
     LOOPS = 4
     WARMUPS = 2
     WRITERS = 2
 
-    def assert_specialized(self, f, opname):
-        instructions = dis.get_instructions(f, adaptive=True)
-        opnames = {instruction.opname for instruction in instructions}
-        self.assertIn(opname, opnames)
-
     @disabling_optimizer
     def assert_races_do_not_crash(
         self, opname, get_items, read, write, *, check_items=False

Python/bytecodes.c

@@ -3563,6 +3563,7 @@ dummy_func(
             PyFunctionObject *init = (PyFunctionObject *)cls->_spec_cache.init;
             PyCodeObject *code = (PyCodeObject *)init->func_code;
             DEOPT_IF(code->co_argcount != oparg+1);
+            DEOPT_IF(_Py_Specialize_CheckCode(code) != _Py_Specialize_Ok);
             DEOPT_IF(!_PyThreadState_HasStackSpace(tstate, code->co_framesize + _Py_InitCleanup.co_framesize));
             STAT_INC(CALL, hit);
             PyObject *self = _PyType_NewManagedObject(tp);

Python/specialize.c

@@ -468,7 +468,6 @@ _PyCode_Quicken(PyCodeObject *code)
 #define SPEC_FAIL_CODE_COMPLEX_PARAMETERS 7
 #define SPEC_FAIL_CODE_NOT_OPTIMIZED 8
 
-
 #define SPEC_FAIL_LOAD_GLOBAL_NON_DICT 17
 #define SPEC_FAIL_LOAD_GLOBAL_NON_STRING_OR_SPLIT 18
 
@@ -1469,15 +1468,20 @@ binary_subscr_fail_kind(PyTypeObject *container_type, PyObject *sub)
 #endif   // Py_STATS
 
 static int
-function_kind(PyCodeObject *code) {
-    int flags = code->co_flags;
-    if ((flags & (CO_VARKEYWORDS | CO_VARARGS)) || code->co_kwonlyargcount) {
-        return SPEC_FAIL_CODE_COMPLEX_PARAMETERS;
-    }
-    if ((flags & CO_OPTIMIZED) == 0) {
-        return SPEC_FAIL_CODE_NOT_OPTIMIZED;
+function_kind(PyCodeObject *code)
+{
+    switch (_Py_Specialize_CheckCode(code)) {
+        case _Py_Specialize_ErrComplexParams: {
+            return SPEC_FAIL_CODE_COMPLEX_PARAMETERS;
+        }
+        case _Py_Specialize_ErrCodeUnoptimized: {
+            return SPEC_FAIL_CODE_NOT_OPTIMIZED;
+        }
+        case _Py_Specialize_Ok: {
+            return SIMPLE_FUNCTION;
+        }
     }
-    return SIMPLE_FUNCTION;
+    Py_UNREACHABLE();
 }
 
 /* Returning false indicates a failure. */