gh-117680: make _PyInstructionSequence a PyObject and use it in tests (#117629)

Author: iritkatriel

Include/internal/pycore_global_objects_fini_generated.h

@@ -353,6 +353,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(co_stacksize)
         STRUCT_FOR_ID(co_varnames)
         STRUCT_FOR_ID(code)
+        STRUCT_FOR_ID(col_offset)
         STRUCT_FOR_ID(command)
         STRUCT_FOR_ID(comment_factory)
         STRUCT_FOR_ID(compile_mode)
@@ -402,6 +403,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(encode)
         STRUCT_FOR_ID(encoding)
         STRUCT_FOR_ID(end)
+        STRUCT_FOR_ID(end_col_offset)
         STRUCT_FOR_ID(end_lineno)
         STRUCT_FOR_ID(end_offset)
         STRUCT_FOR_ID(endpos)
@@ -522,6 +524,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(kw1)
         STRUCT_FOR_ID(kw2)
         STRUCT_FOR_ID(kwdefaults)
+        STRUCT_FOR_ID(label)
         STRUCT_FOR_ID(lambda)
         STRUCT_FOR_ID(last)
         STRUCT_FOR_ID(last_exc)
@@ -585,6 +588,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(namespaces)
         STRUCT_FOR_ID(narg)
         STRUCT_FOR_ID(ndigits)
+        STRUCT_FOR_ID(nested)
         STRUCT_FOR_ID(new_file_name)
         STRUCT_FOR_ID(new_limit)
         STRUCT_FOR_ID(newline)

Include/internal/pycore_global_strings.h

@@ -5,10 +5,13 @@
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
+#include "pycore_symtable.h"
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+
 typedef struct {
     int h_label;
     int h_startdepth;
@@ -26,23 +29,30 @@ typedef struct {
     int i_offset;
 } _PyInstruction;
 
-typedef struct {
+typedef struct instruction_sequence {
+    PyObject_HEAD
     _PyInstruction *s_instrs;
     int s_allocated;
     int s_used;
 
     int s_next_free_label; /* next free label id */
+
     /* Map of a label id to instruction offset (index into s_instrs).
      * If s_labelmap is NULL, then each label id is the offset itself.
      */
-    int *s_labelmap;       /* label id --> instr offset */
+    int *s_labelmap;
     int s_labelmap_size;
+
+    /* PyList of instruction sequences of nested functions */
+    PyObject *s_nested;
 } _PyInstructionSequence;
 
 typedef struct {
     int id;
 } _PyJumpTargetLabel;
 
+PyAPI_FUNC(PyObject*)_PyInstructionSequence_New(void);
+
 int _PyInstructionSequence_UseLabel(_PyInstructionSequence *seq, int lbl);
 int _PyInstructionSequence_Addop(_PyInstructionSequence *seq,
                                  int opcode, int oparg,
@@ -53,6 +63,8 @@ int _PyInstructionSequence_InsertInstruction(_PyInstructionSequence *seq, int po
                                              int opcode, int oparg, _Py_SourceLocation loc);
 void PyInstructionSequence_Fini(_PyInstructionSequence *seq);
 
+extern PyTypeObject _PyInstructionSequence_Type;
+#define _PyInstructionSequence_Check(v) Py_IS_TYPE((v), &_PyInstructionSequence_Type)
 
 #ifdef __cplusplus
 }

Include/internal/pycore_instruction_sequence.h

@@ -3,6 +3,7 @@
 import unittest
 import dis
 import io
+import opcode
 try:
     import _testinternalcapi
 except ImportError:
@@ -68,16 +69,14 @@ class CompilationStepTestCase(unittest.TestCase):
     class Label:
         pass
 
-    def assertInstructionsMatch(self, actual_, expected_):
-        # get two lists where each entry is a label or
-        # an instruction tuple. Normalize the labels to the
-        # instruction count of the target, and compare the lists.
+    def assertInstructionsMatch(self, actual_seq, expected):
+        # get an InstructionSequence and an expected list, where each
+        # entry is a label or an instruction tuple. Construct an expcted
+        # instruction sequence and compare with the one given.
 
-        self.assertIsInstance(actual_, list)
-        self.assertIsInstance(expected_, list)
-
-        actual = self.normalize_insts(actual_)
-        expected = self.normalize_insts(expected_)
+        self.assertIsInstance(expected, list)
+        actual = actual_seq.get_instructions()
+        expected = self.seq_from_insts(expected).get_instructions()
         self.assertEqual(len(actual), len(expected))
 
         # compare instructions
@@ -87,10 +86,8 @@ def assertInstructionsMatch(self, actual_, expected_):
                 continue
             self.assertIsInstance(exp, tuple)
             self.assertIsInstance(act, tuple)
-            # crop comparison to the provided expected values
-            if len(act) > len(exp):
-                act = act[:len(exp)]
-            self.assertEqual(exp, act)
+            idx = max([p[0] for p in enumerate(exp) if p[1] != -1])
+            self.assertEqual(exp[:idx], act[:idx])
 
     def resolveAndRemoveLabels(self, insts):
         idx = 0
@@ -105,35 +102,37 @@ def resolveAndRemoveLabels(self, insts):
 
         return res
 
-    def normalize_insts(self, insts):
-        """ Map labels to instruction index.
-            Map opcodes to opnames.
-        """
-        insts = self.resolveAndRemoveLabels(insts)
-        res = []
-        for item in insts:
-            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
+    def seq_from_insts(self, insts):
+        labels = {item for item in insts if isinstance(item, self.Label)}
+        for i, lbl in enumerate(labels):
+            lbl.value = i
 
-    def complete_insts_info(self, insts):
-        # fill in omitted fields in location, and oparg 0 for ops with no arg.
-        res = []
+        seq = _testinternalcapi.new_instruction_sequence()
         for item in insts:
-            assert isinstance(item, tuple)
-            inst = list(item)
-            opcode = dis.opmap[inst[0]]
-            oparg = inst[1]
-            loc = inst[2:] + [-1] * (6 - len(inst))
-            res.append((opcode, oparg, *loc))
-        return res
+            if isinstance(item, self.Label):
+                seq.use_label(item.value)
+            else:
+                op = item[0]
+                if isinstance(op, str):
+                    op = opcode.opmap[op]
+                arg, *loc = item[1:]
+                if isinstance(arg, self.Label):
+                    arg = arg.value
+                loc = loc + [-1] * (4 - len(loc))
+                seq.addop(op, arg or 0, *loc)
+        return seq
+
+    def check_instructions(self, insts):
+        for inst in insts:
+            if isinstance(inst, self.Label):
+                continue
+            op, arg, *loc = inst
+            if isinstance(op, str):
+                op = opcode.opmap[op]
+            self.assertEqual(op in opcode.hasarg,
+                             arg is not None,
+                             f"{opcode.opname[op]=} {arg=}")
+            self.assertTrue(all(isinstance(l, int) for l in loc))
 
 
 @unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
@@ -147,10 +146,8 @@ def generate_code(self, ast):
 @unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
 class CfgOptimizationTestCase(CompilationStepTestCase):
 
-    def get_optimized(self, insts, consts, nlocals=0):
-        insts = self.normalize_insts(insts)
-        insts = self.complete_insts_info(insts)
-        insts = _testinternalcapi.optimize_cfg(insts, consts, nlocals)
+    def get_optimized(self, seq, consts, nlocals=0):
+        insts = _testinternalcapi.optimize_cfg(seq, consts, nlocals)
         return insts, consts
 
 @unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")

Include/internal/pycore_runtime_init_generated.h

@@ -2,6 +2,7 @@
 import dis
 import io
 import math
+import opcode
 import os
 import unittest
 import sys
@@ -11,6 +12,8 @@
 import types
 import textwrap
 import warnings
+import _testinternalcapi
+
 from test import support
 from test.support import (script_helper, requires_debug_ranges,
                           requires_specialization, get_c_recursion_limit)
@@ -2419,6 +2422,49 @@ def test_return_inside_async_with_block(self):
             """
         self.check_stack_size(snippet, async_=True)
 
+class TestInstructionSequence(unittest.TestCase):
+    def compare_instructions(self, seq, expected):
+        self.assertEqual([(opcode.opname[i[0]],) + i[1:] for i in seq.get_instructions()],
+                         expected)
+
+    def test_basics(self):
+        seq = _testinternalcapi.new_instruction_sequence()
+
+        def add_op(seq, opname, oparg, bl, bc=0, el=0, ec=0):
+            seq.addop(opcode.opmap[opname], oparg, bl, bc, el, el)
+
+        add_op(seq, 'LOAD_CONST', 1, 1)
+        add_op(seq, 'JUMP', lbl1 := seq.new_label(), 2)
+        add_op(seq, 'LOAD_CONST', 1, 3)
+        add_op(seq, 'JUMP', lbl2 := seq.new_label(), 4)
+        seq.use_label(lbl1)
+        add_op(seq, 'LOAD_CONST', 2, 4)
+        seq.use_label(lbl2)
+        add_op(seq, 'RETURN_VALUE', 0, 3)
+
+        expected = [('LOAD_CONST', 1, 1),
+                    ('JUMP', 4, 2),
+                    ('LOAD_CONST', 1, 3),
+                    ('JUMP', 5, 4),
+                    ('LOAD_CONST', 2, 4),
+                    ('RETURN_VALUE', None, 3),
+                   ]
+
+        self.compare_instructions(seq, [ex + (0,0,0) for ex in expected])
+
+    def test_nested(self):
+        seq = _testinternalcapi.new_instruction_sequence()
+        seq.addop(opcode.opmap['LOAD_CONST'], 1, 1, 0, 0, 0)
+        nested = _testinternalcapi.new_instruction_sequence()
+        nested.addop(opcode.opmap['LOAD_CONST'], 2, 2, 0, 0, 0)
+
+        self.compare_instructions(seq, [('LOAD_CONST', 1, 1, 0, 0, 0)])
+        self.compare_instructions(nested, [('LOAD_CONST', 2, 2, 0, 0, 0)])
+
+        seq.add_nested(nested)
+        self.compare_instructions(seq, [('LOAD_CONST', 1, 1, 0, 0, 0)])
+        self.compare_instructions(seq.get_nested()[0], [('LOAD_CONST', 2, 2, 0, 0, 0)])
+
 
 if __name__ == "__main__":
     unittest.main()