From 13e67aaa4a9806c46874139f3dae1bdc3c6c55e6 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Thu, 30 Jun 2022 18:19:55 +0200
Subject: [PATCH] feat(solc): add remove_contract utility function (#1436)

---
 ethers-solc/src/compile/output/info.rs | 35 +++++++++++++++++++++++++-
 ethers-solc/src/compile/output/mod.rs  | 33 ++++++++++++++++++++++++
 ethers-solc/tests/project.rs           | 12 +++++++--
 3 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/ethers-solc/src/compile/output/info.rs b/ethers-solc/src/compile/output/info.rs
index 8e118cea..3094bd9a 100644
--- a/ethers-solc/src/compile/output/info.rs
+++ b/ethers-solc/src/compile/output/info.rs
@@ -1,5 +1,5 @@
 //! Commonly used identifiers for contracts in the compiled output
-use std::{convert::TryFrom, fmt, str::FromStr};
+use std::{borrow::Cow, convert::TryFrom, fmt, str::FromStr};
 
 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
 #[error("{0}")]
@@ -77,6 +77,39 @@ impl From<FullContractInfo> for ContractInfo {
     }
 }
 
+/// The reference type for `ContractInfo`
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub struct ContractInfoRef<'a> {
+    pub path: Option<Cow<'a, str>>,
+    pub name: Cow<'a, str>,
+}
+
+impl<'a> From<ContractInfo> for ContractInfoRef<'a> {
+    fn from(info: ContractInfo) -> Self {
+        ContractInfoRef { path: info.path.map(Into::into), name: info.name.into() }
+    }
+}
+
+impl<'a> From<&'a ContractInfo> for ContractInfoRef<'a> {
+    fn from(info: &'a ContractInfo) -> Self {
+        ContractInfoRef {
+            path: info.path.as_deref().map(Into::into),
+            name: info.name.as_str().into(),
+        }
+    }
+}
+impl<'a> From<FullContractInfo> for ContractInfoRef<'a> {
+    fn from(info: FullContractInfo) -> Self {
+        ContractInfoRef { path: Some(info.path.into()), name: info.name.into() }
+    }
+}
+
+impl<'a> From<&'a FullContractInfo> for ContractInfoRef<'a> {
+    fn from(info: &'a FullContractInfo) -> Self {
+        ContractInfoRef { path: Some(info.path.as_str().into()), name: info.name.as_str().into() }
+    }
+}
+
 /// Represents the common contract argument pattern `<path>:<contractname>`
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct FullContractInfo {
diff --git a/ethers-solc/src/compile/output/mod.rs b/ethers-solc/src/compile/output/mod.rs
index 4724cf7f..bcb5e1e8 100644
--- a/ethers-solc/src/compile/output/mod.rs
+++ b/ethers-solc/src/compile/output/mod.rs
@@ -6,6 +6,7 @@ use crate::{
         Error,
     },
     buildinfo::RawBuildInfo,
+    info::ContractInfoRef,
     sources::{VersionedSourceFile, VersionedSourceFiles},
     ArtifactId, ArtifactOutput, Artifacts, CompilerOutput, ConfigurableArtifacts, SolcIoError,
 };
@@ -378,6 +379,38 @@ impl AggregatedCompilerOutput {
         self.contracts.remove(path, contract)
     }
 
+    /// Removes the contract with matching path and name using the `<path>:<contractname>` pattern
+    /// where `path` is optional.
+    ///
+    /// If the `path` segment is `None`, then the first matching `Contract` is returned, see
+    /// [Self::remove_first]
+    ///
+    ///
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use ethers_solc::Project;
+    /// use ethers_solc::artifacts::*;
+    /// use ethers_solc::info::ContractInfo;
+    /// # fn demo(project: Project) {
+    /// let mut output = project.compile().unwrap().output();
+    /// let info = ContractInfo::new("src/Greeter.sol:Greeter");
+    /// let contract = output.remove_contract(&info).unwrap();
+    /// # }
+    /// ```
+    pub fn remove_contract<'a>(
+        &mut self,
+        info: impl Into<ContractInfoRef<'a>>,
+    ) -> Option<Contract> {
+        let ContractInfoRef { path, name } = info.into();
+        if let Some(path) = path {
+            self.remove(path, name)
+        } else {
+            self.remove_first(name)
+        }
+    }
+
     /// Iterate over all contracts and their names
     pub fn contracts_iter(&self) -> impl Iterator<Item = (&String, &Contract)> {
         self.contracts.contracts()
diff --git a/ethers-solc/tests/project.rs b/ethers-solc/tests/project.rs
index 322e7a20..f126f5c3 100644
--- a/ethers-solc/tests/project.rs
+++ b/ethers-solc/tests/project.rs
@@ -15,6 +15,7 @@ use ethers_solc::{
     },
     buildinfo::BuildInfo,
     cache::{SolFilesCache, SOLIDITY_FILES_CACHE_FILENAME},
+    info::ContractInfo,
     project_util::*,
     remappings::Remapping,
     CompilerInput, ConfigurableArtifacts, ExtraOutputValues, Graph, Project, ProjectCompileOutput,
@@ -1548,9 +1549,16 @@ fn can_compile_sparse_with_link_references() {
     let lib = dup.remove_first("MyLib");
     assert!(lib.is_none());
 
-    let lib = output.remove(my_lib_path.to_string_lossy(), "MyLib");
+    dup = output.clone();
+    let lib = dup.remove(my_lib_path.to_string_lossy(), "MyLib");
     assert!(lib.is_some());
-    let lib = output.remove(my_lib_path.to_string_lossy(), "MyLib");
+    let lib = dup.remove(my_lib_path.to_string_lossy(), "MyLib");
+    assert!(lib.is_none());
+
+    let info = ContractInfo::new(format!("{}:{}", my_lib_path.to_string_lossy(), "MyLib"));
+    let lib = output.remove_contract(&info);
+    assert!(lib.is_some());
+    let lib = output.remove_contract(&info);
     assert!(lib.is_none());
 }