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 numba_rvsdg.core.datastructures.scfg import SCFG
@@ -63,87 +62,48 @@ def from_bytecode(code) -> "ByteFlow":
     def _join_returns(self):
         """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):
         """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):
         """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):
         """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"):

numba_rvsdg/core/datastructures/flow_info.py

@@ -128,15 +128,15 @@ def build_basicblocks(self: "FlowInfo", end_offset=None) -> "SCFG":
             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

@@ -506,14 +506,14 @@ def insert_block(
         # TODO: needs a diagram and documentaion
         # initialize new block
         new_block = block_type(
-            name=new_name, _jump_targets=successors, backedges=set()
+            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:
@@ -524,7 +524,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,
@@ -618,8 +618,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
@@ -631,16 +631,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=set(),
+            _jump_targets=list(successors),
+            backedges=[],
             variable=branch_variable,
             branch_value_table=branch_value_table,
         )
@@ -1084,18 +1080,22 @@ def reverse_lookup(value: type):
         q = 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(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):
-                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()]
@@ -1177,14 +1177,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")
 

numba_rvsdg/core/transformations.py

@@ -58,9 +58,7 @@ def loop_restructure_helper(scfg: SCFG, loop: Set[str]):
         and len(exiting_blocks) == 1
         and backedge_blocks[0] == next(iter(exiting_blocks))
     ):
-        scfg.add_block(
-            scfg.graph.pop(backedge_blocks[0]).declare_backedge(loop_head)
-        )
+        scfg.graph[backedge_blocks[0]].declare_backedge(loop_head)
         return
 
     # The synthetic exiting latch and synthetic exit need to be created
@@ -146,8 +144,8 @@ def reverse_lookup(d, value):
                     # Create the actual control variable block
                     synth_assign_block = SyntheticAssignment(
                         name=synth_assign,
-                        _jump_targets=(synth_exiting_latch,),
-                        backedges=(),
+                        _jump_targets=[synth_exiting_latch],
+                        backedges=[],
                         variable_assignment=variable_assignment,
                     )
                     # Insert the assignment to the scfg
@@ -177,20 +175,18 @@ def reverse_lookup(d, value):
                     # that point to the headers, no need to add a backedge,
                     # since it will be contained in the SyntheticExitingLatch
                     # later on.
-                    block = scfg.graph.pop(name)
+                    block = scfg.graph[name]
                     jts = list(block.jump_targets)
                     for h in headers:
                         if h in jts:
                             jts.remove(h)
-                    scfg.add_block(
-                        block.replace_jump_targets(jump_targets=tuple(jts))
-                    )
+                    block.replace_jump_targets(jump_targets=jts)
                     # Setup the assignment block and initialize it with the
                     # correct jump_targets and variable assignment.
                     synth_assign_block = SyntheticAssignment(
                         name=synth_assign,
-                        _jump_targets=(synth_exiting_latch,),
-                        backedges=(),
+                        _jump_targets=[synth_exiting_latch],
+                        backedges=[],
                         variable_assignment=variable_assignment,
                     )
                     # Add the new block to the SCFG
@@ -199,22 +195,19 @@ def reverse_lookup(d, value):
                     new_jt[new_jt.index(jt)] = synth_assign
             # finally, replace the jump_targets for this block with the new
             # ones
-            scfg.add_block(
-                scfg.graph.pop(name).replace_jump_targets(
-                    jump_targets=tuple(new_jt)
-                )
-            )
+            scfg.graph[name].replace_jump_targets(jump_targets=new_jt)
+
     # Add any new blocks to the loop.
     loop.update(new_blocks)
 
     # Insert the exiting latch, add it to the loop and to the graph.
     synth_exiting_latch_block = SyntheticExitingLatch(
         name=synth_exiting_latch,
-        _jump_targets=(
+        _jump_targets=[
             synth_exit if needs_synth_exit else next(iter(exit_blocks)),
             loop_head,
-        ),
-        backedges=(loop_head,),
+        ],
+        backedges=[loop_head],
         variable=backedge_variable,
         branch_value_table=backedge_value_table,
     )
@@ -225,8 +218,8 @@ def reverse_lookup(d, value):
     if needs_synth_exit:
         synth_exit_block = SyntheticExitBranch(
             name=synth_exit,
-            _jump_targets=tuple(exit_blocks),
-            backedges=(),
+            _jump_targets=list(exit_blocks),
+            backedges=[],
             variable=exit_variable,
             branch_value_table=exit_value_table,
         )
@@ -329,29 +322,23 @@ def update_exiting(
 ):
     # Recursively updates the exiting blocks of a regionblock
     region_exiting = region_block.exiting
-    region_exiting_block: BasicBlock = region_block.subregion.graph.pop(
+    region_exiting_block: BasicBlock = region_block.subregion.graph[
         region_exiting
-    )
+    ]
     jt = list(region_exiting_block._jump_targets)
     for idx, s in enumerate(jt):
         if s is new_region_header:
             jt[idx] = new_region_name
-    region_exiting_block = region_exiting_block.replace_jump_targets(
-        jump_targets=tuple(jt)
-    )
+    region_exiting_block.replace_jump_targets(jump_targets=jt)
     be = list(region_exiting_block.backedges)
     for idx, s in enumerate(be):
         if s is new_region_header:
             be[idx] = new_region_name
-    region_exiting_block = region_exiting_block.replace_backedges(
-        backedges=tuple(be)
-    )
+    region_exiting_block.replace_backedges(backedges=be)
     if isinstance(region_exiting_block, RegionBlock):
         region_exiting_block = update_exiting(
             region_exiting_block, new_region_header, new_region_name
         )
-    region_block.subregion.add_block(region_exiting_block)
-    return region_block
 
 
 def extract_region(
@@ -381,27 +368,26 @@ def extract_region(
             # the SCFG represents should not be the meta region.
             assert scfg.region.kind != "meta"
             continue
-        entry = scfg.graph.pop(name)
+        entry = scfg.graph[name]
         jt = list(entry._jump_targets)
         for idx, s in enumerate(jt):
             if s is region_header:
                 jt[idx] = region_name
-        entry = entry.replace_jump_targets(jump_targets=tuple(jt))
+        entry.replace_jump_targets(jump_targets=jt)
         be = list(entry.backedges)
         for idx, s in enumerate(be):
             if s is region_header:
                 be[idx] = region_name
-        entry = entry.replace_backedges(backedges=tuple(be))
+        entry.replace_backedges(backedges=be)
         # If the entry itself is a region, update it's
         # exiting blocks too, recursively
         if isinstance(entry, RegionBlock):
             entry = update_exiting(entry, region_header, region_name)
-        scfg.add_block(entry)
 
     region = RegionBlock(
         name=region_name,
         _jump_targets=scfg[region_exiting].jump_targets,
-        backedges=(),
+        backedges=[],
         kind=region_kind,
         header=region_header,
         subregion=head_subgraph,