Use a compact index for lazy version maps

crates/uv-resolver/src/resolver/mod.rs

@@ -2445,7 +2445,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                                     return Ok(Some(Response::Dist {
                                         dist,
                                         metadata: MetadataResponse::Found(
-                                            ArchiveMetadata::from_metadata23(metadata.clone()),
+                                            ArchiveMetadata::from_metadata23(metadata),
                                         ),
                                     }));
                                 }
@@ -2468,7 +2468,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                                 return Ok(Some(Response::Dist {
                                     dist,
                                     metadata: MetadataResponse::Found(
-                                        ArchiveMetadata::from_metadata23(metadata.clone()),
+                                        ArchiveMetadata::from_metadata23(metadata),
                                     ),
                                 }));
                             }
@@ -2578,9 +2578,8 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                         if version_map.index() == dist.index() {
                             debug!("Found registry-provided metadata for: {dist}");
 
-                            let metadata = MetadataResponse::Found(
-                                ArchiveMetadata::from_metadata23(metadata.clone()),
-                            );
+                            let metadata =
+                                MetadataResponse::Found(ArchiveMetadata::from_metadata23(metadata));
 
                             let dist = dist.to_owned();
                             if &package_name != dist.name() {

crates/uv-resolver/src/version_map.rs

@@ -1,11 +1,10 @@
+use std::collections::BTreeMap;
 use std::collections::Bound;
-use std::collections::btree_map::{BTreeMap, Entry};
 use std::ops::RangeBounds;
 use std::sync::OnceLock;
 
 use jiff::Timestamp;
 use pubgrub::Ranges;
-use rustc_hash::FxHashMap;
 use tracing::{instrument, trace};
 
 use uv_client::{FlatIndexEntry, OwnedArchive, SimpleDetailMetadata, VersionFiles};
@@ -59,68 +58,45 @@ impl VersionMap {
     ) -> Self {
         let mut stable = false;
         let mut local = false;
-        let mut map = BTreeMap::new();
-        let mut core_metadata = FxHashMap::default();
+        let mut entries = Vec::with_capacity(simple_metadata.iter().size_hint().0);
         // Create stubs for each entry in simple metadata. The full conversion
         // from a `VersionFiles` to a PrioritizedDist for each version
         // isn't done until that specific version is requested.
         for (datum_index, datum) in simple_metadata.iter().enumerate() {
-            // Deserialize the version.
             let version = rkyv::deserialize::<Version, rkyv::rancor::Error>(&datum.version)
                 .expect("archived version always deserializes");
 
-            // Deserialize the metadata.
-            let core_metadatum =
-                rkyv::deserialize::<Option<ResolutionMetadata>, rkyv::rancor::Error>(
-                    &datum.metadata,
-                )
-                .expect("archived metadata always deserializes");
-            if let Some(core_metadatum) = core_metadatum {
-                core_metadata.insert(version.clone(), core_metadatum);
-            }
-
             stable |= version.is_stable();
             local |= version.is_local();
-            map.insert(
-                version,
-                LazyPrioritizedDist::OnlySimple(SimplePrioritizedDist {
-                    datum_index,
-                    dist: OnceLock::new(),
-                }),
+            debug_assert!(
+                entries
+                    .last()
+                    .is_none_or(|entry: &VersionMapLazyEntry| entry.version < version),
+                "simple metadata versions must be sorted and unique"
             );
-        }
-        // If a set of flat distributions have been given, we need to add those
-        // to our map of entries as well.
-        for (version, prioritized_dist) in flat_index.into_iter().flatten() {
-            stable |= version.is_stable();
-            match map.entry(version) {
-                Entry::Vacant(e) => {
-                    e.insert(LazyPrioritizedDist::OnlyFlat(prioritized_dist));
-                }
-                // When there is both a `VersionFiles` (from the "simple"
-                // metadata) and a flat distribution for the same version of
-                // a package, we store both and "merge" them into a single
-                // `PrioritizedDist` upon access later.
-                Entry::Occupied(e) => match e.remove_entry() {
-                    (version, LazyPrioritizedDist::OnlySimple(simple_dist)) => {
-                        map.insert(
-                            version,
-                            LazyPrioritizedDist::Both {
-                                flat: prioritized_dist,
-                                simple: simple_dist,
-                            },
-                        );
-                    }
-                    _ => unreachable!(),
+            entries.push(VersionMapLazyEntry {
+                version,
+                dist: LazyPrioritizedDist {
+                    flat: None,
+                    simple: Some(SimplePrioritizedDist {
+                        datum_index,
+                        dist: OnceLock::new(),
+                    }),
                 },
-            }
+            });
+        }
+        let mut map = VersionMapLazyIndex { entries };
+        // If a set of flat distributions have been given, linearly merge the
+        // already sorted flat entries with the archive-ordered simple vector.
+        if let Some(flat_index) = flat_index {
+            stable |= flat_index.iter().any(|(version, _)| version.is_stable());
+            map = map.merge_flat(flat_index);
         }
         Self {
             inner: VersionMapInner::Lazy(VersionMapLazy {
                 map,
                 stable,
                 local,
-                core_metadata,
                 simple_metadata,
                 no_binary: build_options.no_binary_package(package_name),
                 no_build: build_options.no_build_package(package_name),
@@ -160,10 +136,10 @@ impl VersionMap {
     }
 
     /// Return the [`ResolutionMetadata`] for the given version, if any.
-    pub(crate) fn get_metadata(&self, version: &Version) -> Option<&ResolutionMetadata> {
+    pub(crate) fn get_metadata(&self, version: &Version) -> Option<ResolutionMetadata> {
         match self.inner {
             VersionMapInner::Eager(_) => None,
-            VersionMapInner::Lazy(ref lazy) => lazy.core_metadata.get(version),
+            VersionMapInner::Lazy(ref lazy) => lazy.get_metadata(version),
         }
     }
 
@@ -223,14 +199,15 @@ impl VersionMap {
                     ))
                 }
                 VersionMapInner::Lazy(ref lazy) => {
-                    either::Either::Right(lazy.map.get_key_value(version).into_iter().map(
-                        move |(version, dist)| {
-                            let version_map_dist = VersionMapDistHandle {
-                                inner: VersionMapDistHandleInner::Lazy { lazy, dist },
-                            };
-                            (version, version_map_dist)
-                        },
-                    ))
+                    either::Either::Right(lazy.map.get(version).into_iter().map(move |entry| {
+                        let version_map_dist = VersionMapDistHandle {
+                            inner: VersionMapDistHandleInner::Lazy {
+                                lazy,
+                                dist: &entry.dist,
+                            },
+                        };
+                        (&entry.version, version_map_dist)
+                    }))
                 }
             })
         } else {
@@ -246,12 +223,15 @@ impl VersionMap {
                     ))
                 }
                 VersionMapInner::Lazy(ref lazy) => {
-                    either::Either::Right(lazy.map.range(BoundingRange::from(range)).map(
-                        |(version, dist)| {
+                    either::Either::Right(lazy.map.range(BoundingRange::from(range)).iter().map(
+                        |entry| {
                             let version_map_dist = VersionMapDistHandle {
-                                inner: VersionMapDistHandleInner::Lazy { lazy, dist },
+                                inner: VersionMapDistHandleInner::Lazy {
+                                    lazy,
+                                    dist: &entry.dist,
+                                },
                             };
-                            (version, version_map_dist)
+                            (&entry.version, version_map_dist)
                         },
                     ))
                 }
@@ -370,6 +350,98 @@ struct VersionMapEager {
     local: bool,
 }
 
+/// An entry in the immutable lazy version index.
+#[derive(Debug)]
+struct VersionMapLazyEntry {
+    version: Version,
+    dist: LazyPrioritizedDist,
+}
+
+/// A compact immutable version index, ordered by native PEP 440 versions.
+#[derive(Debug)]
+struct VersionMapLazyIndex {
+    entries: Vec<VersionMapLazyEntry>,
+}
+
+impl VersionMapLazyIndex {
+    /// Merge the sorted flat index into the sorted simple entries in one pass.
+    fn merge_flat(self, flat_index: FlatDistributions) -> Self {
+        let flat_count = flat_index.iter().size_hint().0;
+        let mut merged = Vec::with_capacity(self.entries.len() + flat_count);
+        let mut simple = self.entries.into_iter().peekable();
+        let mut flat = flat_index.into_iter().peekable();
+        let flat_entry = |(version, dist)| VersionMapLazyEntry {
+            version,
+            dist: LazyPrioritizedDist {
+                flat: Some(dist),
+                simple: None,
+            },
+        };
+        while let (Some(simple_entry), Some((flat_version, _))) = (simple.peek(), flat.peek()) {
+            match simple_entry.version.cmp(flat_version) {
+                std::cmp::Ordering::Less => {
+                    if let Some(entry) = simple.next() {
+                        merged.push(entry);
+                    }
+                }
+                std::cmp::Ordering::Greater => {
+                    if let Some(entry) = flat.next() {
+                        merged.push(flat_entry(entry));
+                    }
+                }
+                std::cmp::Ordering::Equal => {
+                    if let (Some(mut entry), Some((_, dist))) = (simple.next(), flat.next()) {
+                        entry.dist.flat = Some(dist);
+                        merged.push(entry);
+                    }
+                }
+            }
+        }
+        merged.extend(simple);
+        merged.extend(flat.map(flat_entry));
+
+        Self { entries: merged }
+    }
+
+    fn get(&self, version: &Version) -> Option<&VersionMapLazyEntry> {
+        let index = self
+            .entries
+            .binary_search_by(|entry| entry.version.cmp(version))
+            .ok()?;
+        self.entries.get(index)
+    }
+
+    fn keys(&self) -> impl DoubleEndedIterator<Item = &Version> {
+        self.entries.iter().map(|entry| &entry.version)
+    }
+
+    fn range(&self, range: BoundingRange<'_>) -> &[VersionMapLazyEntry] {
+        let start = match range.min {
+            Bound::Included(version) => self
+                .entries
+                .partition_point(|entry| entry.version < *version),
+            Bound::Excluded(version) => self
+                .entries
+                .partition_point(|entry| entry.version <= *version),
+            Bound::Unbounded => 0,
+        };
+        let end = match range.max {
+            Bound::Included(version) => self
+                .entries
+                .partition_point(|entry| entry.version <= *version),
+            Bound::Excluded(version) => self
+                .entries
+                .partition_point(|entry| entry.version < *version),
+            Bound::Unbounded => self.entries.len(),
+        };
+        self.entries.get(start..end).unwrap_or_default()
+    }
+
+    fn len(&self) -> usize {
+        self.entries.len()
+    }
+}
+
 /// A map that lazily materializes some prioritized distributions upon access.
 ///
 /// The idea here is that some packages have a lot of versions published, and
@@ -380,14 +452,12 @@ struct VersionMapEager {
 /// provide substantial savings in some cases.
 #[derive(Debug)]
 struct VersionMapLazy {
-    /// A map from version to possibly-initialized distribution.
-    map: BTreeMap<Version, LazyPrioritizedDist>,
+    /// An immutable archive-order index from version to possibly-initialized distribution.
+    map: VersionMapLazyIndex,
     /// Whether the version map contains at least one stable (non-pre-release) version.
     stable: bool,
     /// Whether the version map contains at least one local version.
     local: bool,
-    /// The pre-populated metadata for each version.
-    core_metadata: FxHashMap<Version, ResolutionMetadata>,
     /// The raw simple metadata from which `PrioritizedDist`s should
     /// be constructed.
     simple_metadata: OwnedArchive<SimpleDetailMetadata>,
@@ -414,11 +484,23 @@ struct VersionMapLazy {
 }
 
 impl VersionMapLazy {
+    /// Returns the registry-provided metadata for the given version, if it exists.
+    fn get_metadata(&self, version: &Version) -> Option<ResolutionMetadata> {
+        let archived = self
+            .map
+            .get(version)
+            .and_then(|entry| entry.dist.simple.as_ref())
+            .and_then(|simple| self.simple_metadata.datum(simple.datum_index))
+            .and_then(|datum| datum.metadata.as_ref())?;
+        Some(
+            rkyv::deserialize::<ResolutionMetadata, rkyv::rancor::Error>(archived)
+                .expect("archived metadata always deserializes"),
+        )
+    }
+
     /// Returns the distribution for the given version, if it exists.
     fn get(&self, version: &Version) -> Option<&PrioritizedDist> {
-        let lazy_dist = self.map.get(version)?;
-        let priority_dist = self.get_lazy(lazy_dist)?;
-        Some(priority_dist)
+        self.get_lazy(&self.map.get(version)?.dist)
     }
 
     /// Given a reference to a possibly-initialized distribution that is in
@@ -427,13 +509,11 @@ impl VersionMapLazy {
     /// When both a flat and simple distribution are present internally, they
     /// are merged automatically.
     fn get_lazy<'p>(&'p self, lazy_dist: &'p LazyPrioritizedDist) -> Option<&'p PrioritizedDist> {
-        match *lazy_dist {
-            LazyPrioritizedDist::OnlyFlat(ref dist) => Some(dist),
-            LazyPrioritizedDist::OnlySimple(ref dist) => self.get_simple(None, dist),
-            LazyPrioritizedDist::Both {
-                ref flat,
-                ref simple,
-            } => self.get_simple(Some(flat), simple),
+        match (&lazy_dist.flat, &lazy_dist.simple) {
+            (Some(flat), Some(simple)) => self.get_simple(Some(flat), simple),
+            (Some(flat), None) => Some(flat),
+            (None, Some(simple)) => self.get_simple(None, simple),
+            (None, None) => None,
         }
     }
 
@@ -665,22 +745,13 @@ impl VersionMapLazy {
     }
 }
 
-/// Represents a possibly initialized [`PrioritizedDist`] for
-/// a single version of a package.
+/// Represents a possibly initialized [`PrioritizedDist`] for a package version.
 #[derive(Debug)]
-enum LazyPrioritizedDist {
-    /// Represents an eagerly constructed distribution from a
-    /// `FlatDistributions`.
-    OnlyFlat(PrioritizedDist),
-    /// Represents a lazily constructed distribution from an index into a
-    /// `VersionFiles` from `SimpleDetailMetadata`.
-    OnlySimple(SimplePrioritizedDist),
-    /// Combines the above. This occurs when we have data from both a flat
-    /// distribution and a simple distribution.
-    Both {
-        flat: PrioritizedDist,
-        simple: SimplePrioritizedDist,
-    },
+struct LazyPrioritizedDist {
+    /// An eagerly constructed distribution from [`FlatDistributions`], if present.
+    flat: Option<PrioritizedDist>,
+    /// A lazy index into [`SimpleDetailMetadata`], if present.
+    simple: Option<SimplePrioritizedDist>,
 }
 
 /// Represents a lazily initialized `PrioritizedDist`.