Incorporate heuristics to improve package prioritization #3087
+163
−49
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, PubGrubPriority>); | ||
|
|
||
| impl PubGrubPriorities { | ||
| /// 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 [`PubGrubPriority`] of the given package, if it exists. | ||
| pub(crate) fn get(&self, package: &PubGrubPackage) -> Option<PubGrubPriority> { | ||
| match package { | ||
| PubGrubPackage::Root(_) => Some(PubGrubPriority::Root), | ||
| PubGrubPackage::Python(_) => Some(PubGrubPriority::Root), | ||
| PubGrubPackage::Package(name, _, _) => self.0.get(name).copied(), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| #[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
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
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This needs some comment about the heuristic, a reference to pip and that a lower entry means we pick it first