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 {