Use FxHash (#151)

crates/puffin-resolver/Cargo.toml

@@ -23,12 +23,13 @@ anyhow = { workspace = true }
 bitflags = { workspace = true }
 colored = { workspace = true }
 futures = { workspace = true }
+fxhash = { workspace = true }
 once_cell = { workspace = true }
+petgraph = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 waitmap = { workspace = true }
-petgraph = "0.6.4"
 
 [dev-dependencies]
 once_cell = { version = "1.18.0" }

crates/puffin-resolver/src/resolution.rs

@@ -1,6 +1,7 @@
-use std::collections::{BTreeMap, HashMap};
+use std::hash::BuildHasherDefault;
 
 use colored::Colorize;
+use fxhash::FxHashMap;
 use petgraph::visit::EdgeRef;
 use pubgrub::range::Range;
 use pubgrub::solver::{Kind, State};
@@ -49,11 +50,11 @@ impl PinnedPackage {
 
 /// A set of packages pinned at specific versions.
 #[derive(Debug, Default)]
-pub struct Resolution(BTreeMap<PackageName, PinnedPackage>);
+pub struct Resolution(FxHashMap<PackageName, PinnedPackage>);
 
 impl Resolution {
     /// Create a new resolution from the given pinned packages.
-    pub(crate) fn new(packages: BTreeMap<PackageName, PinnedPackage>) -> Self {
+    pub(crate) fn new(packages: FxHashMap<PackageName, PinnedPackage>) -> Self {
         Self(packages)
     }
 
@@ -87,7 +88,7 @@ impl Graph {
     /// Create a new graph from the resolved `PubGrub` state.
     pub fn from_state(
         selection: &SelectedDependencies<PubGrubPackage, PubGrubVersion>,
-        pins: &HashMap<PackageName, HashMap<Version, File>>,
+        pins: &FxHashMap<PackageName, FxHashMap<Version, File>>,
         state: &State<PubGrubPackage, Range<PubGrubVersion>>,
     ) -> Self {
         // TODO(charlie): petgraph is a really heavy and unnecessary dependency here. We should
@@ -95,7 +96,8 @@ impl Graph {
         let mut graph = petgraph::graph::Graph::with_capacity(selection.len(), selection.len());
 
         // Add every package to the graph.
-        let mut inverse = HashMap::with_capacity(selection.len());
+        let mut inverse =
+            FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default());
         for (package, version) in selection {
             let PubGrubPackage::Package(package_name, None) = package else {
                 continue;

crates/puffin-resolver/src/resolver.rs

@@ -2,7 +2,6 @@
 
 use std::borrow::Borrow;
 use std::collections::hash_map::Entry;
-use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -10,6 +9,7 @@ use anyhow::Result;
 use futures::channel::mpsc::UnboundedReceiver;
 use futures::future::Either;
 use futures::{pin_mut, FutureExt, StreamExt, TryFutureExt};
+use fxhash::{FxHashMap, FxHashSet};
 use pubgrub::error::PubGrubError;
 use pubgrub::range::Range;
 use pubgrub::solver::{Incompatibility, State};
@@ -104,14 +104,14 @@ impl<'a> Resolver<'a> {
         let root = PubGrubPackage::Root;
 
         // Keep track of the packages for which we've requested metadata.
-        let mut requested_packages = HashSet::new();
-        let mut requested_versions = HashSet::new();
-        let mut pins = HashMap::new();
+        let mut requested_packages = FxHashSet::default();
+        let mut requested_versions = FxHashSet::default();
+        let mut pins = FxHashMap::default();
 
         // Start the solve.
         let mut state = State::init(root.clone(), MIN_VERSION.clone());
-        let mut added_dependencies: HashMap<PubGrubPackage, HashSet<PubGrubVersion>> =
-            HashMap::default();
+        let mut added_dependencies: FxHashMap<PubGrubPackage, FxHashSet<PubGrubVersion>> =
+            FxHashMap::default();
         let mut next = root;
 
         loop {
@@ -243,8 +243,8 @@ impl<'a> Resolver<'a> {
     async fn choose_package_version<T: Borrow<PubGrubPackage>, U: Borrow<Range<PubGrubVersion>>>(
         &self,
         mut potential_packages: Vec<(T, U)>,
-        pins: &mut HashMap<PackageName, HashMap<pep440_rs::Version, File>>,
-        in_flight: &mut HashSet<String>,
+        pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, File>>,
+        in_flight: &mut FxHashSet<String>,
         request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
     ) -> Result<(T, Option<PubGrubVersion>), ResolveError> {
         let mut selection = 0usize;
@@ -373,8 +373,8 @@ impl<'a> Resolver<'a> {
         &self,
         package: &PubGrubPackage,
         version: &PubGrubVersion,
-        pins: &mut HashMap<PackageName, HashMap<pep440_rs::Version, File>>,
-        requested_packages: &mut HashSet<PackageName>,
+        pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, File>>,
+        requested_packages: &mut FxHashSet<PackageName>,
         request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
     ) -> Result<Dependencies, ResolveError> {
         match package {

crates/puffin-resolver/src/wheel_finder.rs

@@ -2,12 +2,13 @@
 //!
 //! This is similar to running `pip install` with the `--no-deps` flag.
 
-use std::collections::BTreeMap;
+use std::hash::BuildHasherDefault;
 use std::str::FromStr;
 
 use anyhow::Result;
 use futures::future::Either;
 use futures::{StreamExt, TryFutureExt};
+use fxhash::FxHashMap;
 use tracing::debug;
 
 use pep440_rs::Version;
@@ -81,7 +82,8 @@ impl<'a> WheelFinder<'a> {
         }
 
         // Resolve the requirements.
-        let mut resolution: BTreeMap<PackageName, PinnedPackage> = BTreeMap::new();
+        let mut resolution: FxHashMap<PackageName, PinnedPackage> =
+            FxHashMap::with_capacity_and_hasher(requirements.len(), BuildHasherDefault::default());
 
         while let Some(chunk) = package_stream.next().await {
             for result in chunk {