-
Notifications
You must be signed in to change notification settings - Fork 765
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Incorporate heuristics to improve package prioritization
- Loading branch information
1 parent
d1b07a3
commit f326b43
Showing
3 changed files
with
124 additions
and
32 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,34 +1,126 @@ | ||
| use std::cmp::Reverse; | ||
|
|
||
| use pubgrub::range::Range; | ||
| use rustc_hash::FxHashMap; | ||
|
|
||
| use pep440_rs::Version; | ||
| use uv_normalize::PackageName; | ||
|
|
||
| use crate::pubgrub::package::PubGrubPackage; | ||
|
|
||
| /// A prioritization map to guide the PubGrub resolution process. | ||
| /// | ||
| /// During resolution, PubGrub needs to decide which package to consider next. The priorities | ||
| /// encoded here are used to guide that decision. | ||
| /// | ||
| /// Like `pip`, we prefer packages that are pinned to direct URLs over packages pinned to a single | ||
| /// version over packages that are constrained in some way over packages that are unconstrained. | ||
| /// | ||
| /// See: <https://github.com/pypa/pip/blob/ef78c129b1a966dbbbdb8ebfffc43723e89110d1/src/pip/_internal/resolution/resolvelib/provider.py#L120> | ||
| #[derive(Debug, Default)] | ||
| pub(crate) struct PubGrubPriorities(FxHashMap<PackageName, usize>); | ||
| pub(crate) struct PubGrubPriorities(FxHashMap<PackageName, PubGrubPriority>); | ||
|
|
||
| impl PubGrubPriorities { | ||
| /// Add a package to the priority map. | ||
| pub(crate) fn add(&mut self, package: PackageName) { | ||
| let priority = self.0.len(); | ||
| self.0.entry(package).or_insert(priority); | ||
| /// Add a [`PubGrubPackage`] to the priority map. | ||
| pub(crate) fn insert(&mut self, package: &PubGrubPackage, version: &Range<Version>) { | ||
| let next = self.0.len(); | ||
| match package { | ||
| PubGrubPackage::Root(_) => {} | ||
| PubGrubPackage::Python(_) => {} | ||
| PubGrubPackage::Package(name, _, None) => { | ||
| match self.0.entry(name.clone()) { | ||
| std::collections::hash_map::Entry::Occupied(mut entry) => { | ||
| // Preserve the original index. | ||
| let index = match entry.get() { | ||
| PubGrubPriority::Singleton(Reverse(index)) => *index, | ||
| PubGrubPriority::Unconstrained(Reverse(index)) => *index, | ||
| PubGrubPriority::Constrained(Reverse(index)) => *index, | ||
| PubGrubPriority::DirectUrl(Reverse(index)) => *index, | ||
| PubGrubPriority::Root => next, | ||
| }; | ||
|
|
||
| // Compute the priority. | ||
| let priority = if version.as_singleton().is_some() { | ||
| PubGrubPriority::Singleton(Reverse(index)) | ||
| } else if version == &Range::full() { | ||
| PubGrubPriority::Unconstrained(Reverse(index)) | ||
| } else { | ||
| PubGrubPriority::Constrained(Reverse(index)) | ||
| }; | ||
|
|
||
| // Take the maximum of the new and existing priorities. | ||
| if priority > *entry.get() { | ||
| entry.insert(priority); | ||
| } | ||
| } | ||
| std::collections::hash_map::Entry::Vacant(entry) => { | ||
| // Insert the priority. | ||
| entry.insert(if version.as_singleton().is_some() { | ||
| PubGrubPriority::Singleton(Reverse(next)) | ||
| } else if version == &Range::full() { | ||
| PubGrubPriority::Unconstrained(Reverse(next)) | ||
| } else { | ||
| PubGrubPriority::Constrained(Reverse(next)) | ||
| }); | ||
| } | ||
| } | ||
| } | ||
| PubGrubPackage::Package(name, _, Some(_)) => { | ||
| match self.0.entry(name.clone()) { | ||
| std::collections::hash_map::Entry::Occupied(mut entry) => { | ||
| // Preserve the original index. | ||
| let index = match entry.get() { | ||
| PubGrubPriority::Singleton(Reverse(index)) => *index, | ||
| PubGrubPriority::Unconstrained(Reverse(index)) => *index, | ||
| PubGrubPriority::Constrained(Reverse(index)) => *index, | ||
| PubGrubPriority::DirectUrl(Reverse(index)) => *index, | ||
| PubGrubPriority::Root => next, | ||
| }; | ||
|
|
||
| // Compute the priority. | ||
| let priority = PubGrubPriority::DirectUrl(Reverse(index)); | ||
|
|
||
| // Take the maximum of the new and existing priorities. | ||
| if priority > *entry.get() { | ||
| entry.insert(priority); | ||
| } | ||
| } | ||
| std::collections::hash_map::Entry::Vacant(entry) => { | ||
| // Insert the priority. | ||
| entry.insert(PubGrubPriority::DirectUrl(Reverse(next))); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// Return the priority of the given package, if it exists. | ||
| /// Return the [`PubGrubPriority`] of the given package, if it exists. | ||
| pub(crate) fn get(&self, package: &PubGrubPackage) -> Option<PubGrubPriority> { | ||
| match package { | ||
| PubGrubPackage::Root(_) => Some(Reverse(0)), | ||
| PubGrubPackage::Python(_) => Some(Reverse(0)), | ||
| PubGrubPackage::Package(name, _, _) => self | ||
| .0 | ||
| .get(name) | ||
| .copied() | ||
| .map(|priority| priority + 1) | ||
| .map(Reverse), | ||
| PubGrubPackage::Root(_) => Some(PubGrubPriority::Root), | ||
| PubGrubPackage::Python(_) => Some(PubGrubPriority::Root), | ||
| PubGrubPackage::Package(name, _, _) => self.0.get(name).copied(), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pub(crate) type PubGrubPriority = Reverse<usize>; | ||
| #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] | ||
| pub(crate) enum PubGrubPriority { | ||
| /// The package has no specific priority. | ||
| /// | ||
| /// As such, its priority is based on the order in which the packages were added (FIFO), such | ||
| /// that the first package we visit is prioritized over subsequent packages. | ||
| Unconstrained(Reverse<usize>), | ||
|
|
||
| /// The version range is constrained in some way (e.g., with a `<=` or `>` operator). | ||
| Constrained(Reverse<usize>), | ||
|
|
||
| /// The version range is constrained to a single version (e.g., with the `==` operator). | ||
| Singleton(Reverse<usize>), | ||
|
|
||
| /// The package was specified via a direct URL. | ||
| DirectUrl(Reverse<usize>), | ||
|
|
||
| /// The package is the root package. | ||
| Root, | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters