Let compile.c manage the instrseq

Author: JelleZijlstra

Include/internal/pycore_compile.h

@@ -133,8 +133,8 @@ int _PyCompile_EnterScope(struct _PyCompiler *c, identifier name, int scope_type
 void _PyCompile_ExitScope(struct _PyCompiler *c);
 Py_ssize_t _PyCompile_AddConst(struct _PyCompiler *c, PyObject *o);
 _PyInstructionSequence *_PyCompile_InstrSequence(struct _PyCompiler *c);
-void _PyCompile_SetInstrSequence(struct _PyCompiler *c,
-                                 _PyInstructionSequence *instr_sequence);
+int _PyCompile_StartAnnotationSetup(struct _PyCompiler *c);
+int _PyCompile_EndAnnotationSetup(struct _PyCompiler *c);
 int _PyCompile_FutureFeatures(struct _PyCompiler *c);
 void _PyCompile_DeferredAnnotations(
     struct _PyCompiler *c, PyObject **deferred_annotations,

Python/codegen.c

@@ -795,15 +795,12 @@ codegen_process_deferred_annotations(compiler *c, location loc)
         return SUCCESS;
     }
 
-    instr_sequence *old_instr_seq = INSTR_SEQUENCE(c);
-    instr_sequence *nested_instr_seq = NULL;
     int scope_type = SCOPE_TYPE(c);
-    if (scope_type == COMPILE_SCOPE_MODULE) {
-        nested_instr_seq = (instr_sequence *)_PyInstructionSequence_New();
-        if (nested_instr_seq == NULL) {
+    bool need_separate_block = scope_type == COMPILE_SCOPE_MODULE;
+    if (need_separate_block) {
+        if (_PyCompile_StartAnnotationSetup(c) == ERROR) {
             goto error;
         }
-        _PyCompile_SetInstrSequence(c, nested_instr_seq);
     }
 
     // It's possible that ste_annotations_block is set but
@@ -833,18 +830,12 @@ codegen_process_deferred_annotations(compiler *c, location loc)
         ste->ste_type == ClassBlock ? &_Py_ID(__annotate_func__) : &_Py_ID(__annotate__),
         Store));
 
-    if (nested_instr_seq != NULL) {
-        RETURN_IF_ERROR(
-            _PyInstructionSequence_SetAnnotationsCode(old_instr_seq, nested_instr_seq));
-        _PyCompile_SetInstrSequence(c, old_instr_seq);
+    if (need_separate_block) {
+        RETURN_IF_ERROR(_PyCompile_EndAnnotationSetup(c));
     }
 
     return SUCCESS;
 error:
-    if (nested_instr_seq != NULL) {
-        PyInstructionSequence_Fini(nested_instr_seq);
-        _PyCompile_SetInstrSequence(c, old_instr_seq);
-    }
     Py_XDECREF(deferred_anno);
     Py_XDECREF(conditional_annotation_indices);
     return ERROR;

Python/compile.c

@@ -64,6 +64,7 @@ struct compiler_unit {
     long u_next_conditional_annotation_index;  /* index of the next conditional annotation */
 
     instr_sequence *u_instr_sequence; /* codegen output */
+    instr_sequence *u_stashed_instr_sequence; /* temporarily stashed parent instruction sequence */
 
     int u_nfblocks;
     int u_in_inlined_comp;
@@ -178,6 +179,7 @@ static void
 compiler_unit_free(struct compiler_unit *u)
 {
     Py_CLEAR(u->u_instr_sequence);
+    Py_CLEAR(u->u_stashed_instr_sequence);
     Py_CLEAR(u->u_ste);
     Py_CLEAR(u->u_metadata.u_name);
     Py_CLEAR(u->u_metadata.u_qualname);
@@ -681,6 +683,7 @@ _PyCompile_EnterScope(compiler *c, identifier name, int scope_type,
         compiler_unit_free(u);
         return ERROR;
     }
+    u->u_stashed_instr_sequence = NULL;
 
     /* Push the old compiler_unit on the stack. */
     if (c->u) {
@@ -1231,12 +1234,35 @@ _PyCompile_InstrSequence(compiler *c)
     return c->u->u_instr_sequence;
 }
 
-void
-_PyCompile_SetInstrSequence(compiler *c, instr_sequence *seq)
+int
+_PyCompile_StartAnnotationSetup(struct _PyCompiler *c)
 {
-    c->u->u_instr_sequence = seq;
+    instr_sequence *new_seq = (instr_sequence *)_PyInstructionSequence_New();
+    if (new_seq == NULL) {
+        return ERROR;
+    }
+    assert(c->u->u_stashed_instr_sequence == NULL);
+    c->u->u_stashed_instr_sequence = c->u->u_instr_sequence;
+    c->u->u_instr_sequence = new_seq;
+    return SUCCESS;
 }
 
+int
+_PyCompile_EndAnnotationSetup(struct _PyCompiler *c)
+{
+    assert(c->u->u_stashed_instr_sequence != NULL);
+    instr_sequence *parent_seq = c->u->u_stashed_instr_sequence;
+    instr_sequence *anno_seq = c->u->u_instr_sequence;
+    c->u->u_stashed_instr_sequence = NULL;
+    c->u->u_instr_sequence = parent_seq;
+    if (_PyInstructionSequence_SetAnnotationsCode(parent_seq, anno_seq) == ERROR) {
+        Py_DECREF(anno_seq);
+        return ERROR;
+    }
+    return SUCCESS;
+}
+
+
 int
 _PyCompile_FutureFeatures(compiler *c)
 {