From 262fb911ab7df8e890ebd0efb0773c3e0b5a757f Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Wed, 27 Mar 2024 17:38:19 +0000
Subject: [PATCH] gh-117288: Allocate fewer label IDs in
 _PyCfg_ToInstructionSequence (#117290)

---
 Include/internal/pycore_compile.h |  1 +
 Python/assemble.c                 |  3 +++
 Python/compile.c                  | 34 +++++++++++++++++++++++++++----
 Python/flowgraph.c                |  5 +++--
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
index 0f446a00b4df22..f54f4f7f37acee 100644
--- 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 569454ebf3b9cb..09db2fab48d95c 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 e9507e47dac8fe..43b3cbd4e1894c 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 2f47e47bf9d29d..5437c3875ff7b0 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(