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.

crates/uv-resolver/src/preferences.rs

@@ -1,5 +1,5 @@
+use std::collections::hash_map::Entry;
 use std::str::FromStr;
-use std::sync::Arc;
 
 use rustc_hash::FxHashMap;
 use tracing::trace;
@@ -99,7 +99,7 @@ impl Preference {
 
 /// A set of pinned packages that should be preserved during resolution, if possible.
 #[derive(Debug, Clone, Default)]
-pub struct Preferences(Arc<FxHashMap<PackageName, Pin>>);
+pub struct Preferences(FxHashMap<PackageName, Pin>);
 
 impl Preferences {
     /// Create a map of pinned packages from an iterator of [`Preference`] entries.
@@ -134,7 +134,12 @@ impl Preferences {
             })
             .collect();
 
-        Self(Arc::new(preferences))
+        Self(preferences)
+    }
+
+    /// Return the [`Entry`] for a package in the preferences.
+    pub fn entry(&mut self, package_name: PackageName) -> Entry<PackageName, Pin> {
+        self.0.entry(package_name)
     }
 
     /// Returns an iterator over the preferences.
@@ -168,7 +173,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)]
-struct Pin {
+pub struct Pin {
     version: Version,
     hashes: Vec<HashDigest>,
 }
@@ -184,3 +189,12 @@ impl Pin {
         &self.hashes
     }
 }
+
+impl From<Version> for Pin {
+    fn from(version: Version) -> Self {
+        Self {
+            version,
+            hashes: Vec::new(),
+        }
+    }
+}

crates/uv-resolver/src/resolver/mod.rs

@@ -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 {

crates/uv/tests/lock_scenarios.rs

@@ -16,10 +16,6 @@ mod common;
 
 /// This test ensures that multiple non-conflicting but also non-overlapping
 /// dependency specifications with the same package name are allowed and supported.
-/// At time of writing, this provokes a fork in the resolver, but it arguably
-/// shouldn't since the requirements themselves do not conflict with one another.
-/// However, this does impact resolution. Namely, it leaves the `a>=1` fork free to
-/// choose `a==2.0.0` since it behaves as if the `a<2` constraint doesn't exist.
 ///
 /// ```text
 /// fork-allows-non-conflicting-non-overlapping-dependencies
@@ -71,7 +67,7 @@ fn fork_allows_non_conflicting_non_overlapping_dependencies() -> Result<()> {
 
     ----- stderr -----
     warning: `uv lock` is experimental and may change without warning.
-    Resolved 3 packages in [TIME]
+    Resolved 2 packages in [TIME]
     "###
     );
 
@@ -93,22 +89,12 @@ fn fork_allows_non_conflicting_non_overlapping_dependencies() -> Result<()> {
             { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_allows_non_conflicting_non_overlapping_dependencies_a-1.0.0-py3-none-any.whl#sha256=37c13aa13cca009990929df08bed3d9de26e1d405a5ebd16ec0c3baef6899b23", hash = "sha256:37c13aa13cca009990929df08bed3d9de26e1d405a5ebd16ec0c3baef6899b23" },
         ]
 
-        [[distribution]]
-        name = "package-a"
-        version = "2.0.0"
-        source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }
-        sdist = { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_allows_non_conflicting_non_overlapping_dependencies_a-2.0.0.tar.gz#sha256=0b4ca63d060f4daa2269c08b7083f594e096b94e1bcbde53d212c65b52378358", hash = "sha256:0b4ca63d060f4daa2269c08b7083f594e096b94e1bcbde53d212c65b52378358" }
-        wheels = [
-            { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_allows_non_conflicting_non_overlapping_dependencies_a-2.0.0-py3-none-any.whl#sha256=35168196ad80d0f2822191a47e3a805b4ad527280c8b84e7eed77b7fee505497", hash = "sha256:35168196ad80d0f2822191a47e3a805b4ad527280c8b84e7eed77b7fee505497" },
-        ]
-
         [[distribution]]
         name = "project"
         version = "0.1.0"
         source = { editable = "." }
         dependencies = [
-            { name = "package-a", version = "1.0.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }, marker = "sys_platform == 'darwin'" },
+            { name = "package-a", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-            { name = "package-a", version = "2.0.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }, marker = "sys_platform == 'linux'" },
         ]
         "###
         );
@@ -1676,7 +1662,7 @@ fn fork_marker_selection() -> Result<()> {
 
     ----- stderr -----
     warning: `uv lock` is experimental and may change without warning.
-    Resolved 5 packages in [TIME]
+    Resolved 4 packages in [TIME]
     "###
     );
 
@@ -1698,18 +1684,6 @@ fn fork_marker_selection() -> Result<()> {
             { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_marker_selection_a-0.1.0-py3-none-any.whl#sha256=8aecc639cc090aa80aa263fb3a9644a7cec9da215133299b8fb381cb7a6bcbb7", hash = "sha256:8aecc639cc090aa80aa263fb3a9644a7cec9da215133299b8fb381cb7a6bcbb7" },
         ]
 
-        [[distribution]]
-        name = "package-a"
-        version = "0.2.0"
-        source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }
-        sdist = { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_marker_selection_a-0.2.0.tar.gz#sha256=42abfb3ce2c13ae008e498d27c80ae39ab19e30fd56e175719b67b1c778ea632", hash = "sha256:42abfb3ce2c13ae008e498d27c80ae39ab19e30fd56e175719b67b1c778ea632" }
-        dependencies = [
-            { name = "package-b", version = "2.0.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" } },
-        ]
-        wheels = [
-            { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_marker_selection_a-0.2.0-py3-none-any.whl#sha256=65ff1ce26de8218278abb1ae190fe70d031de79833d85231112208672566b9c4", hash = "sha256:65ff1ce26de8218278abb1ae190fe70d031de79833d85231112208672566b9c4" },
-        ]
-
         [[distribution]]
         name = "package-b"
         version = "1.0.0"
@@ -1733,8 +1707,7 @@ fn fork_marker_selection() -> Result<()> {
         version = "0.1.0"
         source = { editable = "." }
         dependencies = [
-            { name = "package-a", version = "0.1.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" } },
+            { name = "package-a" },
-            { name = "package-a", version = "0.2.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" } },
             { name = "package-b", version = "1.0.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }, marker = "sys_platform == 'darwin'" },
             { name = "package-b", version = "2.0.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }, marker = "sys_platform == 'linux'" },
         ]
@@ -1817,7 +1790,7 @@ fn fork_marker_track() -> Result<()> {
 
     ----- stderr -----
     warning: `uv lock` is experimental and may change without warning.
-    Resolved 6 packages in [TIME]
+    Resolved 5 packages in [TIME]
     "###
     );
 
@@ -1842,18 +1815,6 @@ fn fork_marker_track() -> Result<()> {
             { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_marker_track_a-1.3.1-py3-none-any.whl#sha256=79e82592fe6644839cdb6dc73d3d54fc543f0e0f28cce26e221a6c1e30072104", hash = "sha256:79e82592fe6644839cdb6dc73d3d54fc543f0e0f28cce26e221a6c1e30072104" },
         ]
 
-        [[distribution]]
-        name = "package-a"
-        version = "4.3.0"
-        source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }
-        sdist = { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_marker_track_a-4.3.0.tar.gz#sha256=ce810c2e0922cff256d3050167c0d2a041955d389d21280fd684ab986dfdb1f5", hash = "sha256:ce810c2e0922cff256d3050167c0d2a041955d389d21280fd684ab986dfdb1f5" }
-        dependencies = [
-            { name = "package-b", version = "2.8", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" } },
-        ]
-        wheels = [
-            { url = "https://astral-sh.github.io/packse/0.3.29/files/fork_marker_track_a-4.3.0-py3-none-any.whl#sha256=fb90bca8d00206119df736f59a9c4e18e104a9321b8ea91f19400a119b77ef99", hash = "sha256:fb90bca8d00206119df736f59a9c4e18e104a9321b8ea91f19400a119b77ef99" },
-        ]
-
         [[distribution]]
         name = "package-b"
         version = "2.7"
@@ -1886,8 +1847,7 @@ fn fork_marker_track() -> Result<()> {
         version = "0.1.0"
         source = { editable = "." }
         dependencies = [
-            { name = "package-a", version = "1.3.1", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" } },
+            { name = "package-a" },
-            { name = "package-a", version = "4.3.0", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" } },
             { name = "package-b", version = "2.7", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }, marker = "sys_platform == 'darwin'" },
             { name = "package-b", version = "2.8", source = { registry = "https://astral-sh.github.io/packse/0.3.29/simple-html/" }, marker = "sys_platform == 'linux'" },
         ]