Set fork solution as preference when resolving (#4662)

## Summary

This should both make it faster to solve forks (since we have a guess
for a valid resolution, and will bias towards packages we've already
fetched) and improve consistency between forks.

Closes https://github.com/astral-sh/uv/issues/4617.
This commit is contained in:
Charlie Marsh 2024-07-01 08:25:40 -04:00 committed by GitHub
parent bfadadefaf
commit 2d57309b0f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 58 additions and 53 deletions

View file

@ -1,6 +1,7 @@
//! Given a set of requirements, find a set of compatible packages.
use std::borrow::Cow;
use std::collections::hash_map::Entry;
use std::collections::{BTreeMap, VecDeque};
use std::fmt::{Display, Formatter};
use std::ops::Bound;
@ -323,6 +324,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
priorities: PubGrubPriorities::default(),
added_dependencies: FxHashMap::default(),
markers: MarkerTree::And(vec![]),
preferences: self.preferences.clone(),
};
let mut forked_states = vec![state];
let mut resolutions = vec![];
@ -373,7 +375,23 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
state.markers,
start.elapsed().as_secs_f32()
);
resolutions.push(state.into_resolution());
let resolution = state.into_resolution();
// Walk over the selected versions, and mark them as preferences.
for state in &mut forked_states {
for (package, versions) in &resolution.packages {
if let Entry::Vacant(entry) =
state.preferences.entry(package.name.clone())
{
if let Some(version) = versions.iter().next() {
entry.insert(version.clone().into());
}
}
}
}
resolutions.push(resolution);
continue 'FORK;
};
state.next = highest_priority_pkg;
@ -410,6 +428,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
&state.next,
term_intersection.unwrap_positive(),
&mut state.pins,
&state.preferences,
&state.fork_urls,
visited,
&request_sink,
@ -711,6 +730,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
package: &PubGrubPackage,
range: &Range<Version>,
pins: &mut FilePins,
fork_preferences: &Preferences,
fork_urls: &ForkUrls,
visited: &mut FxHashSet<PackageName>,
request_sink: &Sender<Request>,
@ -734,7 +754,15 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
if let Some(url) = package.name().and_then(|name| fork_urls.get(name)) {
self.choose_version_url(name, range, url)
} else {
self.choose_version_registry(name, range, package, pins, visited, request_sink)
self.choose_version_registry(
name,
range,
package,
fork_preferences,
pins,
visited,
request_sink,
)
}
}
}
@ -841,6 +869,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
name: &PackageName,
range: &Range<Version>,
package: &PubGrubPackage,
fork_preferences: &Preferences,
pins: &mut FilePins,
visited: &mut FxHashSet<PackageName>,
request_sink: &Sender<Request>,
@ -879,7 +908,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
name,
range,
version_maps,
&self.preferences,
fork_preferences,
&self.installed_packages,
&self.exclusions,
) else {
@ -1699,6 +1728,8 @@ struct SolveState {
/// that the marker expression that provoked the fork is true), then that
/// dependency is completely ignored.
markers: MarkerTree,
/// The preferences to respect for the fork.
preferences: Preferences,
}
impl SolveState {