refactor: reuse some common functions

rust-lang · Jul 7, 2024 · 242919d · 242919d
1 parent 62fd2b0
commit 242919d
Showing 1 changed file with 78 additions and 48 deletions.
diff --git a/src/cargo/ops/registry/mod.rs b/src/cargo/ops/registry/mod.rs
@@ -191,25 +191,11 @@ fn get_source_id(
     gctx: &GlobalContext,
     reg_or_index: Option<&RegistryOrIndex>,
 ) -> CargoResult<RegistrySourceIds> {
-    let sid = match reg_or_index {
-        None => SourceId::crates_io(gctx)?,
-        Some(RegistryOrIndex::Index(url)) => SourceId::for_registry(url)?,
-        Some(RegistryOrIndex::Registry(r)) => SourceId::alt_registry(gctx, r)?,
-    };
-    // Load source replacements that are built-in to Cargo.
-    let builtin_replacement_sid = SourceConfigMap::empty(gctx)?
-        .load(sid, &HashSet::new())?
-        .replaced_source_id();
-    let replacement_sid = SourceConfigMap::new(gctx)?
-        .load(sid, &HashSet::new())?
-        .replaced_source_id();
+    let sid = get_initial_source_id(gctx, reg_or_index)?;
+    let (builtin_replacement_sid, replacement_sid) = get_replacement_source_ids(gctx, sid)?;
+
     if reg_or_index.is_none() && replacement_sid != builtin_replacement_sid {
-        // Neither --registry nor --index was passed and the user has configured source-replacement.
-        if let Some(replacement_name) = replacement_sid.alt_registry_key() {
-            bail!("crates-io is replaced with remote registry {replacement_name};\ninclude `--registry {replacement_name}` or `--registry crates-io`");
-        } else {
-            bail!("crates-io is replaced with non-remote-registry source {replacement_sid};\ninclude `--registry crates-io` to use crates.io");
-        }
+        bail!(gen_replacement_error(replacement_sid));
     } else {
         Ok(RegistrySourceIds {
             original: sid,
@@ -218,6 +204,7 @@ fn get_source_id(
     }
 }
 
+// Very similar to get_source_id, but this function is used when the package_id is known.
 fn get_source_id_with_package_id(
     gctx: &GlobalContext,
     package_id: Option<PackageId>,
@@ -229,42 +216,24 @@ fn get_source_id_with_package_id(
         (Some(RegistryOrIndex::Index(url)), None) => (false, SourceId::for_registry(url)?),
         (Some(RegistryOrIndex::Registry(r)), None) => (false, SourceId::alt_registry(gctx, r)?),
         (Some(reg_or_index), Some(package_id)) => {
-            let sid = match reg_or_index {
-                RegistryOrIndex::Index(url) => SourceId::for_registry(url)?,
-                RegistryOrIndex::Registry(r) => SourceId::alt_registry(gctx, r)?,
-            };
+            let sid = get_initial_source_id_from_registry_or_index(gctx, reg_or_index)?;
             let package_source_id = package_id.source_id();
-            // Same registry, use the package's source.
-            if sid == package_source_id {
-                (true, sid)
+            // 1. Same registry, use the package's source.
+            // 2. Use the package's source if the specified registry is a replacement for the package's source.
+            if sid == package_source_id
+                || is_replacement_for_package_source(gctx, sid, package_source_id)?
+            {
+                (true, package_source_id)
             } else {
-                let pkg_source_replacement_sid = SourceConfigMap::new(gctx)?
-                    .load(package_source_id, &HashSet::new())?
-                    .replaced_source_id();
-                // Use the package's source if the specified registry is a replacement for the package's source.
-                if pkg_source_replacement_sid == sid {
-                    (true, package_source_id)
-                } else {
-                    (false, sid)
-                }
+                (false, sid)
             }
         }
     };
-    // Load source replacements that are built-in to Cargo.
-    let builtin_replacement_sid = SourceConfigMap::empty(gctx)?
-        .load(sid, &HashSet::new())?
-        .replaced_source_id();
-    let replacement_sid = SourceConfigMap::new(gctx)?
-        .load(sid, &HashSet::new())?
-        .replaced_source_id();
-    // Check if the user has configured source-replacement for the registry we are querying.
+
+    let (builtin_replacement_sid, replacement_sid) = get_replacement_source_ids(gctx, sid)?;
+
     if reg_or_index.is_none() && replacement_sid != builtin_replacement_sid {
-        // Neither --registry nor --index was passed and the user has configured source-replacement.
-        if let Some(replacement_name) = replacement_sid.alt_registry_key() {
-            bail!("crates-io is replaced with remote registry {replacement_name};\ninclude `--registry {replacement_name}` or `--registry crates-io`");
-        } else {
-            bail!("crates-io is replaced with non-remote-registry source {replacement_sid};\ninclude `--registry crates-io` to use crates.io");
-        }
+        bail!(gen_replacement_error(replacement_sid));
     } else {
         Ok((
             use_package_source_id,
@@ -276,6 +245,67 @@ fn get_source_id_with_package_id(
     }
 }
 
+fn get_initial_source_id(
+    gctx: &GlobalContext,
+    reg_or_index: Option<&RegistryOrIndex>,
+) -> CargoResult<SourceId> {
+    match reg_or_index {
+        None => SourceId::crates_io(gctx),
+        Some(reg_or_index) => get_initial_source_id_from_registry_or_index(gctx, reg_or_index),
+    }
+}
+
+fn get_initial_source_id_from_registry_or_index(
+    gctx: &GlobalContext,
+    reg_or_index: &RegistryOrIndex,
+) -> CargoResult<SourceId> {
+    match reg_or_index {
+        RegistryOrIndex::Index(url) => SourceId::for_registry(url),
+        RegistryOrIndex::Registry(r) => SourceId::alt_registry(gctx, r),
+    }
+}
+
+fn get_replacement_source_ids(
+    gctx: &GlobalContext,
+    sid: SourceId,
+) -> CargoResult<(SourceId, SourceId)> {
+    let builtin_replacement_sid = SourceConfigMap::empty(gctx)?
+        .load(sid, &HashSet::new())?
+        .replaced_source_id();
+    let replacement_sid = SourceConfigMap::new(gctx)?
+        .load(sid, &HashSet::new())?
+        .replaced_source_id();
+    Ok((builtin_replacement_sid, replacement_sid))
+}
+
+fn gen_replacement_error(replacement_sid: SourceId) -> String {
+    // Neither --registry nor --index was passed and the user has configured source-replacement.
+    let error_message = if let Some(replacement_name) = replacement_sid.alt_registry_key() {
+        format!(
+            "crates-io is replaced with remote registry {};\ninclude `--registry {}` or `--registry crates-io`",
+            replacement_name, replacement_name
+        )
+    } else {
+        format!(
+            "crates-io is replaced with non-remote-registry source {};\ninclude `--registry crates-io` to use crates.io",
+            replacement_sid
+        )
+    };
+
+    error_message
+}
+
+fn is_replacement_for_package_source(
+    gctx: &GlobalContext,
+    sid: SourceId,
+    package_source_id: SourceId,
+) -> CargoResult<bool> {
+    let pkg_source_replacement_sid = SourceConfigMap::new(gctx)?
+        .load(package_source_id, &HashSet::new())?
+        .replaced_source_id();
+    Ok(pkg_source_replacement_sid == sid)
+}
+
 struct RegistrySourceIds {
     /// Use when looking up the auth token, or writing out `Cargo.lock`
     original: SourceId,