From de6c86115be2410ac567c01ec633c58aef201658 Mon Sep 17 00:00:00 2001
From: Jonathan Strong <jstrong@shipyard.rs>
Date: Thu, 9 Nov 2023 12:44:45 -0500
Subject: [PATCH] fix editing of deps in manifest

---
 src/publish.rs | 74 ++++++++++++++++++++++++++------------------------
 1 file changed, 39 insertions(+), 35 deletions(-)

diff --git a/src/publish.rs b/src/publish.rs
index ab869dc..f7fe0a4 100644
--- a/src/publish.rs
+++ b/src/publish.rs
@@ -846,54 +846,61 @@ fn parse_manifests(
 //     }
 // }
 
-// conditional dep tables aren't handled right:
-//
-// [target.'cfg(not(target_env = "msvc"))'.dependencies]
-// dep-one = { version = "0.1.0", registry = "old-registry" }
-//
-fn edit_dep_registries(
-    dep_key: &str,
+/// edit registry deps to point to the destination registry.
+///
+/// NOTE: recursive traversing of the toml is needed to handle things
+/// like conditional deps blocks like:
+///
+/// ```toml,ignore
+/// [target.'cfg(not(target_env = "msvc"))'.dependencies]
+/// dep-one = { version = "0.1.0", registry = "old-registry" }
+/// ```
+fn edit_deps(
     manifest: &mut toml_edit::Document,
     config: &Config,
-) -> Result<(), Error> {
-    let src_registry_name = config.src.registry_name.as_str();
-    let dst_registry_name = config.dst.registry_name.as_str();
-    let src_index_url = config.src.index_url.as_str();
-    let dst_index_url = config.dst.index_url.as_str();
+) {
+    use toml_edit::{visit_mut::VisitMut, KeyMut, Item, TableLike};
 
-    let dep_key_ends_with = format!(".{dep_key}");
+    struct DepsVisitor<'a>(&'a Config);
 
-    let it = manifest.iter_mut()
-        .filter(|(k, v)| {
-            k == dep_key || k.ends_with(&dep_key_ends_with)
-        });
+    impl<'a> VisitMut for DepsVisitor<'a> {
+        fn visit_table_like_mut(&mut self, dep: &mut dyn TableLike) {
+            let config = self.0;
+            let src_registry_name = config.src.registry_name.as_str();
+            let dst_registry_name = config.dst.registry_name.as_str();
+            let src_index_url = config.src.index_url.as_str();
+            let dst_index_url = config.dst.index_url.as_str();
+            let mut edited = false;
 
-    for (outer_k, outer_v) in it {
-        let Some(deps) = outer_v.as_table_like_mut() else {
-            anyhow::bail!("failed to cast deps item as table");
-        };
-        for (k, v) in deps.iter_mut() {
-            let Some(t) = v.as_table_like_mut() else { continue };
-
-            if let Some(registry_item) = t.get_mut("registry") {
+            if let Some(registry_item) = dep.get_mut("registry") {
                 if registry_item.as_str().unwrap_or("") == src_registry_name {
-                    trace!(dep_name = ?k, %dep_key, ?src_registry_name, ?dst_registry_name, "modifying registry in Cargo.toml");
                     *registry_item = toml_edit::value(dst_registry_name);
+                    edited = true;
                 }
             }
 
-            if let Some(registry_index_item) = t.get_mut("registry-index") {
+            if let Some(registry_index_item) = dep.get_mut("registry-index") {
                 if registry_index_item.as_str().unwrap_or("") == src_index_url {
-                    trace!(dep_name = ?k, %dep_key, ?src_index_url, ?dst_index_url, "modifying registry-index in Cargo.toml");
                     *registry_index_item = toml_edit::value(dst_index_url);
+                    edited = true;
+                }
+            }
+
+            if !edited {
+                for (_, v) in dep.iter_mut() {
+                    if let Some(t) = v.as_table_like_mut() {
+                        toml_edit::visit_mut::visit_table_like_mut(self, t);
+                    }
                 }
             }
         }
     }
-    Ok(())
+
+    let mut visitor = DepsVisitor(config);
+    visitor.visit_document_mut(&mut *manifest);
 }
 
-fn edit_publish_registry_if_present(
+fn edit_publish_registry(
     manifest: &mut toml_edit::Document,
     src_registry_name: &str,
     dst_registry_name: &str,
@@ -925,11 +932,8 @@ fn prepare_source_dir_for_publish(config: &Config, meta: &mut VersionMeta) -> Re
     let source_dir = meta.source_dir();
     let mut modified_manifest = meta.manifest_files.cargo_toml_orig.parse::<toml_edit::Document>()?;
 
-    edit_dep_registries("dependencies", &mut modified_manifest, &config)?;
-    edit_dep_registries("dev-dependencies", &mut modified_manifest, &config)?;
-    edit_dep_registries("build-dependencies", &mut modified_manifest, &config)?;
-
-    edit_publish_registry_if_present(&mut modified_manifest, &config.src.registry_name, &config.dst.registry_name)?;
+    edit_deps(&mut modified_manifest, &config);
+    edit_publish_registry(&mut modified_manifest, &config.src.registry_name, &config.dst.registry_name)?;
 
     // write modified manifest over Cargo.toml (leaves Cargo.toml.orig as is)
     let modified_manifest_toml = modified_manifest.to_string();