1use std::collections::BTreeMap;
2use std::collections::Bound;
3use std::ops::RangeBounds;
4use std::sync::OnceLock;
5
6use jiff::Timestamp;
7use pubgrub::Ranges;
8use tracing::{instrument, trace};
9
10use uv_client::{FlatIndexEntry, OwnedArchive, SimpleDetailMetadata, VersionFiles};
11use uv_configuration::BuildOptions;
12use uv_distribution_filename::{DistFilename, WheelFilename};
13use uv_distribution_types::{
14 HashComparison, IncompatibleSource, IncompatibleWheel, IndexUrl, PrioritizedDist,
15 RegistryBuiltWheel, RegistrySourceDist, RequiresPython, SourceDistCompatibility,
16 WheelCompatibility,
17};
18use uv_normalize::PackageName;
19use uv_pep440::Version;
20use uv_platform_tags::{IncompatibleTag, TagCompatibility, Tags};
21use uv_pypi_types::{HashDigest, ResolutionMetadata, Yanked};
22use uv_types::HashStrategy;
23use uv_warnings::warn_user_once;
24
25use crate::flat_index::FlatDistributions;
26use crate::yanks::AllowedYanks;
27
28#[derive(Debug)]
30pub struct VersionMap {
31 inner: VersionMapInner,
33}
34
35impl VersionMap {
36 #[instrument(skip_all, fields(package_name))]
46 pub(crate) fn from_simple_metadata(
47 simple_metadata: OwnedArchive<SimpleDetailMetadata>,
48 package_name: &PackageName,
49 index: &IndexUrl,
50 tags: Option<&Tags>,
51 requires_python: &RequiresPython,
52 allowed_yanks: &AllowedYanks,
53 hasher: &HashStrategy,
54 included_version_cutoff: Option<Timestamp>,
55 available_version_cutoff: Option<Timestamp>,
56 flat_index: Option<FlatDistributions>,
57 build_options: &BuildOptions,
58 ) -> Self {
59 let mut stable = false;
60 let mut local = false;
61 let mut entries = Vec::with_capacity(simple_metadata.iter().size_hint().0);
62 for (datum_index, datum) in simple_metadata.iter().enumerate() {
66 let version = rkyv::deserialize::<Version, rkyv::rancor::Error>(&datum.version)
67 .expect("archived version always deserializes");
68
69 stable |= version.is_stable();
70 local |= version.is_local();
71 debug_assert!(
72 entries
73 .last()
74 .is_none_or(|entry: &VersionMapLazyEntry| entry.version < version),
75 "simple metadata versions must be sorted and unique"
76 );
77 entries.push(VersionMapLazyEntry {
78 version,
79 dist: LazyPrioritizedDist {
80 flat: None,
81 simple: Some(SimplePrioritizedDist {
82 datum_index,
83 dist: OnceLock::new(),
84 }),
85 },
86 });
87 }
88 let mut map = VersionMapLazyIndex { entries };
89 if let Some(flat_index) = flat_index {
92 stable |= flat_index.iter().any(|(version, _)| version.is_stable());
93 map = map.merge_flat(flat_index);
94 }
95 Self {
96 inner: VersionMapInner::Lazy(VersionMapLazy {
97 map,
98 stable,
99 local,
100 simple_metadata,
101 no_binary: build_options.no_binary_package(package_name),
102 no_build: build_options.no_build_package(package_name),
103 index: index.clone(),
104 tags: tags.cloned(),
105 allowed_yanks: allowed_yanks.clone(),
106 hasher: hasher.clone(),
107 requires_python: requires_python.clone(),
108 included_version_cutoff,
109 available_version_cutoff,
110 }),
111 }
112 }
113
114 #[instrument(skip_all, fields(package_name))]
115 pub(crate) fn from_flat_metadata(
116 flat_metadata: Vec<FlatIndexEntry>,
117 tags: Option<&Tags>,
118 hasher: &HashStrategy,
119 build_options: &BuildOptions,
120 ) -> Self {
121 let mut stable = false;
122 let mut local = false;
123 let mut map = BTreeMap::new();
124
125 for (version, prioritized_dist) in
126 FlatDistributions::from_entries(flat_metadata, tags, hasher, build_options)
127 {
128 stable |= version.is_stable();
129 local |= version.is_local();
130 map.insert(version, prioritized_dist);
131 }
132
133 Self {
134 inner: VersionMapInner::Eager(VersionMapEager { map, stable, local }),
135 }
136 }
137
138 pub(crate) fn get_metadata(&self, version: &Version) -> Option<ResolutionMetadata> {
140 match self.inner {
141 VersionMapInner::Eager(_) => None,
142 VersionMapInner::Lazy(ref lazy) => lazy.get_metadata(version),
143 }
144 }
145
146 pub(crate) fn get(&self, version: &Version) -> Option<&PrioritizedDist> {
148 match self.inner {
149 VersionMapInner::Eager(ref eager) => eager.map.get(version),
150 VersionMapInner::Lazy(ref lazy) => lazy.get(version),
151 }
152 }
153
154 pub(crate) fn versions(&self) -> impl DoubleEndedIterator<Item = &Version> {
156 match &self.inner {
157 VersionMapInner::Eager(eager) => either::Either::Left(eager.map.keys()),
158 VersionMapInner::Lazy(lazy) => either::Either::Right(lazy.map.keys()),
159 }
160 }
161
162 pub(crate) fn index(&self) -> Option<&IndexUrl> {
164 match &self.inner {
165 VersionMapInner::Eager(_) => None,
166 VersionMapInner::Lazy(lazy) => Some(&lazy.index),
167 }
168 }
169
170 pub(crate) fn included_version_cutoff(&self) -> Option<&Timestamp> {
172 match &self.inner {
173 VersionMapInner::Eager(_) => None,
174 VersionMapInner::Lazy(lazy) => lazy.included_version_cutoff.as_ref(),
175 }
176 }
177
178 pub(crate) fn iter(
185 &self,
186 range: &Ranges<Version>,
187 ) -> impl DoubleEndedIterator<Item = (&Version, VersionMapDistHandle<'_>)> {
188 if let Some(version) = range.as_singleton() {
190 either::Either::Left(match self.inner {
191 VersionMapInner::Eager(ref eager) => {
192 either::Either::Left(eager.map.get_key_value(version).into_iter().map(
193 move |(version, dist)| {
194 let version_map_dist = VersionMapDistHandle {
195 inner: VersionMapDistHandleInner::Eager(dist),
196 };
197 (version, version_map_dist)
198 },
199 ))
200 }
201 VersionMapInner::Lazy(ref lazy) => {
202 either::Either::Right(lazy.map.get(version).into_iter().map(move |entry| {
203 let version_map_dist = VersionMapDistHandle {
204 inner: VersionMapDistHandleInner::Lazy {
205 lazy,
206 dist: &entry.dist,
207 },
208 };
209 (&entry.version, version_map_dist)
210 }))
211 }
212 })
213 } else {
214 either::Either::Right(match self.inner {
215 VersionMapInner::Eager(ref eager) => {
216 either::Either::Left(eager.map.range(BoundingRange::from(range)).map(
217 |(version, dist)| {
218 let version_map_dist = VersionMapDistHandle {
219 inner: VersionMapDistHandleInner::Eager(dist),
220 };
221 (version, version_map_dist)
222 },
223 ))
224 }
225 VersionMapInner::Lazy(ref lazy) => {
226 either::Either::Right(lazy.map.range(BoundingRange::from(range)).iter().map(
227 |entry| {
228 let version_map_dist = VersionMapDistHandle {
229 inner: VersionMapDistHandleInner::Lazy {
230 lazy,
231 dist: &entry.dist,
232 },
233 };
234 (&entry.version, version_map_dist)
235 },
236 ))
237 }
238 })
239 }
240 }
241
242 pub(crate) fn hashes(&self, version: &Version) -> Option<&[HashDigest]> {
244 match self.inner {
245 VersionMapInner::Eager(ref eager) => {
246 eager.map.get(version).map(PrioritizedDist::hashes)
247 }
248 VersionMapInner::Lazy(ref lazy) => lazy.get(version).map(PrioritizedDist::hashes),
249 }
250 }
251
252 pub(crate) fn len(&self) -> usize {
257 match self.inner {
258 VersionMapInner::Eager(VersionMapEager { ref map, .. }) => map.len(),
259 VersionMapInner::Lazy(VersionMapLazy { ref map, .. }) => map.len(),
260 }
261 }
262
263 pub(crate) fn stable(&self) -> bool {
265 match self.inner {
266 VersionMapInner::Eager(ref map) => map.stable,
267 VersionMapInner::Lazy(ref map) => map.stable,
268 }
269 }
270
271 pub(crate) fn local(&self) -> bool {
273 match self.inner {
274 VersionMapInner::Eager(ref map) => map.local,
275 VersionMapInner::Lazy(ref map) => map.local,
276 }
277 }
278}
279
280impl From<FlatDistributions> for VersionMap {
281 fn from(flat_index: FlatDistributions) -> Self {
282 let stable = flat_index.iter().any(|(version, _)| version.is_stable());
283 let local = flat_index.iter().any(|(version, _)| version.is_local());
284 let map = flat_index.into();
285 Self {
286 inner: VersionMapInner::Eager(VersionMapEager { map, stable, local }),
287 }
288 }
289}
290
291pub(crate) struct VersionMapDistHandle<'a> {
304 inner: VersionMapDistHandleInner<'a>,
305}
306
307enum VersionMapDistHandleInner<'a> {
308 Eager(&'a PrioritizedDist),
309 Lazy {
310 lazy: &'a VersionMapLazy,
311 dist: &'a LazyPrioritizedDist,
312 },
313}
314
315impl<'a> VersionMapDistHandle<'a> {
316 pub(crate) fn prioritized_dist(&self) -> Option<&'a PrioritizedDist> {
318 match self.inner {
319 VersionMapDistHandleInner::Eager(dist) => Some(dist),
320 VersionMapDistHandleInner::Lazy { lazy, dist } => Some(lazy.get_lazy(dist)?),
321 }
322 }
323}
324
325#[derive(Debug)]
327#[expect(clippy::large_enum_variant)]
328enum VersionMapInner {
329 Eager(VersionMapEager),
334 Lazy(VersionMapLazy),
340}
341
342#[derive(Debug)]
344struct VersionMapEager {
345 map: BTreeMap<Version, PrioritizedDist>,
347 stable: bool,
349 local: bool,
351}
352
353#[derive(Debug)]
355struct VersionMapLazyEntry {
356 version: Version,
357 dist: LazyPrioritizedDist,
358}
359
360#[derive(Debug)]
362struct VersionMapLazyIndex {
363 entries: Vec<VersionMapLazyEntry>,
364}
365
366impl VersionMapLazyIndex {
367 fn merge_flat(self, flat_index: FlatDistributions) -> Self {
369 let flat_count = flat_index.iter().size_hint().0;
370 let mut merged = Vec::with_capacity(self.entries.len() + flat_count);
371 let mut simple = self.entries.into_iter().peekable();
372 let mut flat = flat_index.into_iter().peekable();
373 let flat_entry = |(version, dist)| VersionMapLazyEntry {
374 version,
375 dist: LazyPrioritizedDist {
376 flat: Some(dist),
377 simple: None,
378 },
379 };
380 while let (Some(simple_entry), Some((flat_version, _))) = (simple.peek(), flat.peek()) {
381 match simple_entry.version.cmp(flat_version) {
382 std::cmp::Ordering::Less => {
383 if let Some(entry) = simple.next() {
384 merged.push(entry);
385 }
386 }
387 std::cmp::Ordering::Greater => {
388 if let Some(entry) = flat.next() {
389 merged.push(flat_entry(entry));
390 }
391 }
392 std::cmp::Ordering::Equal => {
393 if let (Some(mut entry), Some((_, dist))) = (simple.next(), flat.next()) {
394 entry.dist.flat = Some(dist);
395 merged.push(entry);
396 }
397 }
398 }
399 }
400 merged.extend(simple);
401 merged.extend(flat.map(flat_entry));
402
403 Self { entries: merged }
404 }
405
406 fn get(&self, version: &Version) -> Option<&VersionMapLazyEntry> {
407 let index = self
408 .entries
409 .binary_search_by(|entry| entry.version.cmp(version))
410 .ok()?;
411 self.entries.get(index)
412 }
413
414 fn keys(&self) -> impl DoubleEndedIterator<Item = &Version> {
415 self.entries.iter().map(|entry| &entry.version)
416 }
417
418 fn range(&self, range: BoundingRange<'_>) -> &[VersionMapLazyEntry] {
419 let start = match range.min {
420 Bound::Included(version) => self
421 .entries
422 .partition_point(|entry| entry.version < *version),
423 Bound::Excluded(version) => self
424 .entries
425 .partition_point(|entry| entry.version <= *version),
426 Bound::Unbounded => 0,
427 };
428 let end = match range.max {
429 Bound::Included(version) => self
430 .entries
431 .partition_point(|entry| entry.version <= *version),
432 Bound::Excluded(version) => self
433 .entries
434 .partition_point(|entry| entry.version < *version),
435 Bound::Unbounded => self.entries.len(),
436 };
437 self.entries.get(start..end).unwrap_or_default()
438 }
439
440 fn len(&self) -> usize {
441 self.entries.len()
442 }
443}
444
445#[derive(Debug)]
454struct VersionMapLazy {
455 map: VersionMapLazyIndex,
457 stable: bool,
459 local: bool,
461 simple_metadata: OwnedArchive<SimpleDetailMetadata>,
464 no_binary: bool,
466 no_build: bool,
468 index: IndexUrl,
470 tags: Option<Tags>,
473 included_version_cutoff: Option<Timestamp>,
476 available_version_cutoff: Option<Timestamp>,
478 allowed_yanks: AllowedYanks,
480 hasher: HashStrategy,
482 requires_python: RequiresPython,
484}
485
486impl VersionMapLazy {
487 fn get_metadata(&self, version: &Version) -> Option<ResolutionMetadata> {
489 let archived = self
490 .map
491 .get(version)
492 .and_then(|entry| entry.dist.simple.as_ref())
493 .and_then(|simple| self.simple_metadata.datum(simple.datum_index))
494 .and_then(|datum| datum.metadata.as_ref())?;
495 Some(
496 rkyv::deserialize::<ResolutionMetadata, rkyv::rancor::Error>(archived)
497 .expect("archived metadata always deserializes"),
498 )
499 }
500
501 fn get(&self, version: &Version) -> Option<&PrioritizedDist> {
503 self.get_lazy(&self.map.get(version)?.dist)
504 }
505
506 fn get_lazy<'p>(&'p self, lazy_dist: &'p LazyPrioritizedDist) -> Option<&'p PrioritizedDist> {
512 match (&lazy_dist.flat, &lazy_dist.simple) {
513 (Some(flat), Some(simple)) => self.get_simple(Some(flat), simple),
514 (Some(flat), None) => Some(flat),
515 (None, Some(simple)) => self.get_simple(None, simple),
516 (None, None) => None,
517 }
518 }
519
520 fn get_simple<'p>(
525 &'p self,
526 init: Option<&'p PrioritizedDist>,
527 simple: &'p SimplePrioritizedDist,
528 ) -> Option<&'p PrioritizedDist> {
529 let get_or_init = || {
530 let files = rkyv::deserialize::<VersionFiles, rkyv::rancor::Error>(
531 &self
532 .simple_metadata
533 .datum(simple.datum_index)
534 .expect("index to lazy dist is correct")
535 .files,
536 )
537 .expect("archived version files always deserializes");
538 let mut priority_dist = init.cloned().unwrap_or_default();
539 for (filename, file) in files.all() {
540 let (excluded, upload_time) = if let Some(included_version_cutoff) =
543 &self.included_version_cutoff
544 {
545 match file.upload_time_utc_ms.as_ref() {
546 Some(&upload_time)
547 if upload_time >= included_version_cutoff.as_millisecond() =>
548 {
549 trace!(
550 "Excluding `{}` (uploaded {upload_time}) due to exclude-newer ({included_version_cutoff})",
551 file.filename
552 );
553 (true, Some(upload_time))
554 }
555 None => {
556 warn_user_once!(
557 "{} is missing an upload date, but user provided: {included_version_cutoff}",
558 file.filename,
559 );
560 (true, None)
561 }
562 _ => (false, None),
563 }
564 } else if let Some(available_version_cutoff) = &self.available_version_cutoff {
565 match file.upload_time_utc_ms.as_ref() {
566 Some(&upload_time)
567 if upload_time >= available_version_cutoff.as_millisecond() =>
568 {
569 trace!(
570 "Excluding `{}` (uploaded {upload_time}) due to available version cutoff ({available_version_cutoff})",
571 file.filename
572 );
573 (true, Some(upload_time))
574 }
575 _ => (false, None),
576 }
577 } else {
578 (false, None)
579 };
580
581 let yanked = file.yanked.as_deref();
583 let hashes = file.hashes.clone();
584 match filename {
585 DistFilename::WheelFilename(filename) => {
586 let compatibility = self.wheel_compatibility(
587 &filename,
588 &filename.name,
589 &filename.version,
590 hashes.as_slice(),
591 yanked,
592 excluded,
593 upload_time,
594 );
595 let dist = RegistryBuiltWheel {
596 filename,
597 file: Box::new(file),
598 index: self.index.clone(),
599 };
600 priority_dist.insert_built(dist, hashes, compatibility);
601 }
602 DistFilename::SourceDistFilename(filename) => {
603 let compatibility = self.source_dist_compatibility(
604 &filename.name,
605 &filename.version,
606 hashes.as_slice(),
607 yanked,
608 excluded,
609 upload_time,
610 );
611 let dist = RegistrySourceDist {
612 name: filename.name.clone(),
613 version: filename.version.clone(),
614 ext: filename.extension,
615 file: Box::new(file),
616 index: self.index.clone(),
617 wheels: vec![],
618 };
619 priority_dist.insert_source(dist, hashes, compatibility);
620 }
621 }
622 }
623 if priority_dist.is_empty() {
624 None
625 } else {
626 Some(priority_dist)
627 }
628 };
629 simple.dist.get_or_init(get_or_init).as_ref()
630 }
631
632 fn source_dist_compatibility(
633 &self,
634 name: &PackageName,
635 version: &Version,
636 hashes: &[HashDigest],
637 yanked: Option<&Yanked>,
638 excluded: bool,
639 upload_time: Option<i64>,
640 ) -> SourceDistCompatibility {
641 if self.no_build {
643 return SourceDistCompatibility::Incompatible(IncompatibleSource::NoBuild);
644 }
645
646 if excluded {
648 return SourceDistCompatibility::Incompatible(IncompatibleSource::ExcludeNewer(
649 upload_time,
650 ));
651 }
652
653 if let Some(yanked) = yanked {
655 if yanked.is_yanked() && !self.allowed_yanks.contains(name, version) {
656 return SourceDistCompatibility::Incompatible(IncompatibleSource::Yanked(
657 yanked.clone(),
658 ));
659 }
660 }
661
662 let hash_policy = self.hasher.get_package(name, version);
664 let required_hashes = hash_policy.digests();
665 let hash = if required_hashes.is_empty() {
666 HashComparison::Matched
667 } else {
668 if hashes.is_empty() {
669 HashComparison::Missing
670 } else if hash_policy.matches(hashes) {
671 HashComparison::Matched
672 } else {
673 HashComparison::Mismatched
674 }
675 };
676
677 SourceDistCompatibility::Compatible(hash)
678 }
679
680 fn wheel_compatibility(
681 &self,
682 filename: &WheelFilename,
683 name: &PackageName,
684 version: &Version,
685 hashes: &[HashDigest],
686 yanked: Option<&Yanked>,
687 excluded: bool,
688 upload_time: Option<i64>,
689 ) -> WheelCompatibility {
690 if self.no_binary {
692 return WheelCompatibility::Incompatible(IncompatibleWheel::NoBinary);
693 }
694
695 if excluded {
697 return WheelCompatibility::Incompatible(IncompatibleWheel::ExcludeNewer(upload_time));
698 }
699
700 if let Some(yanked) = yanked {
702 if yanked.is_yanked() && !self.allowed_yanks.contains(name, version) {
703 return WheelCompatibility::Incompatible(IncompatibleWheel::Yanked(yanked.clone()));
704 }
705 }
706
707 let priority = if let Some(tags) = &self.tags {
709 match filename.compatibility(tags) {
710 TagCompatibility::Incompatible(tag) => {
711 return WheelCompatibility::Incompatible(IncompatibleWheel::Tag(tag));
712 }
713 TagCompatibility::Compatible(priority) => Some(priority),
714 }
715 } else {
716 if !self.requires_python.matches_wheel_tag(filename) {
719 return WheelCompatibility::Incompatible(IncompatibleWheel::Tag(
720 IncompatibleTag::AbiPythonVersion,
721 ));
722 }
723 None
724 };
725
726 let hash_policy = self.hasher.get_package(name, version);
728 let required_hashes = hash_policy.digests();
729 let hash = if required_hashes.is_empty() {
730 HashComparison::Matched
731 } else {
732 if hashes.is_empty() {
733 HashComparison::Missing
734 } else if hash_policy.matches(hashes) {
735 HashComparison::Matched
736 } else {
737 HashComparison::Mismatched
738 }
739 };
740
741 let build_tag = filename.build_tag().cloned();
743
744 WheelCompatibility::Compatible(hash, priority, build_tag)
745 }
746}
747
748#[derive(Debug)]
750struct LazyPrioritizedDist {
751 flat: Option<PrioritizedDist>,
753 simple: Option<SimplePrioritizedDist>,
755}
756
757#[derive(Debug)]
759struct SimplePrioritizedDist {
760 datum_index: usize,
764 dist: OnceLock<Option<PrioritizedDist>>,
772}
773
774#[derive(Debug)]
776struct BoundingRange<'a> {
777 min: Bound<&'a Version>,
778 max: Bound<&'a Version>,
779}
780
781impl<'a> From<&'a Ranges<Version>> for BoundingRange<'a> {
782 fn from(value: &'a Ranges<Version>) -> Self {
783 let (min, max) = value
784 .bounding_range()
785 .unwrap_or((Bound::Unbounded, Bound::Unbounded));
786 Self { min, max }
787 }
788}
789
790impl<'a> RangeBounds<Version> for BoundingRange<'a> {
791 fn start_bound(&self) -> Bound<&'a Version> {
792 self.min
793 }
794
795 fn end_bound(&self) -> Bound<&'a Version> {
796 self.max
797 }
798}