Arc-wrap PubGrubPackage for cheap cloning in pubgrub

crates/uv-resolver/Cargo.toml

@@ -37,6 +37,7 @@ anstream = { workspace = true }
 anyhow = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true, features = ["derive"], optional = true }
+dashmap = { workspace = true }
 derivative = { workspace = true }
 either = { workspace = true }
 futures = { workspace = true }
@@ -56,7 +57,6 @@ tokio = { workspace = true }
 tokio-stream = { workspace = true }
 tracing = { workspace = true }
 url = { workspace = true }
-dashmap = { workspace = true }
 
 [dev-dependencies]
 uv-interpreter = { workspace = true }

crates/uv-resolver/src/error.rs

@@ -16,7 +16,7 @@ use uv_normalize::PackageName;
 
 use crate::candidate_selector::CandidateSelector;
 use crate::dependency_provider::UvDependencyProvider;
-use crate::pubgrub::{PubGrubPackage, PubGrubPython, PubGrubReportFormatter};
+use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner, PubGrubPython, PubGrubReportFormatter};
 use crate::python_requirement::PythonRequirement;
 use crate::resolver::{
     FxOnceMap, IncompletePackage, UnavailablePackage, UnavailableReason, VersionsResponse,
@@ -114,7 +114,7 @@ impl<T> From<tokio::sync::mpsc::error::SendError<T>> for ResolveError {
 }
 
 /// Given a [`DerivationTree`], collapse any [`External::FromDependencyOf`] incompatibilities
-/// wrap an [`PubGrubPackage::Extra`] package.
+/// wrap an [`PubGrubPackageInner::Extra`] package.
 fn collapse_extra_proxies(
     derivation_tree: &mut DerivationTree<PubGrubPackage, Range<Version>, UnavailableReason>,
 ) {
@@ -126,22 +126,16 @@ fn collapse_extra_proxies(
                 Arc::make_mut(&mut derived.cause2),
             ) {
                 (
-                    DerivationTree::External(External::FromDependencyOf(
-                        PubGrubPackage::Extra { .. },
-                        ..,
-                    )),
+                    DerivationTree::External(External::FromDependencyOf(package, ..)),
                     ref mut cause,
-                ) => {
+                ) if matches!(&**package, PubGrubPackageInner::Extra { .. }) => {
                     collapse_extra_proxies(cause);
                     *derivation_tree = cause.clone();
                 }
                 (
                     ref mut cause,
-                    DerivationTree::External(External::FromDependencyOf(
-                        PubGrubPackage::Extra { .. },
-                        ..,
-                    )),
-                ) => {
+                    DerivationTree::External(External::FromDependencyOf(package, ..)),
+                ) if matches!(&**package, PubGrubPackageInner::Extra { .. }) => {
                     collapse_extra_proxies(cause);
                     *derivation_tree = cause.clone();
                 }
@@ -241,22 +235,22 @@ impl NoSolutionError {
     ) -> Self {
         let mut available_versions = IndexMap::default();
         for package in self.derivation_tree.packages() {
-            match package {
-                PubGrubPackage::Root(_) => {}
-                PubGrubPackage::Python(PubGrubPython::Installed) => {
+            match &**package {
+                PubGrubPackageInner::Root(_) => {}
+                PubGrubPackageInner::Python(PubGrubPython::Installed) => {
                     available_versions.insert(
                         package.clone(),
                         BTreeSet::from([python_requirement.installed().deref().clone()]),
                     );
                 }
-                PubGrubPackage::Python(PubGrubPython::Target) => {
+                PubGrubPackageInner::Python(PubGrubPython::Target) => {
                     available_versions.insert(
                         package.clone(),
                         BTreeSet::from([python_requirement.target().deref().clone()]),
                     );
                 }
-                PubGrubPackage::Extra { .. } => {}
-                PubGrubPackage::Package { name, .. } => {
+                PubGrubPackageInner::Extra { .. } => {}
+                PubGrubPackageInner::Package { name, .. } => {
                     // Avoid including available versions for packages that exist in the derivation
                     // tree, but were never visited during resolution. We _may_ have metadata for
                     // these packages, but it's non-deterministic, and omitting them ensures that
@@ -304,7 +298,7 @@ impl NoSolutionError {
     ) -> Self {
         let mut new = FxHashMap::default();
         for package in self.derivation_tree.packages() {
-            if let PubGrubPackage::Package { name, .. } = package {
+            if let PubGrubPackageInner::Package { name, .. } = &**package {
                 if let Some(reason) = unavailable_packages.get(name) {
                     new.insert(name.clone(), reason.clone());
                 }
@@ -322,7 +316,7 @@ impl NoSolutionError {
     ) -> Self {
         let mut new = FxHashMap::default();
         for package in self.derivation_tree.packages() {
-            if let PubGrubPackage::Package { name, .. } = package {
+            if let PubGrubPackageInner::Package { name, .. } = &**package {
                 if let Some(versions) = incomplete_packages.get(name) {
                     for entry in versions.iter() {
                         let (version, reason) = entry.pair();

crates/uv-resolver/src/pubgrub/dependencies.rs

@@ -10,7 +10,7 @@ use uv_configuration::{Constraints, Overrides};
 use uv_normalize::{ExtraName, PackageName};
 
 use crate::pubgrub::specifier::PubGrubSpecifier;
-use crate::pubgrub::PubGrubPackage;
+use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner};
 use crate::resolver::{Locals, Urls};
 use crate::ResolveError;
 
@@ -102,8 +102,8 @@ fn add_requirements(
         })) {
             let PubGrubRequirement { package, version } = result?;
 
-            match &package {
-                PubGrubPackage::Package { name, .. } => {
+            match &*package {
+                PubGrubPackageInner::Package { name, .. } => {
                     // Detect self-dependencies.
                     if source_name.is_some_and(|source_name| source_name == name) {
                         warn!("{name} has a dependency on itself");
@@ -112,7 +112,7 @@ fn add_requirements(
 
                     dependencies.push((package.clone(), version.clone()));
                 }
-                PubGrubPackage::Extra { name, extra, .. } => {
+                PubGrubPackageInner::Extra { name, extra, .. } => {
                     // Recursively add the dependencies of the current package (e.g., `black` depending on
                     // `black[colorama]`).
                     if source_name.is_some_and(|source_name| source_name == name) {
@@ -158,7 +158,7 @@ fn add_requirements(
                     PubGrubRequirement::from_constraint(constraint, urls, locals)?;
 
                 // Ignore self-dependencies.
-                if let PubGrubPackage::Package { name, .. } = &package {
+                if let PubGrubPackageInner::Package { name, .. } = &*package {
                     // Detect self-dependencies.
                     if source_name.is_some_and(|source_name| source_name == name) {
                         warn!("{name} has a dependency on itself");
@@ -249,12 +249,12 @@ impl PubGrubRequirement {
                 }
 
                 Ok(Self {
-                    package: PubGrubPackage::Package {
+                    package: PubGrubPackage::from(PubGrubPackageInner::Package {
                         name: requirement.name.clone(),
                         extra,
                         marker: None,
                         url: Some(expected.clone()),
-                    },
+                    }),
                     version: Range::full(),
                 })
             }
@@ -275,12 +275,12 @@ impl PubGrubRequirement {
                 }
 
                 Ok(Self {
-                    package: PubGrubPackage::Package {
+                    package: PubGrubPackage::from(PubGrubPackageInner::Package {
                         name: requirement.name.clone(),
                         extra,
                         marker: None,
                         url: Some(expected.clone()),
-                    },
+                    }),
                     version: Range::full(),
                 })
             }
@@ -301,12 +301,12 @@ impl PubGrubRequirement {
                 }
 
                 Ok(Self {
-                    package: PubGrubPackage::Package {
+                    package: PubGrubPackage::from(PubGrubPackageInner::Package {
                         name: requirement.name.clone(),
                         extra,
                         marker: None,
                         url: Some(expected.clone()),
-                    },
+                    }),
                     version: Range::full(),
                 })
             }

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

@@ -1,6 +1,6 @@
 pub(crate) use crate::pubgrub::dependencies::{PubGrubDependencies, PubGrubRequirement};
 pub(crate) use crate::pubgrub::distribution::PubGrubDistribution;
-pub(crate) use crate::pubgrub::package::{PubGrubPackage, PubGrubPython};
+pub(crate) use crate::pubgrub::package::{PubGrubPackage, PubGrubPackageInner, PubGrubPython};
 pub(crate) use crate::pubgrub::priority::{PubGrubPriorities, PubGrubPriority};
 pub(crate) use crate::pubgrub::report::PubGrubReportFormatter;
 pub(crate) use crate::pubgrub::specifier::PubGrubSpecifier;

crates/uv-resolver/src/pubgrub/package.rs

@@ -1,9 +1,36 @@
 use distribution_types::VerbatimParsedUrl;
 use pep508_rs::MarkerTree;
+use std::fmt::{Display, Formatter};
+use std::ops::Deref;
+use std::sync::Arc;
 use uv_normalize::{ExtraName, PackageName};
 
 use crate::resolver::Urls;
 
+/// [`Arc`] wrapper around [`PubGrubPackageInner`] to make cloning (inside PubGrub) cheap.
+#[derive(Debug, Clone, Eq, Hash, PartialEq, PartialOrd, Ord)]
+pub struct PubGrubPackage(Arc<PubGrubPackageInner>);
+
+impl Deref for PubGrubPackage {
+    type Target = PubGrubPackageInner;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl Display for PubGrubPackage {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        Display::fmt(&self.0, f)
+    }
+}
+
+impl From<PubGrubPackageInner> for PubGrubPackage {
+    fn from(package: PubGrubPackageInner) -> Self {
+        Self(Arc::new(package))
+    }
+}
+
 /// A PubGrub-compatible wrapper around a "Python package", with two notable characteristics:
 ///
 /// 1. Includes a [`PubGrubPackage::Root`] variant, to satisfy PubGrub's requirement that a
@@ -12,7 +39,7 @@ use crate::resolver::Urls;
 ///    package (e.g., `black[colorama]`), and mark it as a dependency of the real package (e.g.,
 ///    `black`). We then discard the virtual packages at the end of the resolution process.
 #[derive(Debug, Clone, Eq, Hash, PartialEq, PartialOrd, Ord)]
-pub enum PubGrubPackage {
+pub enum PubGrubPackageInner {
     /// The root package, which is used to start the resolution process.
     Root(Option<PackageName>),
     /// A Python version.
@@ -91,19 +118,19 @@ impl PubGrubPackage {
     ) -> Self {
         let url = urls.get(&name).cloned();
         if let Some(extra) = extra {
-            Self::Extra {
+            Self(Arc::new(PubGrubPackageInner::Extra {
                 name,
                 extra,
                 marker,
                 url,
-            }
+            }))
         } else {
-            Self::Package {
+            Self(Arc::new(PubGrubPackageInner::Package {
                 name,
                 extra,
                 marker,
                 url,
-            }
+            }))
         }
     }
 }
@@ -116,7 +143,7 @@ pub enum PubGrubPython {
     Target,
 }
 
-impl std::fmt::Display for PubGrubPackage {
+impl std::fmt::Display for PubGrubPackageInner {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Root(name) => {

crates/uv-resolver/src/pubgrub/priority.rs

@@ -7,6 +7,7 @@ use pep440_rs::Version;
 use uv_normalize::PackageName;
 
 use crate::pubgrub::package::PubGrubPackage;
+use crate::pubgrub::PubGrubPackageInner;
 
 /// A prioritization map to guide the PubGrub resolution process.
 ///
@@ -24,14 +25,14 @@ impl PubGrubPriorities {
     /// Add a [`PubGrubPackage`] to the priority map.
     pub(crate) fn insert(&mut self, package: &PubGrubPackage, version: &Range<Version>) {
         let next = self.0.len();
-        match package {
-            PubGrubPackage::Root(_) => {}
-            PubGrubPackage::Python(_) => {}
+        match &**package {
+            PubGrubPackageInner::Root(_) => {}
+            PubGrubPackageInner::Python(_) => {}
 
-            PubGrubPackage::Extra {
+            PubGrubPackageInner::Extra {
                 name, url: None, ..
             }
-            | PubGrubPackage::Package {
+            | PubGrubPackageInner::Package {
                 name, url: None, ..
             } => {
                 match self.0.entry(name.clone()) {
@@ -66,10 +67,10 @@ impl PubGrubPriorities {
                     }
                 }
             }
-            PubGrubPackage::Extra {
+            PubGrubPackageInner::Extra {
                 name, url: Some(_), ..
             }
-            | PubGrubPackage::Package {
+            | PubGrubPackageInner::Package {
                 name, url: Some(_), ..
             } => {
                 match self.0.entry(name.clone()) {
@@ -101,11 +102,11 @@ impl PubGrubPriorities {
 
     /// Return the [`PubGrubPriority`] of the given package, if it exists.
     pub(crate) fn get(&self, package: &PubGrubPackage) -> Option<PubGrubPriority> {
-        match package {
-            PubGrubPackage::Root(_) => Some(PubGrubPriority::Root),
-            PubGrubPackage::Python(_) => Some(PubGrubPriority::Root),
-            PubGrubPackage::Extra { name, .. } => self.0.get(name).copied(),
-            PubGrubPackage::Package { name, .. } => self.0.get(name).copied(),
+        match &**package {
+            PubGrubPackageInner::Root(_) => Some(PubGrubPriority::Root),
+            PubGrubPackageInner::Python(_) => Some(PubGrubPriority::Root),
+            PubGrubPackageInner::Extra { name, .. } => self.0.get(name).copied(),
+            PubGrubPackageInner::Package { name, .. } => self.0.get(name).copied(),
         }
     }
 }