Use FxHasher in resolver (#3641)

## Summary

We can use `FxHasher` in a few more places for string and version keys.
This gives a consistent ~2% improvement to warm resolves.

crates/once-map/src/lib.rs

@@ -1,5 +1,5 @@
 use std::borrow::Borrow;
-use std::hash::Hash;
+use std::hash::{BuildHasher, Hash, RandomState};
 use std::sync::Arc;
 
 use dashmap::DashMap;
@@ -14,11 +14,11 @@ use tokio::sync::Notify;
 ///
 /// Note that this always clones the value out of the underlying map. Because
 /// of this, it's common to wrap the `V` in an `Arc<V>` to make cloning cheap.
-pub struct OnceMap<K, V> {
-    items: DashMap<K, Value<V>>,
+pub struct OnceMap<K, V, H = RandomState> {
+    items: DashMap<K, Value<V>, H>,
 }
 
-impl<K: Eq + Hash, V: Clone> OnceMap<K, V> {
+impl<K: Eq + Hash, V: Clone, H: BuildHasher + Clone> OnceMap<K, V, H> {
     /// Register that you want to start a job.
     ///
     /// If this method returns `true`, you need to start a job and call [`OnceMap::done`] eventually
@@ -97,10 +97,10 @@ impl<K: Eq + Hash, V: Clone> OnceMap<K, V> {
     }
 }
 
-impl<K: Eq + Hash + Clone, V> Default for OnceMap<K, V> {
+impl<K: Eq + Hash + Clone, V, H: Default + BuildHasher + Clone> Default for OnceMap<K, V, H> {
     fn default() -> Self {
         Self {
-            items: DashMap::new(),
+            items: DashMap::with_hasher(H::default()),
         }
     }
 }

crates/uv-resolver/src/error.rs

@@ -6,11 +6,10 @@ use std::sync::Arc;
 use indexmap::IndexMap;
 use pubgrub::range::Range;
 use pubgrub::report::{DefaultStringReporter, DerivationTree, External, Reporter};
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
-use dashmap::{DashMap, DashSet};
+use dashmap::DashMap;
 use distribution_types::{BuiltDist, IndexLocations, InstalledDist, ParsedUrlError, SourceDist};
-use once_map::OnceMap;
 use pep440_rs::Version;
 use pep508_rs::Requirement;
 use uv_normalize::PackageName;
@@ -19,7 +18,9 @@ use crate::candidate_selector::CandidateSelector;
 use crate::dependency_provider::UvDependencyProvider;
 use crate::pubgrub::{PubGrubPackage, PubGrubPython, PubGrubReportFormatter};
 use crate::python_requirement::PythonRequirement;
-use crate::resolver::{IncompletePackage, UnavailablePackage, UnavailableReason, VersionsResponse};
+use crate::resolver::{
+    FxOnceMap, IncompletePackage, UnavailablePackage, UnavailableReason, VersionsResponse,
+};
 
 #[derive(Debug, thiserror::Error)]
 pub enum ResolveError {
@@ -235,8 +236,8 @@ impl NoSolutionError {
     pub(crate) fn with_available_versions(
         mut self,
         python_requirement: &PythonRequirement,
-        visited: &DashSet<PackageName>,
-        package_versions: &OnceMap<PackageName, Arc<VersionsResponse>>,
+        visited: &FxHashSet<PackageName>,
+        package_versions: &FxOnceMap<PackageName, Arc<VersionsResponse>>,
     ) -> Self {
         let mut available_versions = IndexMap::default();
         for package in self.derivation_tree.packages() {

crates/uv-resolver/src/resolution/graph.rs

@@ -10,7 +10,6 @@ use distribution_types::{
     Dist, DistributionMetadata, Name, ParsedUrlError, Requirement, ResolvedDist, VersionId,
     VersionOrUrlRef,
 };
-use once_map::OnceMap;
 use pep440_rs::{Version, VersionSpecifier};
 use pep508_rs::MarkerEnvironment;
 use uv_normalize::{ExtraName, PackageName};
@@ -22,6 +21,7 @@ use crate::preferences::Preferences;
 use crate::pubgrub::{PubGrubDistribution, PubGrubPackage};
 use crate::redirect::url_to_precise;
 use crate::resolution::AnnotatedDist;
+use crate::resolver::FxOnceMap;
 use crate::{
     lock, InMemoryIndex, Lock, LockError, Manifest, MetadataResponse, ResolveError,
     VersionsResponse,
@@ -45,8 +45,8 @@ impl ResolutionGraph {
     pub(crate) fn from_state(
         selection: &SelectedDependencies<UvDependencyProvider>,
         pins: &FilePins,
-        packages: &OnceMap<PackageName, Arc<VersionsResponse>>,
-        distributions: &OnceMap<VersionId, Arc<MetadataResponse>>,
+        packages: &FxOnceMap<PackageName, Arc<VersionsResponse>>,
+        distributions: &FxOnceMap<VersionId, Arc<MetadataResponse>>,
         state: &State<UvDependencyProvider>,
         preferences: &Preferences,
         editables: Editables,

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

@@ -1,7 +1,9 @@
+use std::hash::BuildHasherDefault;
 use std::sync::Arc;
 
 use distribution_types::VersionId;
 use once_map::OnceMap;
+use rustc_hash::FxHasher;
 use uv_normalize::PackageName;
 
 use crate::resolver::provider::{MetadataResponse, VersionsResponse};
@@ -14,20 +16,22 @@ pub struct InMemoryIndex(Arc<SharedInMemoryIndex>);
 struct SharedInMemoryIndex {
     /// A map from package name to the metadata for that package and the index where the metadata
     /// came from.
-    packages: OnceMap<PackageName, Arc<VersionsResponse>>,
+    packages: FxOnceMap<PackageName, Arc<VersionsResponse>>,
 
     /// A map from package ID to metadata for that distribution.
-    distributions: OnceMap<VersionId, Arc<MetadataResponse>>,
+    distributions: FxOnceMap<VersionId, Arc<MetadataResponse>>,
 }
 
+pub(crate) type FxOnceMap<K, V> = OnceMap<K, V, BuildHasherDefault<FxHasher>>;
+
 impl InMemoryIndex {
     /// Returns a reference to the package metadata map.
-    pub fn packages(&self) -> &OnceMap<PackageName, Arc<VersionsResponse>> {
+    pub fn packages(&self) -> &FxOnceMap<PackageName, Arc<VersionsResponse>> {
         &self.0.packages
     }
 
     /// Returns a reference to the distribution metadata map.
-    pub fn distributions(&self) -> &OnceMap<VersionId, Arc<MetadataResponse>> {
+    pub fn distributions(&self) -> &FxOnceMap<VersionId, Arc<MetadataResponse>> {
         &self.0.distributions
     }
 }

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

@@ -7,8 +7,8 @@ use std::sync::Arc;
 use std::thread;
 
 use anyhow::Result;
-use dashmap::{DashMap, DashSet};
-use futures::{FutureExt, StreamExt};
+use dashmap::DashMap;
+use futures::{FutureExt, StreamExt, TryFutureExt};
 use itertools::Itertools;
 use pubgrub::error::PubGrubError;
 use pubgrub::range::Range;
@@ -49,6 +49,7 @@ use crate::pubgrub::{
 use crate::python_requirement::PythonRequirement;
 use crate::resolution::ResolutionGraph;
 use crate::resolver::batch_prefetch::BatchPrefetcher;
+pub(crate) use crate::resolver::index::FxOnceMap;
 pub use crate::resolver::index::InMemoryIndex;
 pub use crate::resolver::provider::{
     DefaultResolverProvider, MetadataResponse, PackageVersionsResult, ResolverProvider,
@@ -202,8 +203,6 @@ struct ResolverState<InstalledPackages: InstalledPackagesProvider> {
     unavailable_packages: DashMap<PackageName, UnavailablePackage>,
     /// Incompatibilities for packages that are unavailable at specific versions.
     incomplete_packages: DashMap<PackageName, DashMap<Version, IncompletePackage>>,
-    /// The set of all registry-based packages visited during resolution.
-    visited: DashSet<PackageName>,
     reporter: Option<Arc<dyn Reporter>>,
 }
 
@@ -286,7 +285,6 @@ impl<Provider: ResolverProvider, InstalledPackages: InstalledPackagesProvider>
             index: index.clone(),
             unavailable_packages: DashMap::default(),
             incomplete_packages: DashMap::default(),
-            visited: DashSet::default(),
             selector: CandidateSelector::for_resolution(options, &manifest, markers),
             dependency_mode: options.dependency_mode,
             urls: Urls::from_manifest(&manifest, markers, options.dependency_mode)?,
@@ -332,7 +330,11 @@ impl<Provider: ResolverProvider, InstalledPackages: InstalledPackagesProvider>
         let (request_sink, request_stream) = mpsc::channel(300);
 
         // Run the fetcher.
-        let requests_fut = state.clone().fetch(provider.clone(), request_stream).fuse();
+        let requests_fut = state
+            .clone()
+            .fetch(provider.clone(), request_stream)
+            .map_err(|err| (err, FxHashSet::default()))
+            .fuse();
 
         // Spawn the PubGrub solver on a dedicated thread.
         let solver = state.clone();
@@ -347,7 +349,7 @@ impl<Provider: ResolverProvider, InstalledPackages: InstalledPackagesProvider>
 
         let resolve_fut = async move {
             rx.await
-                .map_err(|_| ResolveError::ChannelClosed)
+                .map_err(|_| (ResolveError::ChannelClosed, FxHashSet::default()))
                 .and_then(|result| result)
         };
 
@@ -357,13 +359,13 @@ impl<Provider: ResolverProvider, InstalledPackages: InstalledPackagesProvider>
                 state.on_complete();
                 Ok(resolution)
             }
-            Err(err) => {
+            Err((err, visited)) => {
                 // Add version information to improve unsat error messages.
                 Err(if let ResolveError::NoSolution(err) = err {
                     ResolveError::NoSolution(
                         err.with_available_versions(
                             &state.python_requirement,
-                            &state.visited,
+                            &visited,
                             state.index.packages(),
                         )
                         .with_selector(state.selector.clone())
@@ -381,11 +383,23 @@ impl<Provider: ResolverProvider, InstalledPackages: InstalledPackagesProvider>
 }
 
 impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackages> {
-    /// Run the PubGrub solver.
     #[instrument(skip_all)]
     fn solve(
         self: Arc<Self>,
         request_sink: Sender<Request>,
+    ) -> Result<ResolutionGraph, (ResolveError, FxHashSet<PackageName>)> {
+        let mut visited = FxHashSet::default();
+        self.solve_tracked(&mut visited, request_sink)
+            .map_err(|err| (err, visited))
+    }
+
+    /// Run the PubGrub solver, updating the `visited` set for each package visited during
+    /// resolution.
+    #[instrument(skip_all)]
+    fn solve_tracked(
+        self: Arc<Self>,
+        visited: &mut FxHashSet<PackageName>,
+        request_sink: Sender<Request>,
     ) -> Result<ResolutionGraph, ResolveError> {
         let root = PubGrubPackage::Root(self.project.clone());
         let mut prefetcher = BatchPrefetcher::default();
@@ -450,6 +464,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 &state.next,
                 term_intersection.unwrap_positive(),
                 &mut state.pins,
+                visited,
                 &request_sink,
             )?;
 
@@ -681,6 +696,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
         package: &PubGrubPackage,
         range: &Range<Version>,
         pins: &mut FilePins,
+        visited: &mut FxHashSet<PackageName>,
         request_sink: &Sender<Request>,
     ) -> Result<Option<ResolverVersion>, ResolveError> {
         match package {
@@ -805,7 +821,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     .packages()
                     .wait_blocking(package_name)
                     .ok_or(ResolveError::Unregistered)?;
-                self.visited.insert(package_name.clone());
+                visited.insert(package_name.clone());
 
                 let version_maps = match *versions_response {
                     VersionsResponse::Found(ref version_maps) => version_maps.as_slice(),