gh-117288: Allocate fewer label IDs in _PyCfg_ToInstructionSequence (#117290)

diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
index 0f446a00b4df22cb830cf0962af577e3e3724192..f54f4f7f37acee72c7fd4c628035842862fe6b8e 100644 (file)
--- a/Include/internal/pycore_compile.h
+++ b/Include/internal/pycore_compile.h
@@ -66,6 +66,7 @@ int _PyCompile_InstructionSequence_UseLabel(_PyCompile_InstructionSequence *seq,
 int _PyCompile_InstructionSequence_Addop(_PyCompile_InstructionSequence *seq,
                                          int opcode, int oparg,
                                          _PyCompilerSrcLocation loc);
+int _PyCompile_InstructionSequence_ApplyLabelMap(_PyCompile_InstructionSequence *seq);
 
 typedef struct {
     PyObject *u_name;

diff --git a/Python/assemble.c b/Python/assemble.c
index 569454ebf3b9cb03d27ba976d4bbf8005b542c36..09db2fab48d95c2259f54dad99f51bc09553b82f 100644 (file)
--- a/Python/assemble.c
+++ b/Python/assemble.c
@@ -736,6 +736,9 @@ _PyAssemble_MakeCodeObject(_PyCompile_CodeUnitMetadata *umd, PyObject *const_cac
                            int nlocalsplus, int code_flags, PyObject *filename)
 {
 
+    if (_PyCompile_InstructionSequence_ApplyLabelMap(instrs) < 0) {
+        return NULL;
+    }
     if (resolve_unconditional_jumps(instrs) < 0) {
         return NULL;
     }

diff --git a/Python/compile.c b/Python/compile.c
index e9507e47dac8fe15773ba510bb4f80204dfedd7f..43b3cbd4e1894c8bf7a859233fc89e5a40947736 100644 (file)
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -235,6 +235,28 @@ _PyCompile_InstructionSequence_UseLabel(instr_sequence *seq, int lbl)
     return SUCCESS;
 }
 
+int
+_PyCompile_InstructionSequence_ApplyLabelMap(instr_sequence *instrs)
+{
+    /* Replace labels by offsets in the code */
+    for (int i=0; i < instrs->s_used; i++) {
+        instruction *instr = &instrs->s_instrs[i];
+        if (HAS_TARGET(instr->i_opcode)) {
+            assert(instr->i_oparg < instrs->s_labelmap_size);
+            instr->i_oparg = instrs->s_labelmap[instr->i_oparg];
+        }
+        _PyCompile_ExceptHandlerInfo *hi = &instr->i_except_handler_info;
+        if (hi->h_label >= 0) {
+            assert(hi->h_label < instrs->s_labelmap_size);
+            hi->h_label = instrs->s_labelmap[hi->h_label];
+        }
+    }
+    /* Clear label map so it's never used again */
+    PyMem_Free(instrs->s_labelmap);
+    instrs->s_labelmap = NULL;
+    instrs->s_labelmap_size = 0;
+    return SUCCESS;
+}
 
 #define MAX_OPCODE 511
 
@@ -7824,11 +7846,8 @@ instr_sequence_to_instructions(instr_sequence *seq)
     for (int i = 0; i < seq->s_used; i++) {
         instruction *instr = &seq->s_instrs[i];
         location loc = instr->i_loc;
-        int arg = HAS_TARGET(instr->i_opcode) ?
-                  seq->s_labelmap[instr->i_oparg] : instr->i_oparg;
-
         PyObject *inst_tuple = Py_BuildValue(
-            "(iiiiii)", instr->i_opcode, arg,
+            "(iiiiii)", instr->i_opcode, instr->i_oparg,
             loc.lineno, loc.end_lineno,
             loc.col_offset, loc.end_col_offset);
         if (inst_tuple == NULL) {
@@ -7855,6 +7874,9 @@ cfg_to_instructions(cfg_builder *g)
     if (_PyCfg_ToInstructionSequence(g, &seq) < 0) {
         return NULL;
     }
+    if (_PyCompile_InstructionSequence_ApplyLabelMap(&seq) < 0) {
+        return NULL;
+    }
     PyObject *res = instr_sequence_to_instructions(&seq);
     instr_sequence_fini(&seq);
     return res;
@@ -8026,6 +8048,10 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
         goto finally;
     }
 
+    if (_PyCompile_InstructionSequence_ApplyLabelMap(INSTR_SEQUENCE(c)) < 0) {
+        return NULL;
+    }
+
     PyObject *insts = instr_sequence_to_instructions(INSTR_SEQUENCE(c));
     if (insts == NULL) {
         goto finally;

diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index 2f47e47bf9d29d03e12cfe69d0e367140f1f947f..5437c3875ff7b026852934bc41e5336a8e2b814a 100644 (file)
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -2717,13 +2717,14 @@ _PyCfg_ToInstructionSequence(cfg_builder *g, _PyCompile_InstructionSequence *seq
     int lbl = 0;
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
         b->b_label = (jump_target_label){lbl};
-        lbl += b->b_iused;
+        lbl += 1;
     }
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
         RETURN_IF_ERROR(_PyCompile_InstructionSequence_UseLabel(seq, b->b_label.id));
         for (int i = 0; i < b->b_iused; i++) {
             cfg_instr *instr = &b->b_instr[i];
-            if (OPCODE_HAS_JUMP(instr->i_opcode) || is_block_push(instr)) {
+            if (HAS_TARGET(instr->i_opcode)) {
+                /* Set oparg to the label id (it will later be mapped to an offset) */
                 instr->i_oparg = instr->i_target->b_label.id;
             }
             RETURN_IF_ERROR(