Store dict to block_lists in CFG

src/parser/cfg.py

@@ -1,8 +1,8 @@
 import json
 from pathlib import Path
 from parser.cfg_object import CFGObject
-from parser.cfg_block import CFGBlock
-from typing import Dict, List
+from parser.cfg_block_list import CFGBlockList
+from typing import Dict, List, Optional
 
 
 def store_sfs_json(blocks: List[Dict], final_path: Path) -> None:
@@ -18,18 +18,30 @@ def store_sfs_json(blocks: List[Dict], final_path: Path) -> None:
 class CFG:
     def __init__(self, nodeType: str):
         self.nodeType = nodeType
-        self.objectCFG : Dict[str, CFGObject] = {}
-        self.subObjects : CFG = None
+        self.objectCFG: Dict[str, CFGObject] = {}
+        self.subObjects: Optional[CFG] = None
 
-    def add_object(self, name:str, cfg_object: CFGObject) -> None:
+        # Points each object/function/subobject name to its block list
+        self.block_list: Dict[str, CFGBlockList] = {}
+
+    def add_object(self, name: str, cfg_object: CFGObject) -> None:
         self.objectCFG[name] = cfg_object
 
+        # Once an object is stored, we keep a dictionary with all the blocklists and the name
+        # of the corresponding structure
+        self.block_list[name] = cfg_object.blocks
+        for function_name, cfg_function in cfg_object.functions.items():
+            self.block_list[function_name] = cfg_function.blocks
+
     def get_object(self, name:str) -> CFGObject:
         return self.objectCFG[name]
 
     def set_subobject(self, subobject: 'CFG'):
         self.subObjects = subobject
 
+        # Add all the definitions in the subobject
+        self.block_list.update(subobject.block_list)
+
     def get_subobject(self) -> 'CFG':
         return self.subObjects
 
@@ -41,7 +53,12 @@ def build_spec_for_objects(self):
             object_dict[o] = specs
 
             functions_dict[o] = self.objectCFG[o].build_spec_for_functions()
-
+
+        if self.subObjects is not None:
+            subobject_dict, subfunction_dict = self.subObjects.build_spec_for_objects()
+            object_dict.update(subobject_dict)
+            functions_dict.update(subfunction_dict)
+
         return object_dict, functions_dict
 
     def get_as_json(self):

src/parser/cfg_block_list.py

@@ -16,10 +16,15 @@ class CFGBlockList:
     def __init__(self):
         self.blocks: Dict[str, CFGBlock] = {}
         self.graph = None
+        self.start_block = None
 
     def add_block(self, block: CFGBlock) -> None:
         block_id = block.get_block_id()
 
+        # Assuming the first block corresponds to the entry point
+        if not self.blocks:
+            self.start_block = block_id
+
         if block_id in self.blocks:
             if block_id in self.blocks:
                 print("WARNING: You are overwritting an existing block")

src/parser/parser.py

@@ -151,13 +151,16 @@ def parser_CFG_from_JSON(json_dict: Dict):
     for obj in object_keys:
         json_object = json_dict.get(obj,False)
         cfg_object = parse_object(obj,json_object)
-        cfg.add_object(obj,cfg_object)
 
         json_functions = json_object.get("functions", {})
         for f in json_functions:
             obj_function = parse_function(f, json_functions[f])
             cfg_object.add_function(obj_function)
 
+        # Important: add the object already initialized with the functions, so that we can construct
+        # link the corresponding block lists in the CFG
+        cfg.add_object(obj, cfg_object)
+
         cfg_object.identify_function_calls_in_blocks()
     # obj_name = obj.get("name")
     # cfg.add_object_name(obj_name)