Cleanup, remove Hir:: usage, use ContractIds

Author: grandizzy

crates/compilers/src/lib.rs

@@ -903,6 +903,7 @@ fn apply_updates(sources: &mut Sources, updates: Updates) {
 }
 
 /// Utility function to change source content ranges with provided updates.
+/// Assumes that the updates are sorted.
 fn replace_source_content<'a>(
     source: &str,
     updates: impl IntoIterator<Item = (Range<usize>, &'a str)>,

crates/compilers/src/preprocessor/data.rs

@@ -12,39 +12,39 @@ use std::{
 };
 
 /// Keeps data about project contracts definitions referenced from tests and scripts.
-/// HIR id -> Contract data definition mapping.
-pub type PreprocessorData = BTreeMap<u32, ContractData>;
+/// Contract id -> Contract data definition mapping.
+pub type PreprocessorData = BTreeMap<ContractId, ContractData>;
 
 /// Collects preprocessor data from referenced contracts.
-pub fn collect_preprocessor_data(
+pub(crate) fn collect_preprocessor_data(
     sess: &Session,
     hir: &Hir<'_>,
-    referenced_contracts: HashSet<u32>,
+    referenced_contracts: &HashSet<ContractId>,
 ) -> PreprocessorData {
     let mut data = PreprocessorData::default();
     for contract_id in referenced_contracts {
-        let contract = Hir::contract(hir, ContractId::new(contract_id));
-        let source = Hir::source(hir, contract.source);
+        let contract = hir.contract(*contract_id);
+        let source = hir.source(contract.source);
 
         let FileName::Real(path) = &source.file.name else {
             continue;
         };
 
         let contract_data =
-            ContractData::new(hir, contract_id, contract, path, source, sess.source_map());
-        data.insert(contract_id, contract_data);
+            ContractData::new(hir, *contract_id, contract, path, source, sess.source_map());
+        data.insert(*contract_id, contract_data);
     }
     data
 }
 
 /// Creates helper libraries for contracts with a non-empty constructor.
 ///
 /// See [`ContractData::build_helper`] for more details.
-pub fn create_deploy_helpers(data: &BTreeMap<u32, ContractData>) -> Sources {
+pub(crate) fn create_deploy_helpers(data: &BTreeMap<ContractId, ContractData>) -> Sources {
     let mut deploy_helpers = Sources::new();
     for (contract_id, contract) in data {
         if let Some(code) = contract.build_helper() {
-            let path = format!("foundry-pp/DeployHelper{contract_id}.sol");
+            let path = format!("foundry-pp/DeployHelper{}.sol", contract_id.get());
             deploy_helpers.insert(path.into(), Source::new(code));
         }
     }
@@ -62,9 +62,9 @@ pub struct ContractConstructorData {
 
 /// Keeps data about a single contract definition.
 #[derive(Debug)]
-pub struct ContractData {
+pub(crate) struct ContractData {
     /// HIR Id of the contract.
-    contract_id: u32,
+    contract_id: ContractId,
     /// Path of the source file.
     path: PathBuf,
     /// Name of the contract
@@ -78,7 +78,7 @@ pub struct ContractData {
 impl ContractData {
     fn new(
         hir: &Hir<'_>,
-        contract_id: u32,
+        contract_id: ContractId,
         contract: &Contract<'_>,
         path: &Path,
         source: &solar_sema::hir::Source<'_>,
@@ -89,25 +89,21 @@ impl ContractData {
         // Process data for contracts with constructor and parameters.
         let constructor_data = contract
             .ctor
-            .map(|ctor_id| Hir::function(hir, ctor_id))
+            .map(|ctor_id| hir.function(ctor_id))
             .filter(|ctor| !ctor.parameters.is_empty())
             .map(|ctor| {
                 let abi_encode_args = ctor
                     .parameters
                     .iter()
-                    .map(|param_id| {
-                        format!("args.{}", Hir::variable(hir, *param_id).name.unwrap().name)
-                    })
+                    .map(|param_id| format!("args.{}", hir.variable(*param_id).name.unwrap().name))
                     .join(", ");
                 let struct_fields = ctor
                     .parameters
                     .iter()
                     .map(|param_id| {
                         let src = source.file.src.as_str();
-                        let loc = SourceMapLocation::from_span(
-                            source_map,
-                            Hir::variable(hir, *param_id).span,
-                        );
+                        let loc =
+                            SourceMapLocation::from_span(source_map, hir.variable(*param_id).span);
                         src[loc.start..loc.end].replace(" memory ", " ").replace(" calldata ", " ")
                     })
                     .join("; ");
@@ -170,6 +166,7 @@ impl ContractData {
         let Self { contract_id, path, name, constructor_data, artifact } = self;
 
         let Some(constructor_details) = constructor_data else { return None };
+        let contract_id = contract_id.get();
         let struct_fields = &constructor_details.struct_fields;
         let abi_encode_args = &constructor_details.abi_encode_args;
         let vm_interface_name = format!("VmContractHelper{contract_id}");

crates/compilers/src/preprocessor/deps.rs

@@ -19,20 +19,20 @@ use std::{
 };
 
 /// Holds data about referenced source contracts and bytecode dependencies.
-pub struct PreprocessorDependencies {
+pub(crate) struct PreprocessorDependencies {
     // Mapping contract id to preprocess -> contract bytecode dependencies.
-    pub preprocessed_contracts: BTreeMap<u32, Vec<BytecodeDependency>>,
+    pub preprocessed_contracts: BTreeMap<ContractId, Vec<BytecodeDependency>>,
     // Referenced contract ids.
-    pub referenced_contracts: HashSet<u32>,
+    pub referenced_contracts: HashSet<ContractId>,
 }
 
 impl PreprocessorDependencies {
     pub fn new(sess: &Session, hir: &Hir<'_>, paths: &[PathBuf], src_dir: &PathBuf) -> Self {
         let mut preprocessed_contracts = BTreeMap::new();
         let mut referenced_contracts = HashSet::new();
-        for contract_id in Hir::contract_ids(hir) {
-            let contract = Hir::contract(hir, contract_id);
-            let source = Hir::source(hir, contract.source);
+        for contract_id in hir.contract_ids() {
+            let contract = hir.contract(contract_id);
+            let source = hir.source(contract.source);
 
             let FileName::Real(path) = &source.file.name else {
                 continue;
@@ -53,7 +53,7 @@ impl PreprocessorDependencies {
             deps_collector.walk_contract(contract);
             // Ignore empty test contracts declared in source files with other contracts.
             if !deps_collector.dependencies.is_empty() {
-                preprocessed_contracts.insert(contract_id.get(), deps_collector.dependencies);
+                preprocessed_contracts.insert(contract_id, deps_collector.dependencies);
             }
             // Record collected referenced contract ids.
             referenced_contracts.extend(deps_collector.referenced_contracts);
@@ -73,13 +73,13 @@ enum BytecodeDependencyKind {
 
 /// Represents a single bytecode dependency.
 #[derive(Debug)]
-pub struct BytecodeDependency {
+pub(crate) struct BytecodeDependency {
     /// Dependency kind.
     kind: BytecodeDependencyKind,
     /// Source map location of this dependency.
     loc: SourceMapLocation,
     /// HIR id of referenced contract.
-    referenced_contract: u32,
+    referenced_contract: ContractId,
 }
 
 /// Walks over contract HIR and collects [`BytecodeDependency`]s and referenced contracts.
@@ -95,7 +95,7 @@ struct BytecodeDependencyCollector<'hir> {
     /// Dependencies collected for current contract.
     dependencies: Vec<BytecodeDependency>,
     /// Unique HIR ids of contracts referenced from current contract.
-    referenced_contracts: HashSet<u32>,
+    referenced_contracts: HashSet<ContractId>,
 }
 
 impl<'hir> BytecodeDependencyCollector<'hir> {
@@ -119,8 +119,8 @@ impl<'hir> BytecodeDependencyCollector<'hir> {
     /// Discards any reference that is not in project src directory (e.g. external
     /// libraries or mock contracts that extend source contracts).
     fn collect_dependency(&mut self, dependency: BytecodeDependency) {
-        let contract = Hir::contract(self.hir, ContractId::new(dependency.referenced_contract));
-        let source = Hir::source(self.hir, contract.source);
+        let contract = self.hir.contract(dependency.referenced_contract);
+        let source = self.hir.source(contract.source);
         let FileName::Real(path) = &source.file.name else {
             return;
         };
@@ -159,20 +159,20 @@ impl<'hir> Visit<'hir> for BytecodeDependencyCollector<'hir> {
                                 self.source_map,
                                 Span::new(expr.span.lo(), expr.span.hi()),
                             ),
-                            referenced_contract: contract_id.get(),
+                            referenced_contract: contract_id,
                         });
                     }
                 }
             }
             ExprKind::Member(member_expr, ident) => {
-                if ident.name.to_string() == "creationCode" {
-                    if let ExprKind::TypeCall(ty) = &member_expr.kind {
-                        if let TypeKind::Custom(contract_id) = &ty.kind {
+                if let ExprKind::TypeCall(ty) = &member_expr.kind {
+                    if let TypeKind::Custom(contract_id) = &ty.kind {
+                        if ident.name.as_str() == "creationCode" {
                             if let Some(contract_id) = contract_id.as_contract() {
                                 self.collect_dependency(BytecodeDependency {
                                     kind: BytecodeDependencyKind::CreationCode,
                                     loc: SourceMapLocation::from_span(self.source_map, expr.span),
-                                    referenced_contract: contract_id.get(),
+                                    referenced_contract: contract_id,
                                 });
                             }
                         }
@@ -187,23 +187,23 @@ impl<'hir> Visit<'hir> for BytecodeDependencyCollector<'hir> {
 
 /// Goes over all test/script files and replaces bytecode dependencies with cheatcode
 /// invocations.
-pub fn remove_bytecode_dependencies(
+pub(crate) fn remove_bytecode_dependencies(
     hir: &Hir<'_>,
     deps: &PreprocessorDependencies,
     data: &PreprocessorData,
 ) -> Updates {
     let mut updates = Updates::default();
     for (contract_id, deps) in &deps.preprocessed_contracts {
-        let contract = Hir::contract(hir, ContractId::new(*contract_id));
-        let source = Hir::source(hir, contract.source);
+        let contract = hir.contract(*contract_id);
+        let source = hir.source(contract.source);
         let FileName::Real(path) = &source.file.name else {
             continue;
         };
 
         let updates = updates.entry(path.clone()).or_default();
         let mut used_helpers = BTreeSet::new();
 
-        let vm_interface_name = format!("VmContractHelper{contract_id}");
+        let vm_interface_name = format!("VmContractHelper{}", contract_id.get());
         let vm = format!("{vm_interface_name}(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D)");
 
         for dep in deps {
@@ -239,7 +239,7 @@ pub fn remove_bytecode_dependencies(
                             dep.loc.end,
                             format!(
                                 "deployCode{id}(DeployHelper{id}.ConstructorArgs",
-                                id = dep.referenced_contract
+                                id = dep.referenced_contract.get()
                             ),
                         ));
                         updates.insert((
@@ -252,6 +252,7 @@ pub fn remove_bytecode_dependencies(
             };
         }
         let helper_imports = used_helpers.into_iter().map(|id| {
+            let id = id.get();
             format!(
                 "import {{DeployHelper{id}, encodeArgs{id}, deployCode{id}}} from \"foundry-pp/DeployHelper{id}.sol\";",
             )

crates/compilers/src/preprocessor/mod.rs

@@ -98,7 +98,7 @@ impl Preprocessor<SolcCompiler> for TestOptimizerPreprocessor {
                 &paths.paths_relative().sources,
             );
             // Collect data of source contracts referenced in tests and scripts.
-            let data = collect_preprocessor_data(&sess, hir, deps.referenced_contracts.clone());
+            let data = collect_preprocessor_data(&sess, hir, &deps.referenced_contracts);
 
             // Extend existing sources with preprocessor deploy helper sources.
             sources.extend(create_deploy_helpers(&data));