fix(forge): fix forge clone when there is nested src

crates/forge/bin/cmd/clone.rs

@@ -71,6 +71,14 @@ pub struct CloneArgs {
     #[arg(long)]
     pub no_remappings_txt: bool,
 
+    /// Keep the original directory structure collected from Etherscan.
+    ///
+    /// If this flag is set, the directory structure of the cloned project will be kept as is.
+    /// By default, the directory structure is re-orgnized to increase the readability, but may
+    /// risk some compilation failures.
+    #[arg(long)]
+    pub keep_directory_structure: bool,
+
     #[command(flatten)]
     pub etherscan: EtherscanOpts,
 
@@ -80,7 +88,14 @@ pub struct CloneArgs {
 
 impl CloneArgs {
     pub async fn run(self) -> Result<()> {
-        let CloneArgs { address, root, opts, etherscan, no_remappings_txt } = self;
+        let CloneArgs {
+            address,
+            root,
+            opts,
+            etherscan,
+            no_remappings_txt,
+            keep_directory_structure,
+        } = self;
 
         // step 0. get the chain and api key from the config
         let config = Config::from(&etherscan);
@@ -99,7 +114,8 @@ impl CloneArgs {
         let root = dunce::canonicalize(&root)?;
 
         // step 3. parse the metadata
-        Self::parse_metadata(&meta, chain, &root, no_remappings_txt).await?;
+        Self::parse_metadata(&meta, chain, &root, no_remappings_txt, keep_directory_structure)
+            .await?;
 
         // step 4. collect the compilation metadata
         // if the etherscan api key is not set, we need to wait for 3 seconds between calls
@@ -214,33 +230,24 @@ impl CloneArgs {
         chain: Chain,
         root: &PathBuf,
         no_remappings_txt: bool,
+        keep_directory_structure: bool,
     ) -> Result<()> {
         // dump sources and update the remapping in configuration
-        let Settings { remappings: original_remappings, .. } = meta.settings()?;
-        let (remappings, strip_old_src) = dump_sources(meta, root)?;
+        let remappings = dump_sources(meta, root, keep_directory_structure)?;
         Config::update_at(root, |config, doc| {
             let profile = config.profile.as_str().as_str();
+
+            // update the remappings in the configuration
             let mut remapping_array = toml_edit::Array::new();
-            // original remappings
-            for r in original_remappings.iter() {
-                // we should update its remapped path in the same way as we dump sources
-                // i.e., remove prefix `contracts` (if any) and add prefix `src`
-                let mut r = r.to_owned();
-                if strip_old_src {
-                    let new_path = PathBuf::from(r.path.clone())
-                        .strip_prefix("contracts")
-                        .map(|p| p.to_path_buf())
-                        .unwrap_or(PathBuf::from(r.path));
-                    r.path = PathBuf::from("src").join(new_path).to_string_lossy().to_string();
-                }
-                remapping_array.push(r.to_string());
-            }
-            // new remappings
             for r in remappings {
                 remapping_array.push(r.to_string());
             }
             doc[Config::PROFILE_SECTION][profile]["remappings"] = toml_edit::value(remapping_array);
 
+            // make sure auto_detect_remappings is false (it is very important because cloned
+            // project may not follow the common remappings)
+            doc[Config::PROFILE_SECTION][profile]["auto_detect_remappings"] =
+                toml_edit::value(false);
             true
         })?;
 
@@ -407,74 +414,136 @@ fn update_config_by_metadata(
 /// Dump the contract sources to the root directory.
 /// The sources are dumped to the `src` directory.
 /// IO errors may be returned.
-/// A list of remappings is returned, as well as a boolean indicating whether the old `contract`
-/// or `src` directories are stripped.
-fn dump_sources(meta: &Metadata, root: &PathBuf) -> Result<(Vec<RelativeRemapping>, bool)> {
+/// A list of remappings is returned
+fn dump_sources(meta: &Metadata, root: &PathBuf, no_reorg: bool) -> Result<Vec<RelativeRemapping>> {
     // get config
     let path_config = ProjectPathsConfig::builder().build_with_root(root);
     // we will canonicalize the sources directory later
     let src_dir = &path_config.sources;
+    let lib_dir = &path_config.libraries[0];
     let contract_name = &meta.contract_name;
     let source_tree = meta.source_tree();
 
     // then we move the sources to the correct directories
     // we will first load existing remappings if necessary
     //  make sure this happens before dumping sources
     let mut remappings: Vec<Remapping> = Remapping::find_many(root);
-    // we also load the original remappings from the metadata
-    remappings.extend(meta.settings()?.remappings);
 
     // first we dump the sources to a temporary directory
     let tmp_dump_dir = root.join("raw_sources");
     source_tree
         .write_to(&tmp_dump_dir)
         .map_err(|e| eyre::eyre!("failed to dump sources: {}", e))?;
 
-    // check whether we need to strip the `contract` or `src` directories in the original sources.
-    // They are the source directories in foundry or hardhat projects. We do not want to preserve
-    // them in the cloned project.
-    // If there is any other directory other than `src` `contracts` or `lib`, we should not strip.
-    let strip_old_src = std::fs::read_dir(tmp_dump_dir.join(contract_name))?.all(|e| {
-        let Ok(e) = e else { return false };
-        let folder_name = e.file_name().to_string_lossy().to_string();
-        ["contracts", "src", "lib"].contains(&folder_name.as_str())
-    });
-    // move contract sources to the `src` directory
+    // check whether we need to re-organize directories in the original sources, since we do not
+    // want to put all the sources in the `src` directory if the original directory structure is
+    // well organized, e.g., a standard foundry project containing `src` and `lib`
+    //
+    // * if the user wants to keep the original directory structure, we should not re-organize.
+    // * if there is any other directory other than `src`, `contracts`, `lib`, `hardhat`,
+    //   `forge-std`,
+    // or not started with `@`, we should not re-organize.
+    let to_reorg = !no_reorg &&
+        std::fs::read_dir(tmp_dump_dir.join(contract_name))?.all(|e| {
+            let Ok(e) = e else { return false };
+            let folder_name = e.file_name();
+            folder_name == "src" ||
+                folder_name == "lib" ||
+                folder_name == "contracts" ||
+                folder_name == "hardhat" ||
+                folder_name == "forge-std" ||
+                folder_name.to_string_lossy().starts_with('@')
+        });
+
+    // ensure `src` and `lib` directories exist
+    eyre::ensure!(Path::exists(&root.join(src_dir)), "`src` directory must exists");
+    eyre::ensure!(Path::exists(&root.join(lib_dir)), "`lib` directory must exists");
+
+    // move source files
     for entry in std::fs::read_dir(tmp_dump_dir.join(contract_name))? {
-        if std::fs::metadata(root.join(src_dir)).is_err() {
-            std::fs::create_dir(root.join(src_dir))?;
-        }
         let entry = entry?;
-        let folder_name = entry.file_name().to_string_lossy().to_string();
-        // special handling for contracts and src directories: we flatten them.
-        if strip_old_src && (folder_name.as_str() == "contracts" || folder_name.as_str() == "src") {
-            // move all sub folders in contracts to src
-            for e in read_dir(entry.path())? {
-                let e = e?;
-                let dest = src_dir.join(e.file_name());
-                std::fs::rename(e.path(), &dest)?;
+        let folder_name = entry.file_name();
+        // special handling when we need to re-organize the directories: we flatten them.
+        if to_reorg {
+            if folder_name == "contracts" || folder_name == "src" || folder_name == "lib" {
+                // move all sub folders in contracts to src or lib
+                let new_dir = if folder_name == "lib" { lib_dir } else { src_dir };
+                for e in read_dir(entry.path())? {
+                    let e = e?;
+                    let dest = new_dir.join(&e.file_name());
+                    eyre::ensure!(!Path::exists(&dest), "destination already exists: {:?}", dest);
+                    std::fs::rename(e.path(), &dest)?;
+                    remappings.push(Remapping {
+                        context: None,
+                        name: format!(
+                            "{}/{}",
+                            folder_name.to_string_lossy(),
+                            e.file_name().to_string_lossy()
+                        ),
+                        path: dest.to_string_lossy().to_string(),
+                    });
+                }
+            } else {
+                assert!(
+                    folder_name == "hardhat" ||
+                        folder_name == "forge-std" ||
+                        folder_name.to_string_lossy().starts_with('@')
+                );
+                // move these other folders to lib
+                let dest = lib_dir.join(&folder_name);
+                if folder_name == "forge-std" {
+                    // let's use the provided forge-std directory
+                    std::fs::remove_dir_all(&dest)?;
+                }
+                eyre::ensure!(!Path::exists(&dest), "destination already exists: {:?}", dest);
+                std::fs::rename(entry.path(), &dest)?;
                 remappings.push(Remapping {
                     context: None,
-                    name: format!("{}/{}", folder_name, e.file_name().to_string_lossy()),
+                    name: folder_name.to_string_lossy().to_string(),
                     path: dest.to_string_lossy().to_string(),
                 });
             }
         } else {
-            // move the other folders to src
-            let dest = src_dir.join(entry.file_name());
+            // directly move the all folders into src
+            let dest = src_dir.join(&folder_name);
+            eyre::ensure!(!Path::exists(&dest), "destination already exists: {:?}", dest);
             std::fs::rename(entry.path(), &dest)?;
-            remappings.push(Remapping {
-                context: None,
-                name: entry.file_name().to_string_lossy().to_string(),
-                path: dest.to_string_lossy().to_string(),
-            });
+            if folder_name != "src" {
+                remappings.push(Remapping {
+                    context: None,
+                    name: folder_name.to_string_lossy().to_string(),
+                    path: dest.to_string_lossy().to_string(),
+                });
+            }
         }
     }
 
     // remove the temporary directory
     std::fs::remove_dir_all(tmp_dump_dir)?;
 
-    Ok((remappings.into_iter().map(|r| r.into_relative(root)).collect(), strip_old_src))
+    // add remappings in the metedata
+    for mut r in meta.settings()?.remappings {
+        if to_reorg {
+            // we should update its remapped path in the same way as we dump sources
+            // i.e., remove prefix `contracts` (if any) and add prefix `src`
+            let new_path = if r.path.starts_with("contracts") {
+                PathBuf::from("src").join(PathBuf::from(&r.path).strip_prefix("contracts")?)
+            } else if r.path.starts_with('@') ||
+                r.path.starts_with("hardhat/") ||
+                r.path.starts_with("forge-std/")
+            {
+                PathBuf::from("lib").join(PathBuf::from(&r.path))
+            } else {
+                PathBuf::from(&r.path)
+            };
+            r.path = new_path.to_string_lossy().to_string();
+            remappings.push(r);
+        } else {
+            remappings.push(r);
+        }
+    }
+
+    Ok(remappings.into_iter().map(|r| r.into_relative(root)).collect())
 }
 
 /// Compile the project in the root directory, and return the compilation result.
@@ -617,7 +686,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     #[ignore = "this test is used to dump mock data from Etherscan"]
     async fn test_dump_mock_data() {
-        let address: Address = "0x9ab6b21cdf116f611110b048987e58894786c244".parse().unwrap();
+        let address: Address = "0x9d27527Ada2CF29fBDAB2973cfa243845a08Bd3F".parse().unwrap();
         let data_folder = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
             .join("../../testdata/etherscan")
             .join(address.to_string());
@@ -644,7 +713,9 @@ mod tests {
         let meta = CloneArgs::collect_metadata_from_client(address, &client).await.unwrap();
         CloneArgs::init_an_empty_project(&project_root, DependencyInstallOpts::default()).unwrap();
         project_root = dunce::canonicalize(&project_root).unwrap();
-        CloneArgs::parse_metadata(&meta, Chain::mainnet(), &project_root, false).await.unwrap();
+        CloneArgs::parse_metadata(&meta, Chain::mainnet(), &project_root, false, false)
+            .await
+            .unwrap();
         CloneArgs::collect_compilation_metadata(
             &meta,
             Chain::mainnet(),
@@ -706,6 +777,12 @@ mod tests {
         one_test_case(address, false).await
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_clone_contract_with_nested_src() {
+        let address = "0x9d27527Ada2CF29fBDAB2973cfa243845a08Bd3F".parse().unwrap();
+        one_test_case(address, false).await
+    }
+
     fn pick_creation_info(address: &str) -> Option<(&'static str, &'static str)> {
         for (addr, contract_name, creation_code) in CREATION_ARRAY.iter() {
             if address == *addr {

crates/forge/tests/cli/cmd.rs

@@ -468,6 +468,48 @@ forgetest!(can_clone_quiet, |prj, cmd| {
     cmd.assert_empty_stdout();
 });
 
+// checks that clone works with --no-remappings-txt
+forgetest!(can_clone_no_remappings_txt, |prj, cmd| {
+    prj.wipe();
+
+    let foundry_toml = prj.root().join(Config::FILE_NAME);
+    assert!(!foundry_toml.exists());
+
+    cmd.args([
+        "clone",
+        "--etherscan-api-key",
+        next_etherscan_api_key().as_str(),
+        "--no-remappings-txt",
+        "0x33e690aEa97E4Ef25F0d140F1bf044d663091DAf",
+    ])
+    .arg(prj.root());
+    cmd.assert_non_empty_stdout();
+
+    let s = read_string(&foundry_toml);
+    let _config: BasicConfig = parse_with_profile(&s).unwrap().unwrap().1;
+});
+
+// checks that clone works with --keep-directory-structure
+forgetest!(can_clone_keep_directory_structure, |prj, cmd| {
+    prj.wipe();
+
+    let foundry_toml = prj.root().join(Config::FILE_NAME);
+    assert!(!foundry_toml.exists());
+
+    cmd.args([
+        "clone",
+        "--etherscan-api-key",
+        next_etherscan_api_key().as_str(),
+        "--keep-directory-structure",
+        "0x33e690aEa97E4Ef25F0d140F1bf044d663091DAf",
+    ])
+    .arg(prj.root());
+    cmd.assert_non_empty_stdout();
+
+    let s = read_string(&foundry_toml);
+    let _config: BasicConfig = parse_with_profile(&s).unwrap().unwrap().1;
+});
+
 // checks that `clean` removes dapptools style paths
 forgetest!(can_clean, |prj, cmd| {
     prj.assert_create_dirs_exists();

testdata/etherscan/0x9d27527Ada2CF29fBDAB2973cfa243845a08Bd3F/creation_data.json

@@ -0,0 +1,5 @@
+{
+  "contractAddress": "0x9d27527ada2cf29fbdab2973cfa243845a08bd3f",
+  "contractCreator": "0x7773ae67403d2e30102a84c48cc939919c4c881c",
+  "txHash": "0xc757719b7ae11ea651b1b23249978a3e8fca94d358610f5809a5dc19fbf850ce"
+}