feat!(solc): add additional remove functions (#1406)

gakonst · Jun 22, 2022 · 2ebdef6 · 2ebdef6
1 parent 42e9666
commit 2ebdef6
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 16 deletions.
diff --git a/ethers-solc/src/artifacts/contract.rs b/ethers-solc/src/artifacts/contract.rs
@@ -89,7 +89,7 @@ impl ContractBytecode {
     /// use ethers_solc::artifacts::*;
     /// # fn demo(project: Project) {
     /// let mut output = project.compile().unwrap().output();
-    /// let contract: ContractBytecode = output.remove("Greeter").unwrap().into();
+    /// let contract: ContractBytecode = output.remove_first("Greeter").unwrap().into();
     /// let contract = contract.unwrap();
     /// # }
     /// ```
@@ -312,7 +312,7 @@ impl CompactContract {
     /// use ethers_solc::artifacts::*;
     /// # fn demo(project: Project) {
     /// let mut output = project.compile().unwrap().output();
-    /// let contract: CompactContract = output.remove("Greeter").unwrap().into();
+    /// let contract: CompactContract = output.remove_first("Greeter").unwrap().into();
     /// let contract = contract.unwrap();
     /// # }
     /// ```

diff --git a/ethers-solc/src/artifacts/mod.rs b/ethers-solc/src/artifacts/mod.rs
@@ -858,7 +858,7 @@ pub struct Metadata {
 
 /// A helper type that ensures lossless (de)serialisation so we can preserve the exact String
 /// metadata value that's being hashed by solc
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct LosslessMetadata {
     /// The complete abi as json value
     pub raw_metadata: String,

diff --git a/ethers-solc/src/compile/output/contracts.rs b/ethers-solc/src/compile/output/contracts.rs
@@ -53,10 +53,10 @@ impl VersionedContracts {
     /// use ethers_solc::artifacts::*;
     /// # fn demo(project: Project) {
     /// let (_, mut contracts) = project.compile().unwrap().output().split();
-    /// let contract = contracts.remove("Greeter").unwrap();
+    /// let contract = contracts.remove_first("Greeter").unwrap();
     /// # }
     /// ```
-    pub fn remove(&mut self, contract: impl AsRef<str>) -> Option<Contract> {
+    pub fn remove_first(&mut self, contract: impl AsRef<str>) -> Option<Contract> {
         let contract_name = contract.as_ref();
         self.0.values_mut().find_map(|all_contracts| {
             let mut contract = None;
@@ -72,11 +72,47 @@ impl VersionedContracts {
         })
     }
 
+    ///  Removes the contract with matching path and name
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use ethers_solc::Project;
+    /// use ethers_solc::artifacts::*;
+    /// # fn demo(project: Project) {
+    /// let (_, mut contracts) = project.compile().unwrap().output().split();
+    /// let contract = contracts.remove("src/Greeter.sol", "Greeter").unwrap();
+    /// # }
+    /// ```
+    pub fn remove(&mut self, path: impl AsRef<str>, contract: impl AsRef<str>) -> Option<Contract> {
+        let contract_name = contract.as_ref();
+        let (key, mut all_contracts) = self.0.remove_entry(path.as_ref())?;
+        let mut contract = None;
+        if let Some((c, mut contracts)) = all_contracts.remove_entry(contract_name) {
+            if !contracts.is_empty() {
+                contract = Some(contracts.remove(0).contract);
+            }
+            if !contracts.is_empty() {
+                all_contracts.insert(c, contracts);
+            }
+        }
+
+        if !all_contracts.is_empty() {
+            self.0.insert(key, all_contracts);
+        }
+        contract
+    }
+
     /// Given the contract file's path and the contract's name, tries to return the contract's
     /// bytecode, runtime bytecode, and abi
-    pub fn get(&self, path: &str, contract: &str) -> Option<CompactContractRef> {
+    pub fn get(
+        &self,
+        path: impl AsRef<str>,
+        contract: impl AsRef<str>,
+    ) -> Option<CompactContractRef> {
+        let contract = contract.as_ref();
         self.0
-            .get(path)
+            .get(path.as_ref())
             .and_then(|contracts| {
                 contracts.get(contract).and_then(|c| c.get(0).map(|c| &c.contract))
             })

diff --git a/ethers-solc/src/compile/output/mod.rs b/ethers-solc/src/compile/output/mod.rs
@@ -354,11 +354,27 @@ impl AggregatedCompilerOutput {
     /// use ethers_solc::artifacts::*;
     /// # fn demo(project: Project) {
     /// let mut output = project.compile().unwrap().output();
-    /// let contract = output.remove("Greeter").unwrap();
+    /// let contract = output.remove_first("Greeter").unwrap();
     /// # }
     /// ```
-    pub fn remove(&mut self, contract: impl AsRef<str>) -> Option<Contract> {
-        self.contracts.remove(contract)
+    pub fn remove_first(&mut self, contract: impl AsRef<str>) -> Option<Contract> {
+        self.contracts.remove_first(contract)
+    }
+
+    /// Removes the contract with matching path and name
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use ethers_solc::Project;
+    /// use ethers_solc::artifacts::*;
+    /// # fn demo(project: Project) {
+    /// let mut output = project.compile().unwrap().output();
+    /// let contract = output.remove("src/Greeter.sol", "Greeter").unwrap();
+    /// # }
+    /// ```
+    pub fn remove(&mut self, path: impl AsRef<str>, contract: impl AsRef<str>) -> Option<Contract> {
+        self.contracts.remove(path, contract)
     }
 
     /// Iterate over all contracts and their names
@@ -373,7 +389,21 @@ impl AggregatedCompilerOutput {
 
     /// Given the contract file's path and the contract's name, tries to return the contract's
     /// bytecode, runtime bytecode, and abi
-    pub fn get(&self, path: &str, contract: &str) -> Option<CompactContractRef> {
+    /// # Example
+    ///
+    /// ```
+    /// use ethers_solc::Project;
+    /// use ethers_solc::artifacts::*;
+    /// # fn demo(project: Project) {
+    /// let output = project.compile().unwrap().output();
+    /// let contract = output.get("src/Greeter.sol", "Greeter").unwrap();
+    /// # }
+    /// ```
+    pub fn get(
+        &self,
+        path: impl AsRef<str>,
+        contract: impl AsRef<str>,
+    ) -> Option<CompactContractRef> {
         self.contracts.get(path, contract)
     }
 

diff --git a/ethers-solc/tests/project.rs b/ethers-solc/tests/project.rs
@@ -1516,24 +1516,40 @@ fn can_compile_sparse_with_link_references() {
     )
     .unwrap();
 
-    tmp.add_source(
-        "mylib.sol",
-        r#"
+    let my_lib_path = tmp
+        .add_source(
+            "mylib.sol",
+            r#"
     pragma solidity =0.8.12;
     library MyLib {
        function doStuff() external pure returns (uint256) {return 1337;}
     }
    "#,
-    )
-    .unwrap();
+        )
+        .unwrap();
 
     let mut compiled = tmp.compile_sparse(TestFileFilter::default()).unwrap();
     assert!(!compiled.has_compiler_errors());
 
+    let mut output = compiled.clone().output();
+
     assert!(compiled.find("ATest").is_some());
     assert!(compiled.find("MyLib").is_some());
     let lib = compiled.remove("MyLib").unwrap();
     assert!(lib.bytecode.is_some());
+    let lib = compiled.remove("MyLib");
+    assert!(lib.is_none());
+
+    let mut dup = output.clone();
+    let lib = dup.remove_first("MyLib");
+    assert!(lib.is_some());
+    let lib = dup.remove_first("MyLib");
+    assert!(lib.is_none());
+
+    let lib = output.remove(my_lib_path.to_string_lossy(), "MyLib");
+    assert!(lib.is_some());
+    let lib = output.remove(my_lib_path.to_string_lossy(), "MyLib");
+    assert!(lib.is_none());
 }
 
 #[test]