From ff4651c6979b7cb63b4184f4a8be87180f8dfb17 Mon Sep 17 00:00:00 2001 From: konstin Date: Thu, 25 Jul 2024 13:40:04 +0200 Subject: [PATCH] Use fork marker preferences in resolutions --- crates/uv-resolver/src/candidate_selector.rs | 216 +++++++++++++------ crates/uv-resolver/src/lock.rs | 4 + crates/uv-resolver/src/preferences.rs | 141 +++++++++--- crates/uv-resolver/src/resolver/mod.rs | 13 +- crates/uv-resolver/src/yanks.rs | 4 +- 5 files changed, 274 insertions(+), 104 deletions(-) diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs index f42dc433d9f1..df03e7ba03f6 100644 --- a/crates/uv-resolver/src/candidate_selector.rs +++ b/crates/uv-resolver/src/candidate_selector.rs @@ -1,12 +1,12 @@ use itertools::Itertools; use pubgrub::range::Range; use std::fmt::{Display, Formatter}; -use tracing::debug; +use tracing::{debug, trace}; use distribution_types::{CompatibleDist, IncompatibleDist, IncompatibleSource}; use distribution_types::{DistributionMetadata, IncompatibleWheel, Name, PrioritizedDist}; use pep440_rs::Version; -use pep508_rs::MarkerEnvironment; +use pep508_rs::{MarkerEnvironment, MarkerTree}; use uv_configuration::IndexStrategy; use uv_normalize::PackageName; use uv_types::InstalledPackagesProvider; @@ -66,9 +66,7 @@ impl CandidateSelector { pub(crate) fn index_strategy(&self) -> &IndexStrategy { &self.index_strategy } -} -impl CandidateSelector { /// Select a [`Candidate`] from a set of candidate versions and files. /// /// Unless present in the provided [`Exclusions`], local distributions from the @@ -84,6 +82,8 @@ impl CandidateSelector { exclusions: &'a Exclusions, markers: &ResolverMarkers, ) -> Option> { + // Check for a preference from a lockfile or a previous fork that satisfies the range and + // is allowed. if let Some(preferred) = self.get_preferred( package_name, range, @@ -93,11 +93,18 @@ impl CandidateSelector { exclusions, markers, ) { + trace!("Using preference {} {}", preferred.name, preferred.version,); return Some(preferred); } + // Check for a locally installed distribution that satisfies the range and is allowed. if !exclusions.contains(package_name) { if let Some(installed) = Self::get_installed(package_name, range, installed_packages) { + trace!( + "Using preference {} {} from installed package", + installed.name, + installed.version, + ); return Some(installed); } } @@ -105,74 +112,160 @@ impl CandidateSelector { self.select_no_preference(package_name, range, version_maps, markers) } - /// Check for a preference (e.g., an existing version from an existing lockfile or - /// from a previous fork) that satisfies the current range. + /// If the package has a preference, an existing version from an existing lockfile or a version + /// from a sibling fork, and the preference satisfies the current range, use that. + /// + /// We try to find a resolution that, depending on the input, does not diverge from the + /// lockfile or matches a sibling fork. We try an exact match for the current markers (fork + /// or specific) first, to ensure stability with repeated locking. If that doesn't work, we + /// fall back to preferences that don't match in hopes of still resolving different forks into + /// the same version; A solution with less different versions is more desirable than one where + /// we may have more recent version in some cases, but overall more versions. fn get_preferred<'a, InstalledPackages: InstalledPackagesProvider>( - &self, + &'a self, package_name: &'a PackageName, range: &Range, version_maps: &'a [VersionMap], preferences: &'a Preferences, installed_packages: &'a InstalledPackages, - exclusions: &'a Exclusions, - markers: &ResolverMarkers, - ) -> Option> { - let version = preferences.version(package_name)?; - - // Respect the version range for this requirement. - if !range.contains(version) { - return None; + exclusions: &Exclusions, + resolver_markers: &ResolverMarkers, + ) -> Option { + // In the branches, we "sort" the preferences by marker-matching through an iterator that + // first has the matching half and then the mismatching half. + match resolver_markers { + ResolverMarkers::SpecificEnvironment(env) => { + // We may hit a combination of fork markers preferences with specific environment + // output in the future when adding support for the PEP 665 successor. + let preferences_match = + preferences.get(package_name).filter(|(marker, _version)| { + // `.unwrap_or(true)` because the universal marker is considered matching. + marker + .map(|marker| marker.evaluate(env, &[])) + .unwrap_or(true) + }); + let preferences_mismatch = + preferences.get(package_name).filter(|(marker, _version)| { + marker + .map(|marker| !marker.evaluate(env, &[])) + .unwrap_or(false) + }); + self.get_preferred_from_iter( + preferences_match.chain(preferences_mismatch), + package_name, + range, + version_maps, + installed_packages, + exclusions, + resolver_markers, + ) + } + ResolverMarkers::Universal { .. } => { + // In universal mode, all preferences are matching. + self.get_preferred_from_iter( + preferences.get(package_name), + package_name, + range, + version_maps, + installed_packages, + exclusions, + resolver_markers, + ) + } + ResolverMarkers::Fork(fork_markers) => { + let preferences_match = + preferences.get(package_name).filter(|(marker, _version)| { + // `.unwrap_or(true)` because the universal marker is considered matching. + marker.map(|marker| marker == fork_markers).unwrap_or(true) + }); + let preferences_mismatch = + preferences.get(package_name).filter(|(marker, _version)| { + marker.map(|marker| marker != fork_markers).unwrap_or(false) + }); + self.get_preferred_from_iter( + preferences_match.chain(preferences_mismatch), + package_name, + range, + version_maps, + installed_packages, + exclusions, + resolver_markers, + ) + } } + } - // Check for a locally installed distribution that matches the preferred version. - if !exclusions.contains(package_name) { - let installed_dists = installed_packages.get_packages(package_name); - match installed_dists.as_slice() { - [] => {} - [dist] => { - if dist.version() == version { - debug!("Found installed version of {dist} that satisfies preference in {range}"); - - return Some(Candidate { - name: package_name, - version, - dist: CandidateDist::Compatible(CompatibleDist::InstalledDist(dist)), - choice_kind: VersionChoiceKind::Preference, - }); + /// Return the first preference that satisfies the current range and is allowed. + fn get_preferred_from_iter<'a, InstalledPackages: InstalledPackagesProvider>( + &'a self, + preferences: impl Iterator, &'a Version)>, + package_name: &'a PackageName, + range: &Range, + version_maps: &'a [VersionMap], + installed_packages: &'a InstalledPackages, + exclusions: &Exclusions, + resolver_markers: &ResolverMarkers, + ) -> Option> { + for (_marker, version) in preferences { + // Respect the version range for this requirement. + if !range.contains(version) { + continue; + } + + // Check for a locally installed distribution that matches the preferred version. + if !exclusions.contains(package_name) { + let installed_dists = installed_packages.get_packages(package_name); + match installed_dists.as_slice() { + [] => {} + [dist] => { + if dist.version() == version { + debug!("Found installed version of {dist} that satisfies preference in {range}"); + + return Some(Candidate { + name: package_name, + version, + dist: CandidateDist::Compatible(CompatibleDist::InstalledDist( + dist, + )), + choice_kind: VersionChoiceKind::Preference, + }); + } + } + // We do not consider installed distributions with multiple versions because + // during installation these must be reinstalled from the remote + _ => { + debug!("Ignoring installed versions of {package_name}: multiple distributions found"); } - } - // We do not consider installed distributions with multiple versions because - // during installation these must be reinstalled from the remote - _ => { - debug!("Ignoring installed versions of {package_name}: multiple distributions found"); } } - } - // Respect the pre-release strategy for this fork. - if version.any_prerelease() - && self.prerelease_strategy.allows(package_name, markers) != AllowPreRelease::Yes - { - return None; - } + // Respect the pre-release strategy for this fork. + if version.any_prerelease() + && self + .prerelease_strategy + .allows(package_name, resolver_markers) + != AllowPreRelease::Yes + { + continue; + } - // Check for a remote distribution that matches the preferred version - if let Some(file) = version_maps - .iter() - .find_map(|version_map| version_map.get(version)) - { - return Some(Candidate::new( - package_name, - version, - file, - VersionChoiceKind::Preference, - )); + // Check for a remote distribution that matches the preferred version + if let Some(file) = version_maps + .iter() + .find_map(|version_map| version_map.get(version)) + { + return Some(Candidate::new( + package_name, + version, + file, + VersionChoiceKind::Preference, + )); + } } - None } - /// Check for a locally installed distribution that satisfies the range. + /// Check for an installed distribution that satisfies the current range and is allowed. fn get_installed<'a, InstalledPackages: InstalledPackagesProvider>( package_name: &'a PackageName, range: &Range, @@ -217,8 +310,8 @@ impl CandidateSelector { version_maps: &'a [VersionMap], markers: &ResolverMarkers, ) -> Option { - tracing::trace!( - "selecting candidate for package {package_name} with range {range:?} with {} remote versions", + trace!( + "Selecting candidate for {package_name} with range {range} with {} remote versions", version_maps.iter().map(VersionMap::len).sum::(), ); let highest = self.use_highest_version(package_name); @@ -408,12 +501,9 @@ impl CandidateSelector { return Some(candidate); } - tracing::trace!( - "exhausted all candidates for package {:?} with range {:?} \ - after {} steps", - package_name, - range, - steps, + trace!( + "Exhausted all candidates for package {package_name} with range {range} \ + after {steps} steps", ); match prerelease { None => None, diff --git a/crates/uv-resolver/src/lock.rs b/crates/uv-resolver/src/lock.rs index 02442953ad4e..dee0608e3851 100644 --- a/crates/uv-resolver/src/lock.rs +++ b/crates/uv-resolver/src/lock.rs @@ -1138,6 +1138,10 @@ impl Distribution { &self.id.version } + pub fn fork_markers(&self) -> Option<&BTreeSet> { + self.fork_markers.as_ref() + } + /// Returns a [`VersionId`] for this package that can be used for resolution. fn version_id(&self, workspace_root: &Path) -> Result { match &self.id.source { diff --git a/crates/uv-resolver/src/preferences.rs b/crates/uv-resolver/src/preferences.rs index a4e1cbbb0c3a..2592b1e4a025 100644 --- a/crates/uv-resolver/src/preferences.rs +++ b/crates/uv-resolver/src/preferences.rs @@ -1,4 +1,4 @@ -use std::collections::hash_map::Entry; +use std::collections::BTreeSet; use std::str::FromStr; use rustc_hash::FxHashMap; @@ -22,7 +22,11 @@ pub enum PreferenceError { pub struct Preference { name: PackageName, version: Version, + /// The markers on the requirement itself (those after the semicolon). marker: Option, + /// If coming from a package with diverging versions, the markers of the forks this preference + /// is part of, otherwise `None`. + fork_markers: Option>, hashes: Vec, } @@ -53,6 +57,8 @@ impl Preference { name: requirement.name, version: specifier.version().clone(), marker: requirement.marker, + // requirements.txt doesn't have fork annotations. + fork_markers: None, hashes: entry .hashes .iter() @@ -72,6 +78,8 @@ impl Preference { name: dist.name().clone(), version: version.clone(), marker: None, + // Installed distributions don't have fork annotations. + fork_markers: None, hashes: Vec::new(), } } @@ -82,6 +90,7 @@ impl Preference { name: dist.id.name.clone(), version: dist.id.version.clone(), marker: None, + fork_markers: dist.fork_markers().cloned(), hashes: Vec::new(), } } @@ -98,58 +107,120 @@ impl Preference { } /// A set of pinned packages that should be preserved during resolution, if possible. +/// +/// The marker is the marker of the fork that resolved to the pin, if any. +/// +/// Preferences should be prioritized first by whether their marker matches and then by the order +/// they are stored, so that a lockfile has higher precedence than sibling forks. #[derive(Debug, Clone, Default)] -pub struct Preferences(FxHashMap); +pub struct Preferences(FxHashMap, Pin)>>); impl Preferences { /// Create a map of pinned packages from an iterator of [`Preference`] entries. - /// Takes ownership of the [`Preference`] entries. /// - /// The provided [`MarkerEnvironment`] will be used to filter the preferences + /// The provided [`MarkerEnvironment`] will be used to filter the preferences /// to an applicable subset. pub fn from_iter>( preferences: PreferenceIterator, markers: Option<&MarkerEnvironment>, ) -> Self { - // TODO(zanieb): We should explicitly ensure that when a package name is seen multiple times - // that the newest or oldest version is preferred depending on the resolution strategy; - // right now, the order is dependent on the given iterator. - let preferences = preferences - .into_iter() - .filter_map(|preference| { - if preference.marker.as_ref().map_or(true, |marker| { - marker.evaluate_optional_environment(markers, &[]) - }) { - Some(( - preference.name, + let mut slf = Self::default(); + for preference in preferences { + // Filter non-matching preferences when resolving for an environment. + if let Some(markers) = markers { + if !preference + .marker + .as_ref() + .map_or(true, |marker| marker.evaluate(markers, &[])) + { + trace!("Excluding {preference} from preferences due to unmatched markers"); + continue; + } + + if !preference + .fork_markers + .as_ref() + .map(|fork_markers| { + fork_markers + .iter() + .any(|marker| marker.evaluate(markers, &[])) + }) + .unwrap_or(true) + { + trace!("Excluding {preference} from preferences due to unmatched fork markers"); + continue; + } + } + + // Flatten the list of markers into individual entries. + if let Some(fork_markers) = &preference.fork_markers { + for fork_marker in fork_markers { + slf.insert( + preference.name.clone(), + Some(fork_marker.clone()), Pin { - version: preference.version, - hashes: preference.hashes, + version: preference.version.clone(), + hashes: preference.hashes.clone(), }, - )) - } else { - trace!("Excluding {preference} from preferences due to unmatched markers"); - None + ); } - }) - .collect(); + } else { + slf.insert( + preference.name, + None, + Pin { + version: preference.version, + hashes: preference.hashes, + }, + ); + } + } - Self(preferences) + slf } - /// Return the [`Entry`] for a package in the preferences. - pub fn entry(&mut self, package_name: PackageName) -> Entry { - self.0.entry(package_name) + /// Insert a preference at the back. + pub(crate) fn insert( + &mut self, + package_name: PackageName, + markers: Option, + pin: impl Into, + ) { + self.0 + .entry(package_name) + .or_default() + .push((markers, pin.into())); } /// Returns an iterator over the preferences. - pub fn iter(&self) -> impl Iterator { - self.0.iter().map(|(name, pin)| (name, pin.version())) + pub fn iter( + &self, + ) -> impl Iterator< + Item = ( + &PackageName, + impl Iterator, &Version)>, + ), + > { + self.0.iter().map(|(name, preferences)| { + ( + name, + preferences + .iter() + .map(|(markers, pin)| (markers.as_ref(), pin.version())), + ) + }) } /// Return the pinned version for a package, if any. - pub(crate) fn version(&self, package_name: &PackageName) -> Option<&Version> { - self.0.get(package_name).map(Pin::version) + pub(crate) fn get( + &self, + package_name: &PackageName, + ) -> impl Iterator, &Version)> { + self.0 + .get(package_name) + .into_iter() + .flatten() + .map(|(markers, pin)| (markers.as_ref(), pin.version())) } /// Return the hashes for a package, if the version matches that of the pin. @@ -160,8 +231,10 @@ impl Preferences { ) -> Option<&[HashDigest]> { self.0 .get(package_name) - .filter(|pin| pin.version() == version) - .map(Pin::hashes) + .into_iter() + .flatten() + .find(|(_markers, pin)| pin.version() == version) + .map(|(_markers, pin)| pin.hashes()) } } @@ -173,7 +246,7 @@ impl std::fmt::Display for Preference { /// The pinned data associated with a package in a locked `requirements.txt` file (e.g., `flask==1.2.3`). #[derive(Debug, Clone)] -pub struct Pin { +pub(crate) struct Pin { version: Version, hashes: Vec, } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index cb1e8575e7e3..064f8636c07e 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -1,7 +1,6 @@ //! 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, BTreeSet, VecDeque}; use std::fmt::{Display, Formatter, Write}; use std::ops::Bound; @@ -389,11 +388,15 @@ impl ResolverState