Allow multiple indexes

charliermarsh · charliermarsh · commit ae705cb00d32 · 2024-09-30T17:21:31.000-04:00
diff --git a/crates/uv-resolver/src/error.rs b/crates/uv-resolver/src/error.rs
@@ -52,6 +52,16 @@ pub enum ResolveError {
         fork_markers: MarkerTree,
     },
 
+    #[error("Requirements contain conflicting indexes for package `{0}`:\n- {}", _1.join("\n- "))]
+    ConflictingIndexesUniversal(PackageName, Vec<String>),
+
+    #[error("Requirements contain conflicting indexes for package `{package_name}` in split `{fork_markers:?}`:\n- {}", indexes.join("\n- "))]
+    ConflictingIndexesFork {
+        package_name: PackageName,
+        indexes: Vec<String>,
+        fork_markers: MarkerTree,
+    },
+
     #[error("Requirements contain conflicting indexes for package `{0}`: `{1}` vs. `{2}`")]
     ConflictingIndexes(PackageName, String, String),
 
diff --git a/crates/uv-resolver/src/fork_indexes.rs b/crates/uv-resolver/src/fork_indexes.rs
@@ -0,0 +1,57 @@
+use std::collections::hash_map::Entry;
+
+use rustc_hash::FxHashMap;
+
+use distribution_types::IndexUrl;
+use uv_normalize::PackageName;
+
+use crate::resolver::ResolverMarkers;
+use crate::ResolveError;
+
+/// See [`crate::resolver::ForkState`].
+#[derive(Default, Debug, Clone)]
+pub(crate) struct ForkIndexes(FxHashMap<PackageName, IndexUrl>);
+
+impl ForkIndexes {
+    /// Get the [`IndexUrl`] previously used for a package in this fork.
+    pub(crate) fn get(&self, package_name: &PackageName) -> Option<&IndexUrl> {
+        self.0.get(package_name)
+    }
+
+    /// Check that this is the only [`IndexUrl`] used for this package in this fork.
+    pub(crate) fn insert(
+        &mut self,
+        package_name: &PackageName,
+        index: &IndexUrl,
+        fork_markers: &ResolverMarkers,
+    ) -> Result<(), ResolveError> {
+        match self.0.entry(package_name.clone()) {
+            Entry::Occupied(previous) => {
+                if previous.get() != index {
+                    let mut conflicts = vec![previous.get().to_string(), index.to_string()];
+                    conflicts.sort();
+                    return match fork_markers {
+                        ResolverMarkers::Universal { .. }
+                        | ResolverMarkers::SpecificEnvironment(_) => {
+                            Err(ResolveError::ConflictingIndexesUniversal(
+                                package_name.clone(),
+                                conflicts,
+                            ))
+                        }
+                        ResolverMarkers::Fork(fork_markers) => {
+                            Err(ResolveError::ConflictingIndexesFork {
+                                package_name: package_name.clone(),
+                                indexes: conflicts,
+                                fork_markers: fork_markers.clone(),
+                            })
+                        }
+                    };
+                }
+            }
+            Entry::Vacant(vacant) => {
+                vacant.insert(index.clone());
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/crates/uv-resolver/src/fork_urls.rs b/crates/uv-resolver/src/fork_urls.rs
@@ -9,7 +9,7 @@ use uv_normalize::PackageName;
 use crate::resolver::ResolverMarkers;
 use crate::ResolveError;
 
-/// See [`crate::resolver::SolveState`].
+/// See [`crate::resolver::ForkState`].
 #[derive(Default, Debug, Clone)]
 pub(crate) struct ForkUrls(FxHashMap<PackageName, VerbatimParsedUrl>);
 
diff --git a/crates/uv-resolver/src/lib.rs b/crates/uv-resolver/src/lib.rs
@@ -34,6 +34,7 @@ mod error;
 mod exclude_newer;
 mod exclusions;
 mod flat_index;
+mod fork_indexes;
 mod fork_urls;
 mod graph_ops;
 mod lock;
diff --git a/crates/uv-resolver/src/resolution/graph.rs b/crates/uv-resolver/src/resolution/graph.rs
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 
 use distribution_types::{
-    Dist, DistributionMetadata, Name, ResolutionDiagnostic, ResolvedDist, VersionId,
+    Dist, DistributionMetadata, IndexUrl, Name, ResolutionDiagnostic, ResolvedDist, VersionId,
     VersionOrUrlRef,
 };
 use indexmap::IndexSet;
@@ -88,6 +88,7 @@ struct PackageRef<'a> {
     package_name: &'a PackageName,
     version: &'a Version,
     url: Option<&'a VerbatimParsedUrl>,
+    index: Option<&'a IndexUrl>,
     extra: Option<&'a ExtraName>,
     group: Option<&'a GroupName>,
 }
@@ -284,6 +285,7 @@ impl ResolutionGraph {
                 package_name: from,
                 version: &edge.from_version,
                 url: edge.from_url.as_ref(),
+                index: edge.from_index.as_ref(),
                 extra: edge.from_extra.as_ref(),
                 group: edge.from_dev.as_ref(),
             }]
@@ -292,6 +294,7 @@ impl ResolutionGraph {
             package_name: &edge.to,
             version: &edge.to_version,
             url: edge.to_url.as_ref(),
+            index: edge.to_index.as_ref(),
             extra: edge.to_extra.as_ref(),
             group: edge.to_dev.as_ref(),
         }];
@@ -320,7 +323,7 @@ impl ResolutionGraph {
         diagnostics: &mut Vec<ResolutionDiagnostic>,
         preferences: &Preferences,
         pins: &FilePins,
-        index: &InMemoryIndex,
+        in_memory: &InMemoryIndex,
         git: &GitResolver,
         package: &'a ResolutionPackage,
         version: &'a Version,
@@ -330,16 +333,18 @@ impl ResolutionGraph {
             extra,
             dev,
             url,
+            index,
         } = &package;
         // Map the package to a distribution.
         let (dist, hashes, metadata) = Self::parse_dist(
             name,
+            index.as_ref(),
             url.as_ref(),
             version,
             pins,
             diagnostics,
             preferences,
-            index,
+            in_memory,
             git,
         )?;
 
@@ -366,7 +371,7 @@ impl ResolutionGraph {
         }
 
         // Add the distribution to the graph.
-        let index = petgraph.add_node(ResolutionGraphNode::Dist(AnnotatedDist {
+        let node = petgraph.add_node(ResolutionGraphNode::Dist(AnnotatedDist {
             dist,
             name: name.clone(),
             version: version.clone(),
@@ -381,22 +386,24 @@ impl ResolutionGraph {
                 package_name: name,
                 version,
                 url: url.as_ref(),
+                index: index.as_ref(),
                 extra: extra.as_ref(),
                 group: dev.as_ref(),
             },
-            index,
+            node,
         );
         Ok(())
     }
 
     fn parse_dist(
         name: &PackageName,
+        index: Option<&IndexUrl>,
         url: Option<&VerbatimParsedUrl>,
         version: &Version,
         pins: &FilePins,
         diagnostics: &mut Vec<ResolutionDiagnostic>,
         preferences: &Preferences,
-        index: &InMemoryIndex,
+        in_memory: &InMemoryIndex,
         git: &GitResolver,
     ) -> Result<(ResolvedDist, Vec<HashDigest>, Option<Metadata>), ResolveError> {
         Ok(if let Some(url) = url {
@@ -406,14 +413,24 @@ impl ResolutionGraph {
             let version_id = VersionId::from_url(&url.verbatim);
 
             // Extract the hashes.
-            let hashes =
-                Self::get_hashes(name, Some(url), &version_id, version, preferences, index);
+            let hashes = Self::get_hashes(
+                name,
+                index,
+                Some(url),
+                &version_id,
+                version,
+                preferences,
+                in_memory,
+            );
 
             // Extract the metadata.
             let metadata = {
-                let response = index.distributions().get(&version_id).unwrap_or_else(|| {
-                    panic!("Every URL distribution should have metadata: {version_id:?}")
-                });
+                let response = in_memory
+                    .distributions()
+                    .get(&version_id)
+                    .unwrap_or_else(|| {
+                        panic!("Every URL distribution should have metadata: {version_id:?}")
+                    });
 
                 let MetadataResponse::Found(archive) = &*response else {
                     panic!("Every URL distribution should have metadata: {version_id:?}")
@@ -449,17 +466,28 @@ impl ResolutionGraph {
             }
 
             // Extract the hashes.
-            let hashes = Self::get_hashes(name, None, &version_id, version, preferences, index);
+            let hashes = Self::get_hashes(
+                name,
+                index,
+                None,
+                &version_id,
+                version,
+                preferences,
+                in_memory,
+            );
 
             // Extract the metadata.
             let metadata = {
-                index.distributions().get(&version_id).and_then(|response| {
-                    if let MetadataResponse::Found(archive) = &*response {
-                        Some(archive.metadata.clone())
-                    } else {
-                        None
-                    }
-                })
+                in_memory
+                    .distributions()
+                    .get(&version_id)
+                    .and_then(|response| {
+                        if let MetadataResponse::Found(archive) = &*response {
+                            Some(archive.metadata.clone())
+                        } else {
+                            None
+                        }
+                    })
             };
 
             (dist, hashes, metadata)
@@ -470,11 +498,12 @@ impl ResolutionGraph {
     /// lockfile.
     fn get_hashes(
         name: &PackageName,
+        index: Option<&IndexUrl>,
         url: Option<&VerbatimParsedUrl>,
         version_id: &VersionId,
         version: &Version,
         preferences: &Preferences,
-        index: &InMemoryIndex,
+        in_memory: &InMemoryIndex,
     ) -> Vec<HashDigest> {
         // 1. Look for hashes from the lockfile.
         if let Some(digests) = preferences.match_hashes(name, version) {
@@ -484,7 +513,7 @@ impl ResolutionGraph {
         }
 
         // 2. Look for hashes for the distribution (i.e., the specific wheel or source distribution).
-        if let Some(metadata_response) = index.distributions().get(version_id) {
+        if let Some(metadata_response) = in_memory.distributions().get(version_id) {
             if let MetadataResponse::Found(ref archive) = *metadata_response {
                 let mut digests = archive.hashes.clone();
                 digests.sort_unstable();
@@ -496,7 +525,13 @@ impl ResolutionGraph {
 
         // 3. Look for hashes from the registry, which are served at the package level.
         if url.is_none() {
-            if let Some(versions_response) = index.packages().get(name) {
+            let versions_response = if let Some(index) = index {
+                in_memory.explicit().get(&(name.clone(), index.clone()))
+            } else {
+                in_memory.implicit().get(name)
+            };
+
+            if let Some(versions_response) = versions_response {
                 if let VersionsResponse::Found(ref version_maps) = *versions_response {
                     if let Some(digests) = version_maps
                         .iter()
diff --git a/crates/uv-resolver/src/resolver/batch_prefetch.rs b/crates/uv-resolver/src/resolver/batch_prefetch.rs
@@ -6,7 +6,7 @@ use rustc_hash::FxHashMap;
 use tokio::sync::mpsc::Sender;
 use tracing::{debug, trace};
 
-use distribution_types::{CompatibleDist, DistributionMetadata, IndexCapabilities};
+use distribution_types::{CompatibleDist, DistributionMetadata, IndexCapabilities, IndexUrl};
 use pep440_rs::Version;
 
 use crate::candidate_selector::CandidateSelector;
@@ -47,11 +47,12 @@ impl BatchPrefetcher {
     pub(crate) fn prefetch_batches(
         &mut self,
         next: &PubGrubPackage,
+        index: Option<&IndexUrl>,
         version: &Version,
         current_range: &Range<Version>,
         python_requirement: &PythonRequirement,
         request_sink: &Sender<Request>,
-        index: &InMemoryIndex,
+        in_memory: &InMemoryIndex,
         capabilities: &IndexCapabilities,
         selector: &CandidateSelector,
         markers: &ResolverMarkers,
@@ -73,10 +74,17 @@ impl BatchPrefetcher {
         let total_prefetch = min(num_tried, 50);
 
         // This is immediate, we already fetched the version map.
-        let versions_response = index
-            .packages()
-            .wait_blocking(name)
-            .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?;
+        let versions_response = if let Some(index) = index {
+            in_memory
+                .explicit()
+                .wait_blocking(&(name.clone(), index.clone()))
+                .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
+        } else {
+            in_memory
+                .implicit()
+                .wait_blocking(name)
+                .ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
+        };
 
         let VersionsResponse::Found(ref version_map) = *versions_response else {
             return Ok(());
@@ -191,7 +199,7 @@ impl BatchPrefetcher {
             );
             prefetch_count += 1;
 
-            if index.distributions().register(candidate.version_id()) {
+            if in_memory.distributions().register(candidate.version_id()) {
                 let request = Request::from(dist);
                 request_sink.blocking_send(request)?;
             }
diff --git a/crates/uv-resolver/src/resolver/fork_map.rs b/crates/uv-resolver/src/resolver/fork_map.rs
@@ -44,6 +44,11 @@ impl<T> ForkMap<T> {
         !self.get(package_name, markers).is_empty()
     }
 
+    /// Returns `true` if the map contains any values for a package.
+    pub(crate) fn contains_key(&self, package_name: &PackageName) -> bool {
+        self.0.contains_key(package_name)
+    }
+
     /// Returns a list of values associated with a package that are compatible with the given fork.
     ///
     /// Compatibility implies that the markers on the requirement that contained this value
diff --git a/crates/uv-resolver/src/resolver/index.rs b/crates/uv-resolver/src/resolver/index.rs
@@ -1,7 +1,7 @@
 use std::hash::BuildHasherDefault;
 use std::sync::Arc;
 
-use distribution_types::VersionId;
+use distribution_types::{IndexUrl, VersionId};
 use once_map::OnceMap;
 use rustc_hash::FxHasher;
 use uv_normalize::PackageName;
@@ -16,7 +16,9 @@ pub struct InMemoryIndex(Arc<SharedInMemoryIndex>);
 struct SharedInMemoryIndex {
     /// A map from package name to the metadata for that package and the index where the metadata
     /// came from.
-    packages: FxOnceMap<PackageName, Arc<VersionsResponse>>,
+    implicit: FxOnceMap<PackageName, Arc<VersionsResponse>>,
+
+    explicit: FxOnceMap<(PackageName, IndexUrl), Arc<VersionsResponse>>,
 
     /// A map from package ID to metadata for that distribution.
     distributions: FxOnceMap<VersionId, Arc<MetadataResponse>>,
@@ -26,8 +28,13 @@ pub(crate) type FxOnceMap<K, V> = OnceMap<K, V, BuildHasherDefault<FxHasher>>;
 
 impl InMemoryIndex {
     /// Returns a reference to the package metadata map.
-    pub fn packages(&self) -> &FxOnceMap<PackageName, Arc<VersionsResponse>> {
-        &self.0.packages
+    pub fn implicit(&self) -> &FxOnceMap<PackageName, Arc<VersionsResponse>> {
+        &self.0.implicit
+    }
+
+    /// Returns a reference to the package metadata map.
+    pub fn explicit(&self) -> &FxOnceMap<(PackageName, IndexUrl), Arc<VersionsResponse>> {
+        &self.0.explicit
     }
 
     /// Returns a reference to the distribution metadata map.
diff --git a/crates/uv-resolver/src/resolver/indexes.rs b/crates/uv-resolver/src/resolver/indexes.rs
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
diff --git a/crates/uv-workspace/src/pyproject.rs b/crates/uv-workspace/src/pyproject.rs
diff --git a/crates/uv/tests/lock.rs b/crates/uv/tests/lock.rs
