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

3 files changed, 36 insertions, 6 deletions

diff --git a/Python/assemble.c b/Python/assemble.c
index 569454e..09db2fa 100644
--- 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 e9507e4..43b3cbd 100644
--- 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 2f47e47..5437c38 100644
--- 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(