uv-resolver: implement basic resolver forking

There are still some TODOs/FIXMEs here, but this makes represents a
chunk of the resolver refactoring to enable forking. We don't do any
merging of resolutions yet, so crucially, this code is broken when no
marker environment is provided. But when a marker environment is
provided, this should behave the same as a non-forking resolver. In
particular, `get_dependencies_forking` is just `get_dependencies`
whenever there's a marker environment.
This commit is contained in:
Andrew Gallant 2024-05-22 14:25:12 -04:00 committed by Andrew Gallant
parent f5f330627b
commit 6f76a66510

View file

@ -297,19 +297,22 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
) -> Result<ResolutionGraph, ResolveError> { ) -> Result<ResolutionGraph, ResolveError> {
let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone())); let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone()));
let mut prefetcher = BatchPrefetcher::default(); let mut prefetcher = BatchPrefetcher::default();
let mut state = SolveState { let state = SolveState {
pubgrub: State::init(root.clone(), MIN_VERSION.clone()), pubgrub: State::init(root.clone(), MIN_VERSION.clone()),
next: root, next: root,
pins: FilePins::default(), pins: FilePins::default(),
priorities: PubGrubPriorities::default(), priorities: PubGrubPriorities::default(),
added_dependencies: FxHashMap::default(), added_dependencies: FxHashMap::default(),
}; };
let mut forked_states = vec![state];
let mut resolutions = vec![];
debug!( debug!(
"Solving with target Python version {}", "Solving with target Python version {}",
self.python_requirement.target() self.python_requirement.target()
); );
'FORK: while let Some(mut state) = forked_states.pop() {
loop { loop {
// Run unit propagation. // Run unit propagation.
state.pubgrub.unit_propagation(state.next)?; state.pubgrub.unit_propagation(state.next)?;
@ -333,14 +336,15 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
prefetcher.log_tried_versions(); prefetcher.log_tried_versions();
} }
let selection = state.pubgrub.partial_solution.extract_solution(); let selection = state.pubgrub.partial_solution.extract_solution();
return ResolutionGraph::from_state( resolutions.push(ResolutionGraph::from_state(
&selection, &selection,
&state.pins, &state.pins,
self.index.packages(), self.index.packages(),
self.index.distributions(), self.index.distributions(),
&state.pubgrub, &state.pubgrub,
&self.preferences, &self.preferences,
); )?);
continue 'FORK;
}; };
state.next = highest_priority_pkg; state.next = highest_priority_pkg;
@ -351,7 +355,9 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
.partial_solution .partial_solution
.term_intersection_for_package(&state.next) .term_intersection_for_package(&state.next)
.ok_or_else(|| { .ok_or_else(|| {
PubGrubError::Failure("a package was chosen but we don't have a term.".into()) PubGrubError::Failure(
"a package was chosen but we don't have a term.".into(),
)
})?; })?;
let decision = self.choose_version( let decision = self.choose_version(
&state.next, &state.next,
@ -406,7 +412,9 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
IncompatibleDist::Source(IncompatibleSource::RequiresPython( IncompatibleDist::Source(IncompatibleSource::RequiresPython(
requires_python, requires_python,
)) ))
| IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(requires_python)), | IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
requires_python,
)),
) = reason ) = reason
{ {
let python_version = requires_python let python_version = requires_python
@ -418,16 +426,16 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
let package = &state.next; let package = &state.next;
for kind in [PubGrubPython::Installed, PubGrubPython::Target] { for kind in [PubGrubPython::Installed, PubGrubPython::Target] {
state state.pubgrub.add_incompatibility(
.pubgrub Incompatibility::from_dependency(
.add_incompatibility(Incompatibility::from_dependency(
package.clone(), package.clone(),
Range::singleton(version.clone()), Range::singleton(version.clone()),
( (
PubGrubPackage::from(PubGrubPackageInner::Python(kind)), PubGrubPackage::from(PubGrubPackageInner::Python(kind)),
python_version.clone(), python_version.clone(),
), ),
)); ),
);
} }
state state
.pubgrub .pubgrub
@ -465,12 +473,15 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
{ {
// Retrieve that package dependencies. // Retrieve that package dependencies.
let package = &state.next; let package = &state.next;
let dependencies = match self.get_dependencies( let forks = self.get_dependencies_forking(
package, package,
&version, &version,
&mut state.priorities, &mut state.priorities,
&request_sink, &request_sink,
)? { )?;
for fork in forks {
let mut state = state.clone();
let dependencies = match fork {
Dependencies::Unavailable(reason) => { Dependencies::Unavailable(reason) => {
state state
.pubgrub .pubgrub
@ -479,6 +490,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
version.clone(), version.clone(),
UnavailableReason::Version(reason), UnavailableReason::Version(reason),
)); ));
forked_states.push(state);
continue; continue;
} }
Dependencies::Available(constraints) Dependencies::Available(constraints)
@ -507,11 +519,14 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
state.pubgrub.partial_solution.add_version( state.pubgrub.partial_solution.add_version(
package.clone(), package.clone(),
version, version.clone(),
dep_incompats, dep_incompats,
&state.pubgrub.incompatibility_store, &state.pubgrub.incompatibility_store,
); );
} else { forked_states.push(state);
}
continue 'FORK;
}
// `dep_incompats` are already in `incompatibilities` so we know there are not satisfied // `dep_incompats` are already in `incompatibilities` so we know there are not satisfied
// terms and can add the decision directly. // terms and can add the decision directly.
state state
@ -520,6 +535,11 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
.add_decision(state.next.clone(), version); .add_decision(state.next.clone(), version);
} }
} }
// This unwrap is okay because every code path above leads to at least
// one resolution being pushed.
//
// TODO: Implement merging of resolutions.
Ok(resolutions.pop().unwrap())
} }
/// Visit a [`PubGrubPackage`] prior to selection. This should be called on a [`PubGrubPackage`] /// Visit a [`PubGrubPackage`] prior to selection. This should be called on a [`PubGrubPackage`]
@ -796,6 +816,67 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
} }
} }
/// Given a candidate package and version, return its dependencies.
#[instrument(skip_all, fields(%package, %version))]
fn get_dependencies_forking(
&self,
package: &PubGrubPackage,
version: &Version,
priorities: &mut PubGrubPriorities,
request_sink: &Sender<Request>,
) -> Result<Vec<Dependencies>, ResolveError> {
type Dep = (PubGrubPackage, Range<Version>);
let result = self.get_dependencies(package, version, priorities, request_sink);
if self.markers.is_some() {
return result.map(|deps| vec![deps]);
}
let deps: Vec<Dep> = match result? {
Dependencies::Available(deps) => deps,
Dependencies::Unavailable(err) => return Ok(vec![Dependencies::Unavailable(err)]),
};
let mut by_grouping: FxHashMap<&PackageName, FxHashMap<&Range<Version>, Vec<&Dep>>> =
FxHashMap::default();
for dep in &deps {
let (ref pkg, ref range) = *dep;
let name = match &**pkg {
// A root can never be a dependency of another package, and a `Python` pubgrub
// package is never returned by `get_dependencies`. So these cases never occur.
PubGrubPackageInner::Root(_) | PubGrubPackageInner::Python(_) => unreachable!(),
PubGrubPackageInner::Package { ref name, .. }
| PubGrubPackageInner::Extra { ref name, .. } => name,
};
by_grouping
.entry(name)
.or_default()
.entry(range)
.or_default()
.push(dep);
}
let mut forks: Vec<Vec<Dep>> = vec![vec![]];
for (_, groups) in by_grouping {
if groups.len() <= 1 {
for deps in groups.into_values() {
for fork in &mut forks {
fork.extend(deps.iter().map(|dep| (*dep).clone()));
}
}
} else {
let mut new_forks: Vec<Vec<Dep>> = vec![];
for deps in groups.into_values() {
let mut new_forks_for_group = forks.clone();
for fork in &mut new_forks_for_group {
fork.extend(deps.iter().map(|dep| (*dep).clone()));
}
new_forks.extend(new_forks_for_group);
}
forks = new_forks;
}
}
Ok(forks.into_iter().map(Dependencies::Available).collect())
}
/// Given a candidate package and version, return its dependencies. /// Given a candidate package and version, return its dependencies.
#[instrument(skip_all, fields(%package, %version))] #[instrument(skip_all, fields(%package, %version))]
fn get_dependencies( fn get_dependencies(