Skip to content

Commit

Permalink
Arc-wrap PubGrubPackage for cheap cloning in pubgrub
Browse files Browse the repository at this point in the history
Pubgrub stores incompatibilities as (package name, version range) tuples, meaning it needs to clone the package name for each incompatibility, and each non-borrowed operation on incompatibilities. #3673 made me realize that `PubGrubPackage` has gotten large (expensive to copy), so like `Version` and other structs, i've added an `Arc` wrapper around it.

It's a pity clippy forbids `.deref()`, it's less opaque than `&**` and has IDE support (clicking on `.deref()` jumps to the right impl).

## Benchmarks

It looks like this matters most for complex resolutions which, i assume because they carry larger `PubGrubPackageInner::Package` and `PubGrubPackageInner::Extra` types.

```bash
hyperfine --warmup 5 "./uv-main pip compile -q ./scripts/requirements/jupyter.in" "./uv-branch pip compile -q ./scripts/requirements/jupyter.in"
hyperfine --warmup 5 "./uv-main pip compile -q ./scripts/requirements/airflow.in" "./uv-branch pip compile -q ./scripts/requirements/airflow.in"
hyperfine --warmup 5 "./uv-main pip compile -q ./scripts/requirements/boto3.in" "./uv-branch pip compile -q ./scripts/requirements/boto3.in"
```

```
Benchmark 1: ./uv-main pip compile -q ./scripts/requirements/jupyter.in
  Time (mean ± σ):      18.2 ms ±   1.6 ms    [User: 14.4 ms, System: 26.0 ms]
  Range (min … max):    15.8 ms …  22.5 ms    181 runs

Benchmark 2: ./uv-branch pip compile -q ./scripts/requirements/jupyter.in
  Time (mean ± σ):      17.8 ms ±   1.4 ms    [User: 14.4 ms, System: 25.3 ms]
  Range (min … max):    15.4 ms …  23.1 ms    159 runs

Summary
  ./uv-branch pip compile -q ./scripts/requirements/jupyter.in ran
    1.02 ± 0.12 times faster than ./uv-main pip compile -q ./scripts/requirements/jupyter.in
```

```
Benchmark 1: ./uv-main pip compile -q ./scripts/requirements/airflow.in
  Time (mean ± σ):     153.7 ms ±   3.5 ms    [User: 165.2 ms, System: 157.6 ms]
  Range (min … max):   150.4 ms … 163.0 ms    19 runs

Benchmark 2: ./uv-branch pip compile -q ./scripts/requirements/airflow.in
  Time (mean ± σ):     123.9 ms ±   4.6 ms    [User: 152.4 ms, System: 133.8 ms]
  Range (min … max):   118.4 ms … 138.1 ms    24 runs

Summary
  ./uv-branch pip compile -q ./scripts/requirements/airflow.in ran
    1.24 ± 0.05 times faster than ./uv-main pip compile -q ./scripts/requirements/airflow.in
```

```
Benchmark 1: ./uv-main pip compile -q ./scripts/requirements/boto3.in
  Time (mean ± σ):     327.0 ms ±   3.8 ms    [User: 344.5 ms, System: 71.6 ms]
  Range (min … max):   322.7 ms … 334.6 ms    10 runs

Benchmark 2: ./uv-branch pip compile -q ./scripts/requirements/boto3.in
  Time (mean ± σ):     311.2 ms ±   3.1 ms    [User: 339.3 ms, System: 63.1 ms]
  Range (min … max):   307.8 ms … 317.0 ms    10 runs

Summary
  ./uv-branch pip compile -q ./scripts/requirements/boto3.in ran
    1.05 ± 0.02 times faster than ./uv-main pip compile -q ./scripts/requirements/boto3.in
```
  • Loading branch information
konstin committed May 21, 2024
1 parent acde048 commit b4e9a12
Show file tree
Hide file tree
Showing 10 changed files with 159 additions and 122 deletions.
2 changes: 1 addition & 1 deletion crates/uv-resolver/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand All @@ -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 }
Expand Down
34 changes: 14 additions & 20 deletions crates/uv-resolver/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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>,
) {
Expand All @@ -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();
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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());
}
Expand All @@ -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();
Expand Down
22 changes: 11 additions & 11 deletions crates/uv-resolver/src/pubgrub/dependencies.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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");
Expand All @@ -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) {
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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(),
})
}
Expand All @@ -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(),
})
}
Expand All @@ -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(),
})
}
Expand Down
2 changes: 1 addition & 1 deletion crates/uv-resolver/src/pubgrub/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down
39 changes: 33 additions & 6 deletions crates/uv-resolver/src/pubgrub/package.rs
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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,
}
}))
}
}
}
Expand All @@ -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) => {
Expand Down
25 changes: 13 additions & 12 deletions crates/uv-resolver/src/pubgrub/priority.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand All @@ -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()) {
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -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(),
}
}
}
Expand Down
Loading

0 comments on commit b4e9a12

Please sign in to comment.