Refactored core datastructures to be compatible with the changes

Author: kc611

numba_rvsdg/core/datastructures/byte_flow.py

@@ -1,5 +1,4 @@
 import dis
-from copy import deepcopy
 from dataclasses import dataclass
 from typing import Generator, Callable
 
@@ -61,90 +60,51 @@ def from_bytecode(code: Callable) -> "ByteFlow":  # type: ignore
         scfg = flowinfo.build_basicblocks()
         return ByteFlow(bc=bc, scfg=scfg)
 
-    def _join_returns(self) -> "ByteFlow":
+    def _join_returns(self) -> None:
         """Joins the return blocks within the corresponding SCFG.
 
-        This method creates a deep copy of the SCFG and performs
-        operation to join return blocks within the control flow.
-        It returns a new ByteFlow object with the updated SCFG.
-
-        Returns
-        -------
-        byteflow: ByteFlow
-            The new ByteFlow object with updated SCFG.
+        This method performs operation to join return blocks within
+        the control flow.
         """
-        scfg = deepcopy(self.scfg)
-        scfg.join_returns()
-        return ByteFlow(bc=self.bc, scfg=scfg)
+        self.scfg.join_returns()
 
-    def _restructure_loop(self) -> "ByteFlow":
+    def _restructure_loop(self) -> None:
         """Restructures the loops within the corresponding SCFG.
 
-        Creates a deep copy of the SCFG and performs the operation to
-        restructure loop constructs within the control flow using
-        the algorithm LOOP RESTRUCTURING from section 4.1 of Bahmann2015.
+        Performs the operation to restructure loop constructs within
+        the control flow using the algorithm LOOP RESTRUCTURING from
+        section 4.1 of Bahmann2015.
         It applies the restructuring operation to both the main SCFG
-        and any subregions within it. It returns a new ByteFlow object
-        with the updated SCFG.
-
-        Returns
-        -------
-        byteflow: ByteFlow
-            The new ByteFlow object with updated SCFG.
+        and any subregions within it.
         """
-        scfg = deepcopy(self.scfg)
-        restructure_loop(scfg.region)
-        for region in _iter_subregions(scfg):
+        restructure_loop(self.scfg.region)
+        for region in _iter_subregions(self.scfg):
             restructure_loop(region)
-        return ByteFlow(bc=self.bc, scfg=scfg)
 
-    def _restructure_branch(self) -> "ByteFlow":
+    def _restructure_branch(self) -> None:
         """Restructures the branches within the corresponding SCFG.
 
-        Creates a deep copy of the SCFG and performs the operation to
-        restructure branch constructs within the control flow. It applies
-        the restructuring operation to both the main SCFG and any
-        subregions within it. It returns a new ByteFlow object with
-        the updated SCFG.
-
-        Returns
-        -------
-        byteflow: ByteFlow
-            The new ByteFlow object with updated SCFG.
+        This method applies restructuring branch operation to both
+        the main SCFG and any subregions within it.
         """
-        scfg = deepcopy(self.scfg)
-        restructure_branch(scfg.region)
-        for region in _iter_subregions(scfg):
+        restructure_branch(self.scfg.region)
+        for region in _iter_subregions(self.scfg):
             restructure_branch(region)
-        return ByteFlow(bc=self.bc, scfg=scfg)
 
-    def restructure(self) -> "ByteFlow":
+    def restructure(self) -> None:
         """Applies join_returns, restructure_loop and restructure_branch
         in the respective order on the SCFG.
 
-        Creates a deep copy of the SCFG and applies a series of
-        restructuring operations to it. The operations include
-        joining return blocks, restructuring loop constructs, and
-        restructuring branch constructs. It returns a new ByteFlow
-        object with the updated SCFG.
-
-        Returns
-        -------
-        byteflow: ByteFlow
-            The new ByteFlow object with updated SCFG.
+        Applies a series of restructuring operations to given SCFG.
+        The operations include joining return blocks, restructuring
+        loop constructs, and restructuring branch constructs.
         """
-        scfg = deepcopy(self.scfg)
         # close
-        scfg.join_returns()
+        self._join_returns()
         # handle loop
-        restructure_loop(scfg.region)
-        for region in _iter_subregions(scfg):
-            restructure_loop(region)
+        self._restructure_loop()
         # handle branch
-        restructure_branch(scfg.region)
-        for region in _iter_subregions(scfg):
-            restructure_branch(region)
-        return ByteFlow(bc=self.bc, scfg=scfg)
+        self._restructure_branch()
 
 
 def _iter_subregions(scfg: SCFG) -> Generator[RegionBlock, SCFG, None]:

numba_rvsdg/core/datastructures/flow_info.py

@@ -126,19 +126,19 @@ def build_basicblocks(
 
         for begin, end in zip(offsets, [*offsets[1:], end_offset]):
             name = names[begin]
-            targets: Tuple[str, ...]
+            targets: list[str]
             term_offset = _prev_inst_offset(end)
             if term_offset not in self.jump_insts:
                 # implicit jump
-                targets = (names[end],)
+                targets = [names[end]]
             else:
-                targets = tuple(names[o] for o in self.jump_insts[term_offset])
+                targets = [names[o] for o in self.jump_insts[term_offset]]
             block = PythonBytecodeBlock(
                 name=name,
                 begin=begin,
                 end=end,
                 _jump_targets=targets,
-                backedges=(),
+                backedges=[],
             )
             scfg.add_block(block)
         return scfg

numba_rvsdg/core/datastructures/scfg.py

@@ -527,14 +527,14 @@ def insert_block(
         # TODO: needs a diagram and documentaion
         # initialize new block
         new_block = block_type(
-            name=new_name, _jump_targets=tuple(successors), backedges=tuple()
+            name=new_name, _jump_targets=list(successors), backedges=[]
         )
         # add block to self
         self.add_block(new_block)
         # Replace any arcs from any of predecessors to any of successors with
         # an arc through the inserted block instead.
         for name in predecessors:
-            block = self.graph.pop(name)
+            block = self.graph[name]
             jt = list(block.jump_targets)
             if successors:
                 for s in successors:
@@ -545,7 +545,7 @@ def insert_block(
                             jt.pop(jt.index(s))
             else:
                 jt.append(new_name)
-            self.add_block(block.replace_jump_targets(jump_targets=tuple(jt)))
+            block.replace_jump_targets(jump_targets=jt)
 
     def insert_SyntheticExit(
         self,
@@ -639,8 +639,8 @@ def insert_block_and_control_blocks(
                 variable_assignment[branch_variable] = branch_variable_value
                 synth_assign_block = SyntheticAssignment(
                     name=synth_assign,
-                    _jump_targets=(new_name,),
-                    backedges=(),
+                    _jump_targets=[new_name],
+                    backedges=[],
                     variable_assignment=variable_assignment,
                 )
                 # add block
@@ -652,16 +652,12 @@ def insert_block_and_control_blocks(
                 # replace previous successor with synth_assign
                 jt[jt.index(s)] = synth_assign
             # finally, replace the jump_targets
-            self.add_block(
-                self.graph.pop(name).replace_jump_targets(
-                    jump_targets=tuple(jt)
-                )
-            )
+            self.graph[name].replace_jump_targets(jump_targets=jt)
         # initialize new block, which will hold the branching table
         new_block = SyntheticHead(
             name=new_name,
-            _jump_targets=tuple(successors),
-            backedges=tuple(),
+            _jump_targets=list(successors),
+            backedges=[],
             variable=branch_variable,
             branch_value_table=branch_value_table,
         )
@@ -1079,7 +1075,7 @@ def to_yaml(scfg: "SCFG") -> str:
         return ys
 
     @staticmethod
-    def to_dict(scfg: "SCFG") -> Dict[str, Dict[str, Any]]:
+    def to_dict(scfg: "SCFG") -> Dict[str, Dict[str, BasicBlock]]:
         """Helper method to convert the SCFG object to a dictionary
         representation.
 
@@ -1109,23 +1105,25 @@ def reverse_lookup(value: type) -> str:
                 raise TypeError("Block type not found.")
 
         seen = set()
-        q: Set[Tuple[str, BasicBlock]] = set()
+        q: Set[str] = set()
         # Order of elements doesn't matter since they're going to
         # be sorted at the end.
-        q.update(scfg.graph.items())
+        graph_dict = {}
+        q.update(set(scfg.graph.keys()))
+        graph_dict.update(scfg.graph)
 
         while q:
-            key, value = q.pop()
+            key = q.pop()
+            value = scfg.graph[key]
             if key in seen:
                 continue
             seen.add(key)
 
             block_type = reverse_lookup(type(value))
             blocks[key] = {"type": block_type}
             if isinstance(value, RegionBlock):
-                assert value.subregion is not None
-                assert value.parent_region is not None
-                q.update(value.subregion.graph.items())
+                q.update(value.subregion.graph.keys())
+                graph_dict.update(value.subregion.graph)
                 blocks[key]["kind"] = value.kind
                 blocks[key]["contains"] = sorted(
                     [idx.name for idx in value.subregion.graph.values()]
@@ -1211,14 +1209,14 @@ def extract_block_info(
             List of backedges of the requested block.
         """
         block_info = blocks[current_name].copy()
-        block_edges = tuple(block_ref_dict[idx] for idx in edges[current_name])
+        block_edges = [block_ref_dict[idx] for idx in edges[current_name]]
 
         if backedges.get(current_name):
-            block_backedges = tuple(
+            block_backedges = [
                 block_ref_dict[idx] for idx in backedges[current_name]
-            )
+            ]
         else:
-            block_backedges = ()
+            block_backedges = []
 
         block_type = block_info.pop("type")