[WIP] Added region aware testing

numba_rvsdg/core/datastructures/basic_block.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass, replace
 
 from numba_rvsdg.core.utils import _next_inst_offset
+from numba_rvsdg.core.datastructures import block_names
 
 
 @dataclass(frozen=True)
@@ -384,3 +385,19 @@ def replace_exiting(self, new_exiting):
             The new exiting block of the region represented by the RegionBlock.
         """
         object.__setattr__(self, "exiting", new_exiting)
+
+
+block_type_names = {
+    block_names.BASIC: BasicBlock,
+    block_names.PYTHON_BYTECODE: PythonBytecodeBlock,
+    block_names.SYNTH_HEAD: SyntheticHead,
+    block_names.SYNTH_BRANCH: SyntheticBranch,
+    block_names.SYNTH_TAIL: SyntheticTail,
+    block_names.SYNTH_EXIT: SyntheticExit,
+    block_names.SYNTH_ASSIGN: SyntheticAssignment,
+    block_names.SYNTH_RETURN: SyntheticReturn,
+    block_names.SYNTH_EXIT_LATCH: SyntheticExitingLatch,
+    block_names.SYNTH_EXIT_BRANCH: SyntheticExitBranch,
+    block_names.SYNTH_FILL: SyntheticFill,
+    block_names.REGION: RegionBlock,
+}

numba_rvsdg/core/datastructures/block_names.py

@@ -10,3 +10,21 @@
 SYNTH_RETURN = "synth_return"
 SYNTH_EXIT_LATCH = "synth_exit_latch"
 SYNTH_FILL = "synth_fill"
+SYNTH_EXIT_BRANCH = "synth_exit_branch"
+
+REGION = "region"
+
+block_types = {
+    BASIC,
+    PYTHON_BYTECODE,
+    SYNTH_HEAD,
+    SYNTH_BRANCH,
+    SYNTH_TAIL,
+    SYNTH_EXIT,
+    SYNTH_ASSIGN,
+    SYNTH_RETURN,
+    SYNTH_EXIT_LATCH,
+    SYNTH_EXIT_BRANCH,
+    SYNTH_FILL,
+    REGION,
+}

numba_rvsdg/core/datastructures/scfg.py

@@ -1,6 +1,5 @@
 import dis
 import yaml
-from textwrap import dedent
 from typing import Set, Tuple, Dict, List, Iterator
 from dataclasses import dataclass, field
 from collections import deque
@@ -15,9 +14,22 @@
     SyntheticTail,
     SyntheticReturn,
     SyntheticFill,
+    PythonBytecodeBlock,
     RegionBlock,
+    block_type_names,
+)
+from numba_rvsdg.core.datastructures.block_names import (
+    block_types,
+    PYTHON_BYTECODE,
+    SYNTH_HEAD,
+    SYNTH_BRANCH,
+    SYNTH_TAIL,
+    SYNTH_EXIT,
+    SYNTH_ASSIGN,
+    SYNTH_RETURN,
+    SYNTH_EXIT_LATCH,
+    SYNTH_EXIT_BRANCH,
 )
-from numba_rvsdg.core.datastructures import block_names
 
 
 @dataclass(frozen=True)
@@ -604,9 +616,7 @@ def insert_block_and_control_blocks(
             # predecessors to a successor and insert it between the predecessor
             # and the newly created block
             for s in set(jt).intersection(successors):
-                synth_assign = self.name_gen.new_block_name(
-                    block_names.SYNTH_ASSIGN
-                )
+                synth_assign = self.name_gen.new_block_name(SYNTH_ASSIGN)
                 variable_assignment = {}
                 variable_assignment[branch_variable] = branch_variable_value
                 synth_assign_block = SyntheticAssignment(
@@ -652,9 +662,7 @@ def join_returns(self):
         ]
         # close if more than one is found
         if len(return_nodes) > 1:
-            return_solo_name = self.name_gen.new_block_name(
-                block_names.SYNTH_RETURN
-            )
+            return_solo_name = self.name_gen.new_block_name(SYNTH_RETURN)
             self.insert_SyntheticReturn(
                 return_solo_name, return_nodes, tuple()
             )
@@ -685,29 +693,21 @@ def join_tails_and_exits(self, tails: Set[str], exits: Set[str]):
         if len(tails) == 1 and len(exits) == 2:
             # join only exits
             solo_tail_name = next(iter(tails))
-            solo_exit_name = self.name_gen.new_block_name(
-                block_names.SYNTH_EXIT
-            )
+            solo_exit_name = self.name_gen.new_block_name(SYNTH_EXIT)
             self.insert_SyntheticExit(solo_exit_name, tails, exits)
             return solo_tail_name, solo_exit_name
 
         if len(tails) >= 2 and len(exits) == 1:
             # join only tails
-            solo_tail_name = self.name_gen.new_block_name(
-                block_names.SYNTH_TAIL
-            )
+            solo_tail_name = self.name_gen.new_block_name(SYNTH_TAIL)
             solo_exit_name = next(iter(exits))
             self.insert_SyntheticTail(solo_tail_name, tails, exits)
             return solo_tail_name, solo_exit_name
 
         if len(tails) >= 2 and len(exits) >= 2:
             # join both tails and exits
-            solo_tail_name = self.name_gen.new_block_name(
-                block_names.SYNTH_TAIL
-            )
-            solo_exit_name = self.name_gen.new_block_name(
-                block_names.SYNTH_EXIT
-            )
+            solo_tail_name = self.name_gen.new_block_name(SYNTH_TAIL)
+            solo_exit_name = self.name_gen.new_block_name(SYNTH_EXIT)
             self.insert_SyntheticTail(solo_tail_name, tails, exits)
             self.insert_SyntheticExit(solo_exit_name, {solo_tail_name}, exits)
             return solo_tail_name, solo_exit_name
@@ -731,7 +731,7 @@ def bcmap_from_bytecode(bc: dis.Bytecode):
         return {inst.offset: inst for inst in bc}
 
     @staticmethod
-    def from_yaml(yaml_string):
+    def from_yaml(yaml_string: str):
         """Static method that creates an SCFG object from a YAML
         representation.
 
@@ -781,23 +781,110 @@ def from_dict(graph_dict: dict):
             Dictionary of block names in YAML string corresponding to their
             representation/unique name IDs in the SCFG.
         """
-        scfg_graph = {}
         name_gen = NameGenerator()
-        block_dict = {}
-        for index in graph_dict.keys():
-            block_dict[index] = name_gen.new_block_name(block_names.BASIC)
-        for index, attributes in graph_dict.items():
-            jump_targets = attributes["jt"]
-            backedges = attributes.get("be", ())
-            name = block_dict[index]
-            block = BasicBlock(
-                name=name,
-                backedges=tuple(block_dict[idx] for idx in backedges),
-                _jump_targets=tuple(block_dict[idx] for idx in jump_targets),
-            )
-            scfg_graph[name] = block
-        scfg = SCFG(scfg_graph, name_gen=name_gen)
-        return scfg, block_dict
+        block_ref_dict = {}
+
+        blocks = graph_dict["blocks"]
+        edges = graph_dict["edges"]
+        backedges = graph_dict["backedges"]
+        if backedges is None:
+            backedges = {}
+
+        for key, block in blocks.items():
+            assert block["type"] in block_types
+            block_ref_dict[key] = key
+
+        # Find head of the graph, i.e. node which isn't in anyones contains
+        # and no edges point towards it (backedges are allowed)
+        heads = set(blocks.keys())
+        for block_name, block_data in blocks.items():
+            if block_data.get("contains"):
+                heads.difference_update(block_data["contains"])
+            jump_targets = set(edges[block_name])
+            if backedges.get(block_name):
+                jump_targets.difference_update(set(backedges[block_name]))
+            heads.difference_update(jump_targets)
+        assert len(heads) > 0
+
+        seen = set()
+
+        def make_scfg(curr_heads: set, exiting: str = None):
+            scfg_graph = {}
+            queue = curr_heads
+            while queue:
+                current_name = queue.pop()
+                if current_name in seen:
+                    continue
+                seen.add(current_name)
+
+                block_info = blocks[current_name]
+                block_edges = tuple(
+                    block_ref_dict[idx] for idx in edges[current_name]
+                )
+                if backedges and backedges.get(current_name):
+                    block_backedges = tuple(
+                        block_ref_dict[idx] for idx in backedges[current_name]
+                    )
+                else:
+                    block_backedges = ()
+                block_type = block_info.get("type")
+                block_class = block_type_names[block_type]
+                if block_type == "region":
+                    scfg = make_scfg(
+                        {block_info["header"]}, block_info["exiting"]
+                    )
+                    block = RegionBlock(
+                        name=current_name,
+                        _jump_targets=block_edges,
+                        backedges=block_backedges,
+                        kind=block_info["kind"],
+                        header=block_info["header"],
+                        exiting=block_info["exiting"],
+                        subregion=scfg,
+                    )
+                elif block_type in [
+                    SYNTH_BRANCH,
+                    SYNTH_HEAD,
+                    SYNTH_EXIT_LATCH,
+                    SYNTH_EXIT_BRANCH,
+                ]:
+                    block = block_class(
+                        name=current_name,
+                        backedges=block_backedges,
+                        _jump_targets=block_edges,
+                        branch_value_table=block_info["branch_value_table"],
+                        variable=block_info["variable"],
+                    )
+                elif block_type in [SYNTH_ASSIGN]:
+                    block = SyntheticAssignment(
+                        name=current_name,
+                        _jump_targets=block_edges,
+                        backedges=block_backedges,
+                        variable_assignment=block_info["variable_assignment"],
+                    )
+                elif block_type in [PYTHON_BYTECODE]:
+                    block = PythonBytecodeBlock(
+                        name=current_name,
+                        _jump_targets=block_edges,
+                        backedges=block_backedges,
+                        begin=block_info["begin"],
+                        end=block_info["end"],
+                    )
+                else:
+                    block = block_class(
+                        name=current_name,
+                        backedges=block_backedges,
+                        _jump_targets=block_edges,
+                    )
+                scfg_graph[current_name] = block
+                if current_name != exiting:
+                    queue.update(edges[current_name])
+
+            scfg = SCFG(scfg_graph, name_gen=name_gen)
+            return scfg
+
+        scfg = make_scfg(heads)
+        return scfg, block_ref_dict
 
     def to_yaml(self):
         """Converts the SCFG object to a YAML string representation.
@@ -813,23 +900,33 @@ def to_yaml(self):
             A YAML string representing the SCFG.
         """
         # Convert to yaml
-        scfg_graph = self.graph
         yaml_string = """"""
 
-        for key, value in scfg_graph.items():
-            jump_targets = [i for i in value._jump_targets]
-            jump_targets = str(jump_targets).replace("'", '"')
-            back_edges = [i for i in value.backedges]
-            jump_target_str = f"""
-                "{key}":
-                    jt: {jump_targets}"""
-
-            if back_edges:
-                back_edges = str(back_edges).replace("'", '"')
-                jump_target_str += f"""
-                    be: {back_edges}"""
-            yaml_string += dedent(jump_target_str)
-
+        graph_dict = self.to_dict()
+
+        blocks = graph_dict["blocks"]
+        edges = graph_dict["edges"]
+        backedges = graph_dict["backedges"]
+
+        yaml_string += "\nblocks:"
+        for _block in sorted(blocks.keys()):
+            yaml_string += f"""
+    '{_block}':"""
+            block_dict = blocks[_block]
+            for key, value in block_dict.items():
+                yaml_string += f"""
+        {key}: {value}"""
+
+        yaml_string += "\nedges:"
+        for _block in sorted(blocks.keys()):
+            yaml_string += f"""
+    '{_block}': {edges[_block]}"""
+
+        yaml_string += "\nbackedges:"
+        for _block in sorted(blocks.keys()):
+            if backedges[_block]:
+                yaml_string += f"""
+    '{_block}': {backedges[_block]}"""
         return yaml_string
 
     def to_dict(self):
@@ -845,14 +942,46 @@ def to_dict(self):
         graph_dict: Dict[Dict[...]]
             A dictionary representing the SCFG.
         """
-        scfg_graph = self.graph
-        graph_dict = {}
-        for key, value in scfg_graph.items():
-            curr_dict = {}
-            curr_dict["jt"] = [i for i in value._jump_targets]
-            if value.backedges:
-                curr_dict["be"] = [i for i in value.backedges]
-            graph_dict[key] = curr_dict
+        blocks = {}
+        edges = {}
+        backedges = {}
+
+        def reverse_lookup(value):
+            for k, v in block_type_names.items():
+                if v == value:
+                    return k
+            else:
+                raise TypeError("Block type not found.")
+
+        for key, value in self:
+            block_type = reverse_lookup(type(value))
+            blocks[key] = {"type": block_type}
+            if block_type == "region":
+                blocks[key]["kind"] = value.kind
+                blocks[key]["contains"] = sorted(
+                    [idx.name for idx in value.subregion.graph.values()]
+                )
+                blocks[key]["header"] = value.header
+                blocks[key]["exiting"] = value.exiting
+                blocks[key]["parent_region"] = value.parent_region.name
+            elif block_type in [
+                SYNTH_BRANCH,
+                SYNTH_HEAD,
+                SYNTH_EXIT_LATCH,
+                SYNTH_EXIT_BRANCH,
+            ]:
+                blocks[key]["branch_value_table"] = value.branch_value_table
+                blocks[key]["variable"] = value.variable
+            elif block_type in [SYNTH_ASSIGN]:
+                blocks[key]["variable_assignment"] = value.variable_assignment
+            elif block_type in [PYTHON_BYTECODE]:
+                blocks[key]["begin"] = value.begin
+                blocks[key]["end"] = value.end
+            edges[key] = sorted([i for i in value._jump_targets])
+            backedges[key] = sorted([i for i in value.backedges])
+
+        graph_dict = {"blocks": blocks, "edges": edges, "backedges": backedges}
+
         return graph_dict
 
     def view(self, name: str = None):