diff --git a/crytic_compile/compilation_unit.py b/crytic_compile/compilation_unit.py index 3198b739..4e3550ee 100644 --- a/crytic_compile/compilation_unit.py +++ b/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 ################################################################################### ################################################################################### diff --git a/crytic_compile/crytic_compile.py b/crytic_compile/crytic_compile.py index 0a21342b..0d09c4cf 100644 --- a/crytic_compile/crytic_compile.py +++ b/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]: diff --git a/crytic_compile/platform/standard.py b/crytic_compile/platform/standard.py index bc3f236c..62d3856b 100644 --- a/crytic_compile/platform/standard.py +++ b/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", []) diff --git a/crytic_compile/platform/waffle.py b/crytic_compile/platform/waffle.py index 547b1ce7..d0389e88 100755 --- a/crytic_compile/platform/waffle.py +++ b/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)