diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 8ebebc263ff52..e821c5c25a4dd 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -2445,7 +2445,7 @@ impl ResolverState ResolverState ResolverState 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::(&datum.version) .expect("archived version always deserializes"); - // Deserialize the metadata. - let core_metadatum = - rkyv::deserialize::, 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 { 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, +} + +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 { + 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, + /// 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, /// The raw simple metadata from which `PrioritizedDist`s should /// be constructed. simple_metadata: OwnedArchive, @@ -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 { + 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::(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, + /// A lazy index into [`SimpleDetailMetadata`], if present. + simple: Option, } /// Represents a lazily initialized `PrioritizedDist`.