From 9514a3afeaf4b07dd6960aa0a37b0fbe8e0fabd0 Mon Sep 17 00:00:00 2001
From: alexcere <48130030+alexcere@users.noreply.github.com>
Date: Tue, 10 Sep 2024 16:26:30 +0200
Subject: [PATCH] Modify CFG block list dict accordingly

---
 src/liveness/layout_generation.py    | 22 +++++++++-------------
 src/parser/cfg.py                    |  2 ++
 src/parser/optimizable_block_list.py |  6 ++++++
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/liveness/layout_generation.py b/src/liveness/layout_generation.py
index fbe43501..994216cf 100644
--- a/src/liveness/layout_generation.py
+++ b/src/liveness/layout_generation.py
@@ -6,7 +6,7 @@
 """
 import heapq
 import itertools
-from typing import Dict, List, Type, Any, Set, Tuple
+from typing import Dict, List, Type, Any, Set, Tuple, Optional
 import networkx as nx
 from pathlib import Path
 
@@ -179,15 +179,15 @@ def var_order_repr(block_name: str, var_info: Dict[str, int]):
     return '\n'.join(text_format)
 
 
-def print_stacks(block_name: str, block: CFGBlock):
-    text_format = [f"{block_name}:", f"Src: {block.input_stack}", f"Tgt: {block.output_stack}"]
+def print_stacks(block_name: str, json_dict: Dict[str, Any]):
+    text_format = [f"{block_name}:", f"Src: {json_dict[block_name]['src_ws']}", f"Tgt: {json_dict[block_name]['tgt_ws']}"]
     return '\n'.join(text_format)
 
 
 class LayoutGeneration:
 
     def __init__(self, object_id, block_list: CFGBlockList, liveness_info: Dict[str, LivenessAnalysisInfo], name: Path,
-                 cfg_graph: nx.Graph = None):
+                 cfg_graph: Optional[nx.Graph] = None):
         self._id = object_id
         self._block_list = block_list
         self._liveness_info = liveness_info
@@ -309,10 +309,6 @@ def _construct_code_from_block(self, block: CFGBlock, input_stacks: Dict[str, Li
                 block_json[subblock_name]["src_ws"] = input_stack
                 block_json[subblock_name]["tgt_ws"] = output_stack
 
-        # TODO: temporal hack to output the input and output stacks
-        block.input_stack = input_stack
-        block.output_stack = output_stack
-
         return block_json
 
     def _construct_code_from_block_list(self):
@@ -341,8 +337,8 @@ def _construct_code_from_block_list(self):
             # Retrieve the block
             current_block = self._block_list.get_block(block_name)
 
-            block_specification = self._construct_code_from_block(current_block, input_stacks, output_stacks, combined_stacks)
-
+            block_specification = self._construct_code_from_block(current_block, input_stacks,
+                                                                  output_stacks, combined_stacks)
             json_info.update(block_specification)
 
             successors = [possible_successor for possible_successor in
@@ -359,9 +355,9 @@ def build_layout(self):
         """
         json_info = self._construct_code_from_block_list()
 
-        renamed_graph = information_on_graph(self._cfg_graph, {block_name: print_stacks(block_name, block)
-                                                               for block_name, block in
-                                                               self._block_list.blocks.items()})
+        renamed_graph = information_on_graph(self._cfg_graph, {block_name: print_stacks(block_name, json_info[block_name])
+                                                               for block_name in
+                                                               self._block_list.blocks})
 
         nx.nx_agraph.write_dot(renamed_graph, Path(self._dir.parent).joinpath(self._dir.stem + "_stacks.dot"))
         return json_info
diff --git a/src/parser/cfg.py b/src/parser/cfg.py
index f8e5fed1..b9f4c400 100644
--- a/src/parser/cfg.py
+++ b/src/parser/cfg.py
@@ -103,10 +103,12 @@ def modify_cfg_block_list(self, f: Callable[[CFGBlockList], CFGBlockList]) -> No
         """
         for object_id, cfg_object in self.objectCFG.items():
             cfg_object.blocks = f(cfg_object.blocks)
+            self.block_list[object_id] = cfg_object.blocks
 
             # We also consider the information per function
             for function_name, cfg_function in cfg_object.functions.items():
                 cfg_function.blocks = f(cfg_function.blocks)
+                self.block_list[function_name] = cfg_function.blocks
 
             subobject = self.get_subobject()
 
diff --git a/src/parser/optimizable_block_list.py b/src/parser/optimizable_block_list.py
index 1879c92f..e15de56b 100644
--- a/src/parser/optimizable_block_list.py
+++ b/src/parser/optimizable_block_list.py
@@ -110,6 +110,12 @@ def compute_sub_block_list(block_list: CFGBlockList) -> CFGBlockList:
             new_block_list.add_block(cfg_block)
 
     update_edges_cfg(new_block_list, modified_blocks)
+
+    # Finally, we update the initial start block with the new format
+    modified_start_block = modified_blocks.get(block_list.start_block, None)
+    if modified_start_block is not None:
+        new_block_list.start_block = modified_start_block[0]
+
     return new_block_list