Merge pull request #336 from crytic/dev-filenames-lookup

Refactor filenames/filenames_lookup

Author: montyly

crytic_compile/compilation_unit.py

@@ -3,8 +3,7 @@
 """
 import uuid
 from collections import defaultdict
-from typing import TYPE_CHECKING, Dict, Set
-
+from typing import TYPE_CHECKING, Dict, Set, Optional
 
 from crytic_compile.compiler.compiler import CompilerVersion
 from crytic_compile.source_unit import SourceUnit
@@ -35,6 +34,9 @@ def __init__(self, crytic_compile: "CryticCompile", unique_id: str):
         # set containing all the filenames of this compilation unit
         self._filenames: Set[Filename] = set()
 
+        # mapping from absolute/relative/used to filename
+        self._filenames_lookup: Optional[Dict[str, Filename]] = None
+
         # compiler.compiler
         self._compiler_version: CompilerVersion = CompilerVersion(
             compiler="N/A", version="N/A", optimized=False
@@ -118,7 +120,6 @@ def create_source_unit(self, filename: Filename) -> SourceUnit:
         if not filename in self._source_units:
             source_unit = SourceUnit(self, filename)  # type: ignore
             self.filenames.add(filename)
-            self.crytic_compile.filenames.add(filename)
             self._source_units[filename] = source_unit
         return self._source_units[filename]
 
@@ -194,6 +195,36 @@ def relative_filename_from_absolute_filename(self, absolute_filename: str) -> st
             raise ValueError("f{absolute_filename} does not exist in {d}")
         return d_file[absolute_filename]
 
+    def filename_lookup(self, filename: str) -> Filename:
+        """Return a crytic_compile.naming.Filename from a any filename
+
+        Args:
+            filename (str): filename (used/absolute/relative)
+
+        Raises:
+            ValueError: If the filename is not in the project
+
+        Returns:
+            Filename: Associated Filename object
+        """
+        # pylint: disable=import-outside-toplevel
+        from crytic_compile.platform.truffle import Truffle
+
+        if isinstance(self.crytic_compile.platform, Truffle) and filename.startswith("project:/"):
+            filename = filename[len("project:/") :]
+
+        if self._filenames_lookup is None:
+            self._filenames_lookup = {}
+            for file in self._filenames:
+                self._filenames_lookup[file.absolute] = file
+                self._filenames_lookup[file.relative] = file
+                self._filenames_lookup[file.used] = file
+        if filename not in self._filenames_lookup:
+            raise ValueError(
+                f"{filename} does not exist in {[f.absolute for f in self._filenames_lookup.values()]}"
+            )
+        return self._filenames_lookup[filename]
+
     # endregion
     ###################################################################################
     ###################################################################################

crytic_compile/crytic_compile.py

@@ -19,7 +19,6 @@
 from crytic_compile.platform.all_export import PLATFORMS_EXPORT
 from crytic_compile.platform.solc import Solc
 from crytic_compile.platform.standard import export_to_standard
-from crytic_compile.platform.truffle import Truffle
 from crytic_compile.utils.naming import Filename
 from crytic_compile.utils.npm import get_package_name
 from crytic_compile.utils.zip import load_from_zip
@@ -78,12 +77,6 @@ def __init__(self, target: Union[str, AbstractPlatform], **kwargs: str) -> None:
         # dependencies is needed for platform conversion
         self._dependencies: Set = set()
 
-        # set containing all the filenames
-        self._filenames: Set[Filename] = set()
-
-        # mapping from absolute/relative/used to filename
-        self._filenames_lookup: Optional[Dict[str, Filename]] = None
-
         self._src_content: Dict = {}
 
         # Mapping each file to
@@ -152,27 +145,21 @@ def is_in_multiple_compilation_unit(self, contract: str) -> bool:
 
     ###################################################################################
     ###################################################################################
-    # region Filenames
+    # region Utils
     ###################################################################################
     ###################################################################################
-
     @property
     def filenames(self) -> Set[Filename]:
-        """All the project filenames
-
-        Returns:
-            Set[Filename]: Project's filenames
         """
-        return self._filenames
+        Return the set of all the filenames used
 
-    @filenames.setter
-    def filenames(self, all_filenames: Set[Filename]) -> None:
-        """Set the filenames
-
-        Args:
-            all_filenames (Set[Filename]): New filenames
+        Returns:
+             Set[Filename]: list of filenames
         """
-        self._filenames = all_filenames
+        filenames: Set[Filename] = set()
+        for compile_unit in self._compilation_units.values():
+            filenames = filenames.union(compile_unit.filenames)
+        return filenames
 
     def filename_lookup(self, filename: str) -> Filename:
         """Return a crytic_compile.naming.Filename from a any filename
@@ -186,21 +173,13 @@ def filename_lookup(self, filename: str) -> Filename:
         Returns:
             Filename: Associated Filename object
         """
+        for compile_unit in self.compilation_units.values():
+            try:
+                return compile_unit.filename_lookup(filename)
+            except ValueError:
+                pass
 
-        if isinstance(self.platform, Truffle) and filename.startswith("project:/"):
-            filename = filename[len("project:/") :]
-
-        if self._filenames_lookup is None:
-            self._filenames_lookup = {}
-            for file in self._filenames:
-                self._filenames_lookup[file.absolute] = file
-                self._filenames_lookup[file.relative] = file
-                self._filenames_lookup[file.used] = file
-        if filename not in self._filenames_lookup:
-            raise ValueError(
-                f"{filename} does not exist in {[f.absolute for f in self._filenames_lookup.values()]}"
-            )
-        return self._filenames_lookup[filename]
+        raise ValueError(f"{filename} does not exist")
 
     @property
     def dependencies(self) -> Set[str]:

crytic_compile/platform/standard.py

@@ -416,16 +416,15 @@ def _load_from_compile_0_0_1(crytic_compile: "CryticCompile", loaded_json: Dict)
                     crytic_compile.dependencies.add(filename.short)
                     crytic_compile.dependencies.add(filename.used)
 
-            compilation_unit.filenames = {
-                _convert_dict_to_filename(filename)
-                for filename in compilation_unit_json["filenames"]
-            }
+        compilation_unit.filenames = {
+            _convert_dict_to_filename(filename) for filename in compilation_unit_json["filenames"]
+        }
 
-            for path, ast in compilation_unit_json["asts"].items():
-                # The following might create lookup issue?
-                filename = crytic_compile.filename_lookup(path)
-                source_unit = compilation_unit.create_source_unit(filename)
-                source_unit.ast = ast
+        for path, ast in compilation_unit_json["asts"].items():
+            # The following might create lookup issue?
+            filename = crytic_compile.filename_lookup(path)
+            source_unit = compilation_unit.create_source_unit(filename)
+            source_unit.ast = ast
 
 
 def _load_from_compile_current(crytic_compile: "CryticCompile", loaded_json: Dict) -> None:
@@ -466,16 +465,15 @@ def _load_from_compile_current(crytic_compile: "CryticCompile", loaded_json: Dic
                     crytic_compile.dependencies.add(filename.short)
                     crytic_compile.dependencies.add(filename.used)
 
-            for path, ast in compilation_unit_json["asts"].items:
-                # The following might create lookup issue?
-                filename = convert_filename(path, lambda x: x, crytic_compile)
-                source_unit = compilation_unit.create_source_unit(filename)
-                source_unit.ast = ast
+        compilation_unit.filenames = {
+            _convert_dict_to_filename(filename) for filename in compilation_unit_json["filenames"]
+        }
 
-            compilation_unit.filenames = {
-                _convert_dict_to_filename(filename)
-                for filename in compilation_unit_json["filenames"]
-            }
+        for path, ast in compilation_unit_json["asts"].items:
+            # The following might create lookup issue?
+            filename = convert_filename(path, lambda x: x, crytic_compile)
+            source_unit = compilation_unit.create_source_unit(filename)
+            source_unit.ast = ast
 
 
 def load_from_compile(crytic_compile: "CryticCompile", loaded_json: Dict) -> Tuple[int, List[str]]:
@@ -490,7 +488,6 @@ def load_from_compile(crytic_compile: "CryticCompile", loaded_json: Dict) -> Tup
         Tuple[int, List[str]]: (underlying platform types, guessed unit tests)
     """
     crytic_compile.package_name = loaded_json.get("package", None)
-
     if "compilation_units" not in loaded_json:
         _load_from_compile_legacy1(crytic_compile, loaded_json)
 
@@ -502,10 +499,6 @@ def load_from_compile(crytic_compile: "CryticCompile", loaded_json: Dict) -> Tup
     else:
         _load_from_compile_current(crytic_compile, loaded_json)
 
-    # Set our filenames
-    for compilation_unit in crytic_compile.compilation_units.values():
-        crytic_compile.filenames |= set(compilation_unit.filenames)
-
     crytic_compile.working_dir = loaded_json["working_dir"]
 
     return loaded_json["type"], loaded_json.get("unit_tests", [])

crytic_compile/platform/waffle.py

@@ -192,7 +192,6 @@ def compile(self, crytic_compile: "CryticCompile", **kwargs: str) -> None:
             source_unit = compilation_unit.create_source_unit(filename)
 
             source_unit.ast = target_all["sources"][contract[0]]["AST"]
-            crytic_compile.filenames.add(filename)
             compilation_unit.filenames.add(filename)
             compilation_unit.filename_to_contracts[filename].add(contract_name)
             source_unit.contracts_names.add(contract_name)