1use std::collections::Bound;
2use std::collections::btree_map::{BTreeMap, Entry};
3use std::ops::RangeBounds;
4use std::sync::OnceLock;
5
6use jiff::Timestamp;
7use pubgrub::Ranges;
8use rustc_hash::FxHashMap;
9use tracing::{instrument, trace};
10
11use uv_client::{FlatIndexEntry, OwnedArchive, SimpleDetailMetadata, VersionFiles};
12use uv_configuration::BuildOptions;
13use uv_distribution_filename::{DistFilename, WheelFilename};
14use uv_distribution_types::{
15 HashComparison, IncompatibleSource, IncompatibleWheel, IndexUrl, PrioritizedDist,
16 RegistryBuiltWheel, RegistrySourceDist, RequiresPython, SourceDistCompatibility,
17 WheelCompatibility,
18};
19use uv_normalize::PackageName;
20use uv_pep440::Version;
21use uv_platform_tags::{IncompatibleTag, TagCompatibility, Tags};
22use uv_pypi_types::{HashDigest, ResolutionMetadata, Yanked};
23use uv_types::HashStrategy;
24use uv_warnings::warn_user_once;
25
26use crate::flat_index::FlatDistributions;
27use crate::yanks::AllowedYanks;
28
29#[derive(Debug)]
31pub struct VersionMap {
32 inner: VersionMapInner,
34}
35
36impl VersionMap {
37 #[instrument(skip_all, fields(package_name))]
47 pub(crate) fn from_simple_metadata(
48 simple_metadata: OwnedArchive<SimpleDetailMetadata>,
49 package_name: &PackageName,
50 index: &IndexUrl,
51 tags: Option<&Tags>,
52 requires_python: &RequiresPython,
53 allowed_yanks: &AllowedYanks,
54 hasher: &HashStrategy,
55 included_version_cutoff: Option<Timestamp>,
56 available_version_cutoff: Option<Timestamp>,
57 flat_index: Option<FlatDistributions>,
58 build_options: &BuildOptions,
59 ) -> Self {
60 let mut stable = false;
61 let mut local = false;
62 let mut map = BTreeMap::new();
63 let mut core_metadata = FxHashMap::default();
64 for (datum_index, datum) in simple_metadata.iter().enumerate() {
68 let version = rkyv::deserialize::<Version, rkyv::rancor::Error>(&datum.version)
70 .expect("archived version always deserializes");
71
72 let core_metadatum =
74 rkyv::deserialize::<Option<ResolutionMetadata>, rkyv::rancor::Error>(
75 &datum.metadata,
76 )
77 .expect("archived metadata always deserializes");
78 if let Some(core_metadatum) = core_metadatum {
79 core_metadata.insert(version.clone(), core_metadatum);
80 }
81
82 stable |= version.is_stable();
83 local |= version.is_local();
84 map.insert(
85 version,
86 LazyPrioritizedDist::OnlySimple(SimplePrioritizedDist {
87 datum_index,
88 dist: OnceLock::new(),
89 }),
90 );
91 }
92 for (version, prioritized_dist) in flat_index.into_iter().flatten() {
95 stable |= version.is_stable();
96 match map.entry(version) {
97 Entry::Vacant(e) => {
98 e.insert(LazyPrioritizedDist::OnlyFlat(prioritized_dist));
99 }
100 Entry::Occupied(e) => match e.remove_entry() {
105 (version, LazyPrioritizedDist::OnlySimple(simple_dist)) => {
106 map.insert(
107 version,
108 LazyPrioritizedDist::Both {
109 flat: prioritized_dist,
110 simple: simple_dist,
111 },
112 );
113 }
114 _ => unreachable!(),
115 },
116 }
117 }
118 Self {
119 inner: VersionMapInner::Lazy(VersionMapLazy {
120 map,
121 stable,
122 local,
123 core_metadata,
124 simple_metadata,
125 no_binary: build_options.no_binary_package(package_name),
126 no_build: build_options.no_build_package(package_name),
127 index: index.clone(),
128 tags: tags.cloned(),
129 allowed_yanks: allowed_yanks.clone(),
130 hasher: hasher.clone(),
131 requires_python: requires_python.clone(),
132 included_version_cutoff,
133 available_version_cutoff,
134 }),
135 }
136 }
137
138 #[instrument(skip_all, fields(package_name))]
139 pub(crate) fn from_flat_metadata(
140 flat_metadata: Vec<FlatIndexEntry>,
141 tags: Option<&Tags>,
142 hasher: &HashStrategy,
143 build_options: &BuildOptions,
144 ) -> Self {
145 let mut stable = false;
146 let mut local = false;
147 let mut map = BTreeMap::new();
148
149 for (version, prioritized_dist) in
150 FlatDistributions::from_entries(flat_metadata, tags, hasher, build_options)
151 {
152 stable |= version.is_stable();
153 local |= version.is_local();
154 map.insert(version, prioritized_dist);
155 }
156
157 Self {
158 inner: VersionMapInner::Eager(VersionMapEager { map, stable, local }),
159 }
160 }
161
162 pub fn get_metadata(&self, version: &Version) -> Option<&ResolutionMetadata> {
164 match self.inner {
165 VersionMapInner::Eager(_) => None,
166 VersionMapInner::Lazy(ref lazy) => lazy.core_metadata.get(version),
167 }
168 }
169
170 pub(crate) fn get(&self, version: &Version) -> Option<&PrioritizedDist> {
172 match self.inner {
173 VersionMapInner::Eager(ref eager) => eager.map.get(version),
174 VersionMapInner::Lazy(ref lazy) => lazy.get(version),
175 }
176 }
177
178 pub(crate) fn versions(&self) -> impl DoubleEndedIterator<Item = &Version> {
180 match &self.inner {
181 VersionMapInner::Eager(eager) => either::Either::Left(eager.map.keys()),
182 VersionMapInner::Lazy(lazy) => either::Either::Right(lazy.map.keys()),
183 }
184 }
185
186 pub(crate) fn index(&self) -> Option<&IndexUrl> {
188 match &self.inner {
189 VersionMapInner::Eager(_) => None,
190 VersionMapInner::Lazy(lazy) => Some(&lazy.index),
191 }
192 }
193
194 pub(crate) fn included_version_cutoff(&self) -> Option<&Timestamp> {
196 match &self.inner {
197 VersionMapInner::Eager(_) => None,
198 VersionMapInner::Lazy(lazy) => lazy.included_version_cutoff.as_ref(),
199 }
200 }
201
202 pub(crate) fn iter(
209 &self,
210 range: &Ranges<Version>,
211 ) -> impl DoubleEndedIterator<Item = (&Version, VersionMapDistHandle<'_>)> {
212 if let Some(version) = range.as_singleton() {
214 either::Either::Left(match self.inner {
215 VersionMapInner::Eager(ref eager) => {
216 either::Either::Left(eager.map.get_key_value(version).into_iter().map(
217 move |(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.get_key_value(version).into_iter().map(
227 move |(version, dist)| {
228 let version_map_dist = VersionMapDistHandle {
229 inner: VersionMapDistHandleInner::Lazy { lazy, dist },
230 };
231 (version, version_map_dist)
232 },
233 ))
234 }
235 })
236 } else {
237 either::Either::Right(match self.inner {
238 VersionMapInner::Eager(ref eager) => {
239 either::Either::Left(eager.map.range(BoundingRange::from(range)).map(
240 |(version, dist)| {
241 let version_map_dist = VersionMapDistHandle {
242 inner: VersionMapDistHandleInner::Eager(dist),
243 };
244 (version, version_map_dist)
245 },
246 ))
247 }
248 VersionMapInner::Lazy(ref lazy) => {
249 either::Either::Right(lazy.map.range(BoundingRange::from(range)).map(
250 |(version, dist)| {
251 let version_map_dist = VersionMapDistHandle {
252 inner: VersionMapDistHandleInner::Lazy { lazy, dist },
253 };
254 (version, version_map_dist)
255 },
256 ))
257 }
258 })
259 }
260 }
261
262 pub(crate) fn hashes(&self, version: &Version) -> Option<&[HashDigest]> {
264 match self.inner {
265 VersionMapInner::Eager(ref eager) => {
266 eager.map.get(version).map(PrioritizedDist::hashes)
267 }
268 VersionMapInner::Lazy(ref lazy) => lazy.get(version).map(PrioritizedDist::hashes),
269 }
270 }
271
272 pub(crate) fn len(&self) -> usize {
277 match self.inner {
278 VersionMapInner::Eager(VersionMapEager { ref map, .. }) => map.len(),
279 VersionMapInner::Lazy(VersionMapLazy { ref map, .. }) => map.len(),
280 }
281 }
282
283 pub(crate) fn stable(&self) -> bool {
285 match self.inner {
286 VersionMapInner::Eager(ref map) => map.stable,
287 VersionMapInner::Lazy(ref map) => map.stable,
288 }
289 }
290
291 pub(crate) fn local(&self) -> bool {
293 match self.inner {
294 VersionMapInner::Eager(ref map) => map.local,
295 VersionMapInner::Lazy(ref map) => map.local,
296 }
297 }
298}
299
300impl From<FlatDistributions> for VersionMap {
301 fn from(flat_index: FlatDistributions) -> Self {
302 let stable = flat_index.iter().any(|(version, _)| version.is_stable());
303 let local = flat_index.iter().any(|(version, _)| version.is_local());
304 let map = flat_index.into();
305 Self {
306 inner: VersionMapInner::Eager(VersionMapEager { map, stable, local }),
307 }
308 }
309}
310
311pub(crate) struct VersionMapDistHandle<'a> {
324 inner: VersionMapDistHandleInner<'a>,
325}
326
327enum VersionMapDistHandleInner<'a> {
328 Eager(&'a PrioritizedDist),
329 Lazy {
330 lazy: &'a VersionMapLazy,
331 dist: &'a LazyPrioritizedDist,
332 },
333}
334
335impl<'a> VersionMapDistHandle<'a> {
336 pub(crate) fn prioritized_dist(&self) -> Option<&'a PrioritizedDist> {
338 match self.inner {
339 VersionMapDistHandleInner::Eager(dist) => Some(dist),
340 VersionMapDistHandleInner::Lazy { lazy, dist } => Some(lazy.get_lazy(dist)?),
341 }
342 }
343}
344
345#[derive(Debug)]
347#[expect(clippy::large_enum_variant)]
348enum VersionMapInner {
349 Eager(VersionMapEager),
354 Lazy(VersionMapLazy),
360}
361
362#[derive(Debug)]
364struct VersionMapEager {
365 map: BTreeMap<Version, PrioritizedDist>,
367 stable: bool,
369 local: bool,
371}
372
373#[derive(Debug)]
382struct VersionMapLazy {
383 map: BTreeMap<Version, LazyPrioritizedDist>,
385 stable: bool,
387 local: bool,
389 core_metadata: FxHashMap<Version, ResolutionMetadata>,
391 simple_metadata: OwnedArchive<SimpleDetailMetadata>,
394 no_binary: bool,
396 no_build: bool,
398 index: IndexUrl,
400 tags: Option<Tags>,
403 included_version_cutoff: Option<Timestamp>,
406 available_version_cutoff: Option<Timestamp>,
408 allowed_yanks: AllowedYanks,
410 hasher: HashStrategy,
412 requires_python: RequiresPython,
414}
415
416impl VersionMapLazy {
417 fn get(&self, version: &Version) -> Option<&PrioritizedDist> {
419 let lazy_dist = self.map.get(version)?;
420 let priority_dist = self.get_lazy(lazy_dist)?;
421 Some(priority_dist)
422 }
423
424 fn get_lazy<'p>(&'p self, lazy_dist: &'p LazyPrioritizedDist) -> Option<&'p PrioritizedDist> {
430 match *lazy_dist {
431 LazyPrioritizedDist::OnlyFlat(ref dist) => Some(dist),
432 LazyPrioritizedDist::OnlySimple(ref dist) => self.get_simple(None, dist),
433 LazyPrioritizedDist::Both {
434 ref flat,
435 ref simple,
436 } => self.get_simple(Some(flat), simple),
437 }
438 }
439
440 fn get_simple<'p>(
445 &'p self,
446 init: Option<&'p PrioritizedDist>,
447 simple: &'p SimplePrioritizedDist,
448 ) -> Option<&'p PrioritizedDist> {
449 let get_or_init = || {
450 let files = rkyv::deserialize::<VersionFiles, rkyv::rancor::Error>(
451 &self
452 .simple_metadata
453 .datum(simple.datum_index)
454 .expect("index to lazy dist is correct")
455 .files,
456 )
457 .expect("archived version files always deserializes");
458 let mut priority_dist = init.cloned().unwrap_or_default();
459 for (filename, file) in files.all() {
460 let (excluded, upload_time) = if let Some(included_version_cutoff) =
463 &self.included_version_cutoff
464 {
465 match file.upload_time_utc_ms.as_ref() {
466 Some(&upload_time)
467 if upload_time >= included_version_cutoff.as_millisecond() =>
468 {
469 trace!(
470 "Excluding `{}` (uploaded {upload_time}) due to exclude-newer ({included_version_cutoff})",
471 file.filename
472 );
473 (true, Some(upload_time))
474 }
475 None => {
476 warn_user_once!(
477 "{} is missing an upload date, but user provided: {included_version_cutoff}",
478 file.filename,
479 );
480 (true, None)
481 }
482 _ => (false, None),
483 }
484 } else if let Some(available_version_cutoff) = &self.available_version_cutoff {
485 match file.upload_time_utc_ms.as_ref() {
486 Some(&upload_time)
487 if upload_time >= available_version_cutoff.as_millisecond() =>
488 {
489 trace!(
490 "Excluding `{}` (uploaded {upload_time}) due to available version cutoff ({available_version_cutoff})",
491 file.filename
492 );
493 (true, Some(upload_time))
494 }
495 _ => (false, None),
496 }
497 } else {
498 (false, None)
499 };
500
501 let yanked = file.yanked.as_deref();
503 let hashes = file.hashes.clone();
504 match filename {
505 DistFilename::WheelFilename(filename) => {
506 let compatibility = self.wheel_compatibility(
507 &filename,
508 &filename.name,
509 &filename.version,
510 hashes.as_slice(),
511 yanked,
512 excluded,
513 upload_time,
514 );
515 let dist = RegistryBuiltWheel {
516 filename,
517 file: Box::new(file),
518 index: self.index.clone(),
519 };
520 priority_dist.insert_built(dist, hashes, compatibility);
521 }
522 DistFilename::SourceDistFilename(filename) => {
523 let compatibility = self.source_dist_compatibility(
524 &filename.name,
525 &filename.version,
526 hashes.as_slice(),
527 yanked,
528 excluded,
529 upload_time,
530 );
531 let dist = RegistrySourceDist {
532 name: filename.name.clone(),
533 version: filename.version.clone(),
534 ext: filename.extension,
535 file: Box::new(file),
536 index: self.index.clone(),
537 wheels: vec![],
538 };
539 priority_dist.insert_source(dist, hashes, compatibility);
540 }
541 }
542 }
543 if priority_dist.is_empty() {
544 None
545 } else {
546 Some(priority_dist)
547 }
548 };
549 simple.dist.get_or_init(get_or_init).as_ref()
550 }
551
552 fn source_dist_compatibility(
553 &self,
554 name: &PackageName,
555 version: &Version,
556 hashes: &[HashDigest],
557 yanked: Option<&Yanked>,
558 excluded: bool,
559 upload_time: Option<i64>,
560 ) -> SourceDistCompatibility {
561 if self.no_build {
563 return SourceDistCompatibility::Incompatible(IncompatibleSource::NoBuild);
564 }
565
566 if excluded {
568 return SourceDistCompatibility::Incompatible(IncompatibleSource::ExcludeNewer(
569 upload_time,
570 ));
571 }
572
573 if let Some(yanked) = yanked {
575 if yanked.is_yanked() && !self.allowed_yanks.contains(name, version) {
576 return SourceDistCompatibility::Incompatible(IncompatibleSource::Yanked(
577 yanked.clone(),
578 ));
579 }
580 }
581
582 let hash_policy = self.hasher.get_package(name, version);
584 let required_hashes = hash_policy.digests();
585 let hash = if required_hashes.is_empty() {
586 HashComparison::Matched
587 } else {
588 if hashes.is_empty() {
589 HashComparison::Missing
590 } else if hash_policy.matches(hashes) {
591 HashComparison::Matched
592 } else {
593 HashComparison::Mismatched
594 }
595 };
596
597 SourceDistCompatibility::Compatible(hash)
598 }
599
600 fn wheel_compatibility(
601 &self,
602 filename: &WheelFilename,
603 name: &PackageName,
604 version: &Version,
605 hashes: &[HashDigest],
606 yanked: Option<&Yanked>,
607 excluded: bool,
608 upload_time: Option<i64>,
609 ) -> WheelCompatibility {
610 if self.no_binary {
612 return WheelCompatibility::Incompatible(IncompatibleWheel::NoBinary);
613 }
614
615 if excluded {
617 return WheelCompatibility::Incompatible(IncompatibleWheel::ExcludeNewer(upload_time));
618 }
619
620 if let Some(yanked) = yanked {
622 if yanked.is_yanked() && !self.allowed_yanks.contains(name, version) {
623 return WheelCompatibility::Incompatible(IncompatibleWheel::Yanked(yanked.clone()));
624 }
625 }
626
627 let priority = if let Some(tags) = &self.tags {
629 match filename.compatibility(tags) {
630 TagCompatibility::Incompatible(tag) => {
631 return WheelCompatibility::Incompatible(IncompatibleWheel::Tag(tag));
632 }
633 TagCompatibility::Compatible(priority) => Some(priority),
634 }
635 } else {
636 if !self.requires_python.matches_wheel_tag(filename) {
639 return WheelCompatibility::Incompatible(IncompatibleWheel::Tag(
640 IncompatibleTag::AbiPythonVersion,
641 ));
642 }
643 None
644 };
645
646 let hash_policy = self.hasher.get_package(name, version);
648 let required_hashes = hash_policy.digests();
649 let hash = if required_hashes.is_empty() {
650 HashComparison::Matched
651 } else {
652 if hashes.is_empty() {
653 HashComparison::Missing
654 } else if hash_policy.matches(hashes) {
655 HashComparison::Matched
656 } else {
657 HashComparison::Mismatched
658 }
659 };
660
661 let build_tag = filename.build_tag().cloned();
663
664 WheelCompatibility::Compatible(hash, priority, build_tag)
665 }
666}
667
668#[derive(Debug)]
671enum LazyPrioritizedDist {
672 OnlyFlat(PrioritizedDist),
675 OnlySimple(SimplePrioritizedDist),
678 Both {
681 flat: PrioritizedDist,
682 simple: SimplePrioritizedDist,
683 },
684}
685
686#[derive(Debug)]
688struct SimplePrioritizedDist {
689 datum_index: usize,
693 dist: OnceLock<Option<PrioritizedDist>>,
701}
702
703#[derive(Debug)]
705struct BoundingRange<'a> {
706 min: Bound<&'a Version>,
707 max: Bound<&'a Version>,
708}
709
710impl<'a> From<&'a Ranges<Version>> for BoundingRange<'a> {
711 fn from(value: &'a Ranges<Version>) -> Self {
712 let (min, max) = value
713 .bounding_range()
714 .unwrap_or((Bound::Unbounded, Bound::Unbounded));
715 Self { min, max }
716 }
717}
718
719impl<'a> RangeBounds<Version> for BoundingRange<'a> {
720 fn start_bound(&self) -> Bound<&'a Version> {
721 self.min
722 }
723
724 fn end_bound(&self) -> Bound<&'a Version> {
725 self.max
726 }
727}