Get Graal line offsets working

Author: rocky

xdis/bytecode.py

@@ -27,7 +27,7 @@
 from io import StringIO
 from linecache import getline
 from types import CodeType
-from typing import Iterable, Iterator, Optional, Tuple, Union
+from typing import Iterable, Iterator, Optional, Tuple
 
 from xdis.cross_dis import (
     format_code_info,
@@ -454,15 +454,8 @@ def next_offset(op: int, opc, offset: int) -> int:
 
 
 def get_instructions_bytes(
-    bytecode,
+    code_object,
     opc,
-    varnames=None,
-    names=None,
-    constants=None,
-    cells=None,
-    linestarts=None,
-    line_offset=0,
-    exception_entries=None,
 ):
     """
     Iterate over the instructions in a bytecode string.
@@ -472,7 +465,20 @@ def get_instructions_bytes(
     e.g., variable names, constants, can be specified using optional
     arguments.
     """
+
+    bytecode: bytes = code_object.co_code
+    constants: tuple = code_object.co_consts
+    names: tuple = code_object.co_names
+    varnames: tuple = code_object.co_varnames
+    cells: tuple = code_object.co_cells if hasattr(code_object, "co_cells") else tuple()
+    exception_entries = code_object.exception_entries if hasattr(code_object, "exception_entries") else tuple()
+    # freevars: tuple = code_object.co_freevars
+
     labels = opc.findlabels(bytecode, opc)
+    if hasattr(opc, "findlinestarts"):
+        line_starts = dict(opc.findlinestarts(code_object, dup_lines=True))
+    else:
+        line_starts = None
 
     # PERFORMANCE FIX: Build exception labels ONCE, not on every iteration
     # The old code was O(n^2) because it rebuilt the same list every call to
@@ -497,7 +503,7 @@ def get_instructions_bytes(
                 names=names,
                 constants=constants,
                 cells=cells,
-                linestarts=linestarts,
+                linestarts=line_starts,
                 line_offset=0,
                 exception_entries=exception_entries,
                 labels=labels,
@@ -552,17 +558,7 @@ def __init__(
 
     def __iter__(self):
         co = self.codeobj
-        return get_instructions_bytes(
-            co.co_code,
-            self.opc,
-            co.co_varnames,
-            co.co_names,
-            co.co_consts,
-            self._cell_names,
-            self._linestarts,
-            line_offset=self._line_offset,
-            exception_entries=self.exception_entries,
-        )
+        return get_instructions_bytes(co, self.opc)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self._original_object!r})"
@@ -596,10 +592,8 @@ def dis(
             offset = -1
         output = StringIO()
         if self.opc.version_tuple > (2, 0):
-            cells = self._cell_names
             line_starts = self._linestarts
         else:
-            cells = None
             line_starts = None
 
         first_line_number = co.co_firstlineno if hasattr(co, "co_firstlineno") else None
@@ -610,17 +604,8 @@ def dis(
         if isinstance(filename, UnicodeForPython3):
             filename = str(filename)
 
-        if isinstance(co.co_code, str):
-            co_code = co.co_code.encode('latin-1')
-        else:
-            co_code = co.co_code
-
         self.disassemble_bytes(
-            co_code,
-            varnames=co.co_varnames,
-            names=co.co_names,
-            constants=co.co_consts,
-            cells=cells,
+            co,
             line_starts=line_starts,
             line_offset=self._line_offset,
             file=output,
@@ -629,7 +614,6 @@ def dis(
             filename=filename,
             show_source=show_source,
             first_line_number=first_line_number,
-            exception_entries=self.exception_entries,
         )
         return output.getvalue()
 
@@ -647,20 +631,15 @@ def distb(self, tb=None) -> None:
 
     def disassemble_bytes(
         self,
-        bytecode: Union[CodeType, bytes, str],
+        code_object: CodeType,
         lasti: int = -1,
-        varnames=None,
-        names=None,
-        constants=None,
-        cells=None,
         line_starts=None,
         file=sys.stdout,
         line_offset=0,
         asm_format="classic",
         filename: Optional[str] = None,
         show_source=True,
         first_line_number: Optional[int] = None,
-        exception_entries=None,
     ) -> list:
         # Omit the line number column entirely if we have no line number info
         show_lineno = line_starts is not None or self.opc.version_tuple < (2, 3)
@@ -703,17 +682,7 @@ def show_source_text(line_number: Optional[int]) -> None:
         else:
             get_instructions_fn = get_instructions_bytes
 
-        for instr in get_instructions_fn(
-            bytecode,
-            self.opc,
-            varnames,
-            names,
-            constants,
-            cells,
-            line_starts,
-            line_offset=line_offset,
-            exception_entries=exception_entries,
-        ):
+        for instr in get_instructions_fn(code_object, self.opc):
             # Python 1.x into early 2.0 uses SET_LINENO
             if last_was_set_lineno:
                 instr = Instruction(
@@ -789,7 +758,7 @@ def show_source_text(line_number: Optional[int]) -> None:
             new_source_line = show_lineno and (
                 extended_arg_starts_line
                 or instr.starts_line is not None
-                and instr.offset > 0
+                and instr.offset >= 0
             )
             if new_source_line:
                 file.write("\n")
@@ -833,7 +802,7 @@ def show_source_text(line_number: Optional[int]) -> None:
             pass
         return instructions
 
-    def get_instructions(self, x, first_line=None):
+    def get_instructions(self, x):
         """Iterator for the opcodes in methods, functions or code
 
         Generates a series of Instruction named tuples giving the details of
@@ -845,22 +814,7 @@ def get_instructions(self, x, first_line=None):
         the disassembled code object.
         """
         co = get_code_object(x)
-        cell_names = co.co_cellvars + co.co_freevars
-        line_starts = dict(self.opc.findlinestarts(co))
-        if first_line is not None:
-            line_offset = first_line - co.co_firstlineno
-        else:
-            line_offset = 0
-        return get_instructions_bytes(
-            co.co_code,
-            self.opc,
-            co.co_varnames,
-            co.co_names,
-            co.co_consts,
-            cell_names,
-            line_starts,
-            line_offset,
-        )
+        return get_instructions_bytes(co.co_code, self.opc)
 
 
 def list2bytecode(

xdis/bytecode_graal.py

@@ -1,7 +1,10 @@
 # from xdis.bytecode import get_optype
+from typing import Optional
+
 from xdis.bytecode import Bytecode
 from xdis.cross_dis import get_code_object
 from xdis.instruction import Instruction
+from xdis.lineoffsets_graal import SourceMap, find_linestarts_graal
 from xdis.opcodes.base_graal import (
     BINARY_OPS,
     COLLECTION_KIND,
@@ -11,15 +14,8 @@
 
 
 def get_instructions_bytes_graal(
-    bytecode,
+    code_object,
     opc,
-    varnames,
-    names,
-    constants,
-    cells,
-    freevars,
-    line_offset,
-    exception_entries,
 ):
     """
     Iterate over the instructions in a bytecode string.
@@ -29,19 +25,26 @@ def get_instructions_bytes_graal(
     e.g., variable names, constants, can be specified using optional
     arguments.
     """
-    # source_map = ???
+    bytecode: bytes = code_object.co_code
+    constants: tuple = code_object.co_consts
+    names: tuple = code_object.co_names
+    varnames: tuple = code_object.co_varnames
+    # cells: tuple = code_object.co_cells
+    # freevars: tuple = code_object.co_freevars
 
     i = 0
     n = len(bytecode)
 
-    extended_arg_count = 0
     labels = opc.findlabels(bytecode, opc)
+    linestarts = find_linestarts_graal(code_object, opc, dup_lines=True)
 
+    extended_arg_count = 0
     while i < n:
         opcode = bytecode[i]
         opname = opc.opname[opcode]
         optype = get_optype_graal(opcode, opc)
         offset = i
+        starts_line = linestarts.get(offset, None)
 
         arg_count = opc.arg_counts[opcode]
         is_jump_target = i in labels
@@ -191,7 +194,7 @@ def get_instructions_bytes_graal(
                 arg = 0
             break
 
-        inst_size = (i - offset + 1) + (extended_arg_count * 2)
+        inst_size = (arg_count + 1) + (extended_arg_count * 2)
         start_offset = offset if opc.oppop[opcode] == 0 else None
 
         # for (int i = 0 i < exceptionHandlerRanges.length; i += 4) {
@@ -242,7 +245,7 @@ def get_instructions_bytes_graal(
 
         yield Instruction(
             is_jump_target=is_jump_target,
-            starts_line=False,  # starts_line,
+            starts_line= starts_line,
             offset=offset,
             opname=opname,
             opcode=opcode,
@@ -280,8 +283,7 @@ def __init__(self, x, opc, first_line=None, current_offset=None) -> None:
             self._line_offset = first_line - co.co_firstlineno
         pass
 
-        # self._linestarts = dict(opc.findlinestarts(co, dup_lines=dup_lines))
-        self._linestarts = None
+        self._linestarts = find_linestarts_graal(co, opc, dup_lines=True)
         self._original_object = x
         self.opc = opc
         self.opnames = opc.opname
@@ -298,11 +300,6 @@ def __iter__(self):
         return get_instructions_bytes_graal(
             co.co_code,
             self.opc,
-            co.co_varnames,
-            co.co_names,
-            co.co_consts,
-            co.co_cellvars,
-            co.co_freevars,
         )
 
     def __repr__(self) -> str:
@@ -319,13 +316,5 @@ def get_instructions(self, x):
         Otherwise, the source line information (if any) is taken directly from
         the disassembled code object.
         """
-        co = get_code_object(x)
-        return get_instructions_bytes_graal(
-            co.co_code,
-            self.opc,
-            co.co_varnames,
-            co.co_names,
-            co.co_consts,
-            co.co_cellvars,
-            co.co_freevars,
-        )
+        code_object = get_code_object(x)
+        return get_instructions_bytes_graal(code_object, self.opc)

xdis/codetype/__init__.py

@@ -33,7 +33,7 @@
 __docformat__ = "restructuredtext"
 
 
-def codeType2Portable(code, version_triple=PYTHON_VERSION_TRIPLE, is_graal: bool=False):
+def codeType2Portable(code, version_triple=PYTHON_VERSION_TRIPLE, is_graal: bool=False, other_fields: dict={}):
     """Converts a native types.CodeType code object into a
     corresponding more flexible xdis Code type.
     """
@@ -149,20 +149,7 @@ def codeType2Portable(code, version_triple=PYTHON_VERSION_TRIPLE, is_graal: bool
                 version_triple=version_triple,
             )
         elif version_triple[:2] >= (3, 11):
-            other_fields = {}
             if is_graal:
-                other_fields["condition_profileCount"]=code.condition_profileCount if hasattr(code, "condition_profileCount") else -1,
-                other_fields["endColumn"]=code.endColumn if hasattr(code, "endColumn") else -1,
-                other_fields["endLine"]=code.endLine if hasattr(code, "endLine") else -1,
-                other_fields["exception_handler_ranges"]=code.exception_handler_ranges if hasattr(code, "exception_handler_ranges") else tuple(),
-                other_fields["generalizeInputsMap"]=code.generalizeInputsMap if hasattr(code, "generalizeInputsMap") else {},
-                other_fields["generalizeVarsMap"]=code.generalizeVarsMap if hasattr(code, "generalizeVarsMap") else {},
-                other_fields["outputCanQuicken"]=code.outputCanQuicken if hasattr(code, "outputCanQuicken") else b"",
-                other_fields["primitiveConstants"]=code.primitiveConstants if hasattr(code, "primitiveConstants") else tuple(),
-                other_fields["srcOffsetTable"]=code.srcOffsetTable if hasattr(code, "srcOffsetTable") else b"",
-                other_fields["startColumn"]=code.startColumn if hasattr(code, "startColumn") else -1,
-                other_fields["startLine"]=code.startLine if hasattr(code, "startLine") else -1,
-                other_fields["variableShouldUnbox"]=code.variableShouldUnbox if hasattr(code, "variableShouldUnbox") else b"",
 
                 return Code311Graal(
                     co_argcount=code.co_argcount,
@@ -363,7 +350,7 @@ def to_portable(
         version_triple=version_triple,
         other_fields=other_fields,
     )
-    return codeType2Portable(code, version_triple, is_graal=is_graal)
+    return codeType2Portable(code, version_triple, is_graal=is_graal, other_fields=other_fields)
 
 
 if __name__ == "__main__":

xdis/disasm.py

@@ -137,7 +137,7 @@ def disco(
     co,
     timestamp,
     out=sys.stdout,
-    magic_int=None,
+    magic_int: int=-1,
     source_size=None,
     sip_hash=None,
     asm_format: str = "classic",