From dae773cf06b92a201590f1f6963b1267ca24d1ed Mon Sep 17 00:00:00 2001
From: kc611 <ckaustubhm06@gmail.com>
Date: Thu, 29 Jun 2023 18:15:33 +0530
Subject: [PATCH] Added region identification and printing capabilities to
 to_yaml and to_dict

---
 .../core/datastructures/basic_block.py        | 17 +++++++
 .../core/datastructures/block_names.py        |  4 ++
 numba_rvsdg/core/datastructures/scfg.py       | 51 ++++++++++++++-----
 3 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/numba_rvsdg/core/datastructures/basic_block.py b/numba_rvsdg/core/datastructures/basic_block.py
index 263e09c..393b105 100644
--- a/numba_rvsdg/core/datastructures/basic_block.py
+++ b/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,
+}
diff --git a/numba_rvsdg/core/datastructures/block_names.py b/numba_rvsdg/core/datastructures/block_names.py
index a36db2c..030e924 100644
--- a/numba_rvsdg/core/datastructures/block_names.py
+++ b/numba_rvsdg/core/datastructures/block_names.py
@@ -10,6 +10,9 @@
 SYNTH_RETURN = "synth_return"
 SYNTH_EXIT_LATCH = "synth_exit_latch"
 SYNTH_FILL = "synth_fill"
+SYNTH_EXIT_BRANCH = "synth_exit_branch"
+
+REGION = "region"
 
 all_block_names = {
     BASIC,
@@ -22,4 +25,5 @@
     SYNTH_RETURN,
     SYNTH_EXIT_LATCH,
     SYNTH_FILL,
+    REGION,
 }
diff --git a/numba_rvsdg/core/datastructures/scfg.py b/numba_rvsdg/core/datastructures/scfg.py
index b28ad0b..af39e65 100644
--- a/numba_rvsdg/core/datastructures/scfg.py
+++ b/numba_rvsdg/core/datastructures/scfg.py
@@ -15,6 +15,7 @@
     SyntheticReturn,
     SyntheticFill,
     RegionBlock,
+    block_type_names,
 )
 from numba_rvsdg.core.datastructures import block_names
 
@@ -787,7 +788,6 @@ def from_dict(graph_dict: dict):
         blocks = graph_dict["blocks"]
         edges = graph_dict["edges"]
         backedges = graph_dict["backedges"]
-        graph_dict["regions"]
 
         for key, block in blocks.items():
             assert block["type"] in block_names.all_block_names
@@ -825,15 +825,29 @@ def to_yaml(self):
             A YAML string representing the SCFG.
         """
         # Convert to yaml
-        scfg_graph = self.graph
         yaml_string = """"""
 
         blocks = {}
         edges = {}
         backedges = {}
 
-        for key, value in scfg_graph.items():
-            blocks[key] = {"type": "basic"}
+        # This does a dictionary reverse lookup, to determine the key for a given
+        # value.
+        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"] = [
+                    idx.name for idx in value.subregion.graph.values()
+                ]
             edges[key] = [i for i in value._jump_targets]
             backedges[key] = [i for i in value.backedges]
 
@@ -844,6 +858,11 @@ def to_yaml(self):
             block_type = blocks[_block]["type"]
             yaml_string += f"""
         type: {block_type}"""
+            if block_type == "region":
+                yaml_string += f"""
+        kind: {blocks[_block]["kind"]}"""
+                yaml_string += f"""
+        contains: {blocks[_block]["contains"]}"""
 
         yaml_string += "\nedges:"
         for _block in blocks.keys():
@@ -855,7 +874,6 @@ def to_yaml(self):
             if backedges[_block]:
                 yaml_string += f"""
     {_block}: {backedges[_block]}"""
-        yaml_string += "\nregions:"
         return yaml_string
 
     def to_dict(self):
@@ -876,19 +894,26 @@ def to_dict(self):
         blocks = {}
         edges = {}
         backedges = {}
-        regions = {}
+
+        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 scfg_graph.items():
-            blocks[key] = {"type": "basic"}
+            block_type = reverse_lookup(type(value))
+            blocks[key] = {"type": block_type}
+            if block_type == "region":
+                blocks[key]["kind"] = value.kind
+                blocks[key]["contains"] = [
+                    idx.name for idx in value.subregion.graph.values()
+                ]
             edges[key] = [i for i in value._jump_targets]
             backedges[key] = [i for i in value.backedges]
 
-        graph_dict = {
-            "blocks": blocks,
-            "edges": edges,
-            "backedges": backedges,
-            "regions": regions,
-        }
+        graph_dict = {"blocks": blocks, "edges": edges, "backedges": backedges}
 
         return graph_dict