gh-87092: Make jump target label equal to the offset of the target in the instructions sequence (#102093)

Author: iritkatriel

Lib/test/support/bytecode_helper.py

@@ -50,18 +50,13 @@ class CompilationStepTestCase(unittest.TestCase):
     HAS_TARGET = set(dis.hasjrel + dis.hasjabs + dis.hasexc)
     HAS_ARG_OR_TARGET = HAS_ARG.union(HAS_TARGET)
 
-    def setUp(self):
-        self.last_label = 0
-
-    def Label(self):
-        self.last_label += 1
-        return self.last_label
+    class Label:
+        pass
 
     def assertInstructionsMatch(self, actual_, expected_):
         # get two lists where each entry is a label or
-        # an instruction tuple. Compare them, while mapping
-        # each actual label to a corresponding expected label
-        # based on their locations.
+        # an instruction tuple. Normalize the labels to the
+        # instruction count of the target, and compare the lists.
 
         self.assertIsInstance(actual_, list)
         self.assertIsInstance(expected_, list)
@@ -82,39 +77,35 @@ def assertInstructionsMatch(self, actual_, expected_):
                 act = act[:len(exp)]
             self.assertEqual(exp, act)
 
+    def resolveAndRemoveLabels(self, insts):
+        idx = 0
+        res = []
+        for item in insts:
+            assert isinstance(item, (self.Label, tuple))
+            if isinstance(item, self.Label):
+                item.value = idx
+            else:
+                idx += 1
+                res.append(item)
+
+        return res
+
     def normalize_insts(self, insts):
         """ Map labels to instruction index.
-            Remove labels which are not used as jump targets.
             Map opcodes to opnames.
         """
-        labels_map = {}
-        targets = set()
-        idx = 1
-        for item in insts:
-            assert isinstance(item, (int, tuple))
-            if isinstance(item, tuple):
-                opcode, oparg, *_ = item
-                if dis.opmap.get(opcode, opcode) in self.HAS_TARGET:
-                    targets.add(oparg)
-                idx += 1
-            elif isinstance(item, int):
-                assert item not in labels_map, "label reused"
-                labels_map[item] = idx
-
+        insts = self.resolveAndRemoveLabels(insts)
         res = []
         for item in insts:
-            if isinstance(item, int) and item in targets:
-                if not res or labels_map[item] != res[-1]:
-                    res.append(labels_map[item])
-            elif isinstance(item, tuple):
-                opcode, oparg, *loc = item
-                opcode = dis.opmap.get(opcode, opcode)
-                if opcode in self.HAS_TARGET:
-                    arg = labels_map[oparg]
-                else:
-                    arg = oparg if opcode in self.HAS_TARGET else None
-                opcode = dis.opname[opcode]
-                res.append((opcode, arg, *loc))
+            assert isinstance(item, tuple)
+            opcode, oparg, *loc = item
+            opcode = dis.opmap.get(opcode, opcode)
+            if isinstance(oparg, self.Label):
+                arg = oparg.value
+            else:
+                arg = oparg if opcode in self.HAS_ARG else None
+            opcode = dis.opname[opcode]
+            res.append((opcode, arg, *loc))
         return res
 
 
@@ -129,20 +120,18 @@ class CfgOptimizationTestCase(CompilationStepTestCase):
 
     def complete_insts_info(self, insts):
         # fill in omitted fields in location, and oparg 0 for ops with no arg.
-        instructions = []
+        res = []
         for item in insts:
-            if isinstance(item, int):
-                instructions.append(item)
-            else:
-                assert isinstance(item, tuple)
-                inst = list(reversed(item))
-                opcode = dis.opmap[inst.pop()]
-                oparg = inst.pop() if opcode in self.HAS_ARG_OR_TARGET else 0
-                loc = inst + [-1] * (4 - len(inst))
-                instructions.append((opcode, oparg, *loc))
-        return instructions
+            assert isinstance(item, tuple)
+            inst = list(reversed(item))
+            opcode = dis.opmap[inst.pop()]
+            oparg = inst.pop() if opcode in self.HAS_ARG_OR_TARGET else 0
+            loc = inst + [-1] * (4 - len(inst))
+            res.append((opcode, oparg, *loc))
+        return res
 
     def get_optimized(self, insts, consts):
+        insts = self.normalize_insts(insts)
         insts = self.complete_insts_info(insts)
         insts = optimize_cfg(insts, consts)
         return insts, consts

Lib/test/test_compiler_codegen.py

@@ -37,11 +37,11 @@ def test_for_loop(self):
             ('GET_ITER', None, 1),
             loop_lbl := self.Label(),
             ('FOR_ITER', exit_lbl := self.Label(), 1),
-            ('STORE_NAME', None, 1),
+            ('STORE_NAME', 1, 1),
             ('PUSH_NULL', None, 2),
-            ('LOAD_NAME', None, 2),
-            ('LOAD_NAME', None, 2),
-            ('CALL', None, 2),
+            ('LOAD_NAME', 2, 2),
+            ('LOAD_NAME', 1, 2),
+            ('CALL', 1, 2),
             ('POP_TOP', None),
             ('JUMP', loop_lbl),
             exit_lbl,

Lib/test/test_peepholer.py

@@ -984,6 +984,7 @@ def cfg_optimization_test(self, insts, expected_insts,
         if expected_consts is None:
             expected_consts = consts
         opt_insts, opt_consts = self.get_optimized(insts, consts)
+        expected_insts = self.normalize_insts(expected_insts)
         self.assertInstructionsMatch(opt_insts, expected_insts)
         self.assertEqual(opt_consts, expected_consts)
 
@@ -996,11 +997,11 @@ def test_conditional_jump_forward_non_const_condition(self):
             ('LOAD_CONST', 3, 14),
         ]
         expected = [
-            ('LOAD_NAME', '1', 11),
+            ('LOAD_NAME', 1, 11),
             ('POP_JUMP_IF_TRUE', lbl := self.Label(), 12),
-            ('LOAD_CONST', '2', 13),
+            ('LOAD_CONST', 2, 13),
             lbl,
-            ('LOAD_CONST', '3', 14)
+            ('LOAD_CONST', 3, 14)
         ]
         self.cfg_optimization_test(insts, expected, consts=list(range(5)))
 
@@ -1018,7 +1019,7 @@ def test_conditional_jump_forward_const_condition(self):
         expected = [
             ('NOP', None, 11),
             ('NOP', None, 12),
-            ('LOAD_CONST', '3', 14)
+            ('LOAD_CONST', 3, 14)
         ]
         self.cfg_optimization_test(insts, expected, consts=list(range(5)))
 
@@ -1031,9 +1032,9 @@ def test_conditional_jump_backward_non_const_condition(self):
         ]
         expected = [
             lbl := self.Label(),
-            ('LOAD_NAME', '1', 11),
+            ('LOAD_NAME', 1, 11),
             ('POP_JUMP_IF_TRUE', lbl, 12),
-            ('LOAD_CONST', '2', 13)
+            ('LOAD_CONST', 2, 13)
         ]
         self.cfg_optimization_test(insts, expected, consts=list(range(5)))
 

Python/compile.c

@@ -9704,56 +9704,77 @@ instructions_to_cfg(PyObject *instructions, cfg_builder *g)
 {
     assert(PyList_Check(instructions));
 
-    Py_ssize_t instr_size = PyList_GET_SIZE(instructions);
-    for (Py_ssize_t i = 0; i < instr_size; i++) {
+    Py_ssize_t num_insts = PyList_GET_SIZE(instructions);
+    bool *is_target = PyMem_Calloc(num_insts, sizeof(bool));
+    if (is_target == NULL) {
+        return ERROR;
+    }
+    for (Py_ssize_t i = 0; i < num_insts; i++) {
         PyObject *item = PyList_GET_ITEM(instructions, i);
-        if (PyLong_Check(item)) {
-            int lbl_id = PyLong_AsLong(item);
+        if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 6) {
+            PyErr_SetString(PyExc_ValueError, "expected a 6-tuple");
+            goto error;
+        }
+        int opcode = PyLong_AsLong(PyTuple_GET_ITEM(item, 0));
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+        if (HAS_TARGET(opcode)) {
+            int oparg = PyLong_AsLong(PyTuple_GET_ITEM(item, 1));
             if (PyErr_Occurred()) {
-                return ERROR;
+                goto error;
             }
-            if (lbl_id <= 0 || lbl_id > instr_size) {
-                /* expect label in a reasonable range */
+            if (oparg < 0 || oparg >= num_insts) {
                 PyErr_SetString(PyExc_ValueError, "label out of range");
-                return ERROR;
+                goto error;
             }
-            jump_target_label lbl = {lbl_id};
+            is_target[oparg] = true;
+        }
+    }
+
+    for (int i = 0; i < num_insts; i++) {
+        if (is_target[i]) {
+            jump_target_label lbl = {i};
             RETURN_IF_ERROR(cfg_builder_use_label(g, lbl));
         }
-        else {
-            if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 6) {
-                PyErr_SetString(PyExc_ValueError, "expected a 6-tuple");
-                return ERROR;
-            }
-            int opcode = PyLong_AsLong(PyTuple_GET_ITEM(item, 0));
-            if (PyErr_Occurred()) {
-                return ERROR;
-            }
-            int oparg = PyLong_AsLong(PyTuple_GET_ITEM(item, 1));
-            if (PyErr_Occurred()) {
-                return ERROR;
-            }
-            location loc;
-            loc.lineno = PyLong_AsLong(PyTuple_GET_ITEM(item, 2));
-            if (PyErr_Occurred()) {
-                return ERROR;
-            }
-            loc.end_lineno = PyLong_AsLong(PyTuple_GET_ITEM(item, 3));
-            if (PyErr_Occurred()) {
-                return ERROR;
-            }
-            loc.col_offset = PyLong_AsLong(PyTuple_GET_ITEM(item, 4));
-            if (PyErr_Occurred()) {
-                return ERROR;
-            }
-            loc.end_col_offset = PyLong_AsLong(PyTuple_GET_ITEM(item, 5));
-            if (PyErr_Occurred()) {
-                return ERROR;
-            }
-            RETURN_IF_ERROR(cfg_builder_addop(g, opcode, oparg, loc));
+        PyObject *item = PyList_GET_ITEM(instructions, i);
+        if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 6) {
+            PyErr_SetString(PyExc_ValueError, "expected a 6-tuple");
+            goto error;
+        }
+        int opcode = PyLong_AsLong(PyTuple_GET_ITEM(item, 0));
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+        int oparg = PyLong_AsLong(PyTuple_GET_ITEM(item, 1));
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+        location loc;
+        loc.lineno = PyLong_AsLong(PyTuple_GET_ITEM(item, 2));
+        if (PyErr_Occurred()) {
+            goto error;
         }
+        loc.end_lineno = PyLong_AsLong(PyTuple_GET_ITEM(item, 3));
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+        loc.col_offset = PyLong_AsLong(PyTuple_GET_ITEM(item, 4));
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+        loc.end_col_offset = PyLong_AsLong(PyTuple_GET_ITEM(item, 5));
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+        RETURN_IF_ERROR(cfg_builder_addop(g, opcode, oparg, loc));
     }
+
+    PyMem_Free(is_target);
     return SUCCESS;
+error:
+    PyMem_Free(is_target);
+    return ERROR;
 }
 
 static PyObject *
@@ -9763,20 +9784,12 @@ cfg_to_instructions(cfg_builder *g)
     if (instructions == NULL) {
         return NULL;
     }
-    int lbl = 1;
+    int lbl = 0;
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
-        b->b_label = lbl++;
+        b->b_label = lbl;
+        lbl += b->b_iused;
     }
     for (basicblock *b = g->g_entryblock; b != NULL; b = b->b_next) {
-        PyObject *lbl = PyLong_FromLong(b->b_label);
-        if (lbl == NULL) {
-            goto error;
-        }
-        if (PyList_Append(instructions, lbl) != 0) {
-            Py_DECREF(lbl);
-            goto error;
-        }
-        Py_DECREF(lbl);
         for (int i = 0; i < b->b_iused; i++) {
             struct instr *instr = &b->b_instr[i];
             location loc = instr->i_loc;