Skip to content

Commit

Permalink
Remove extra collects when flattening requirements (#10837)
Browse files Browse the repository at this point in the history
  • Loading branch information
@charliermarsh
charliermarsh authored Jan 22, 2025
1 parent 47a31fc commit 62f8a48
Show file tree
Hide file tree
Showing 2 changed files with 138 additions and 135 deletions.
87 changes: 45 additions & 42 deletions crates/uv-resolver/src/pubgrub/dependencies.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::borrow::Cow;
use std::iter;

use either::Either;
Expand Down Expand Up @@ -26,7 +27,7 @@ pub(crate) struct PubGrubDependency {
impl PubGrubDependency {
pub(crate) fn from_requirement<'a>(
conflicts: &Conflicts,
requirement: &'a Requirement,
requirement: Cow<'a, Requirement>,
dev: Option<&'a GroupName>,
source_name: Option<&'a PackageName>,
) -> impl Iterator<Item = Self> + 'a {
Expand Down Expand Up @@ -89,58 +90,60 @@ impl PubGrubDependency {
};

// Add the package, plus any extra variants.
iter.map(|(extra, group)| PubGrubRequirement::from_requirement(requirement, extra, group))
.map(move |requirement| {
let PubGrubRequirement {
iter.map(move |(extra, group)| {
PubGrubRequirement::from_requirement(&requirement, extra, group)
})
.map(move |requirement| {
let PubGrubRequirement {
package,
version,
url,
} = requirement;
match &*package {
PubGrubPackageInner::Package { .. } => PubGrubDependency {
package,
version,
url,
},
PubGrubPackageInner::Marker { .. } => PubGrubDependency {
package,
version,
url,
} = requirement;
match &*package {
PubGrubPackageInner::Package { .. } => PubGrubDependency {
},
PubGrubPackageInner::Extra { name, .. } => {
// Detect self-dependencies.
if dev.is_none() {
debug_assert!(
source_name.is_none_or(|source_name| source_name != name),
"extras not flattened for {name}"
);
}
PubGrubDependency {
package,
version,
url,
},
PubGrubPackageInner::Marker { .. } => PubGrubDependency {
}
}
PubGrubPackageInner::Dev { name, .. } => {
// Detect self-dependencies.
if dev.is_none() {
debug_assert!(
source_name.is_none_or(|source_name| source_name != name),
"group not flattened for {name}"
);
}
PubGrubDependency {
package,
version,
url,
},
PubGrubPackageInner::Extra { name, .. } => {
// Detect self-dependencies.
if dev.is_none() {
debug_assert!(
source_name.is_none_or(|source_name| source_name != name),
"extras not flattened for {name}"
);
}
PubGrubDependency {
package,
version,
url,
}
}
PubGrubPackageInner::Dev { name, .. } => {
// Detect self-dependencies.
if dev.is_none() {
debug_assert!(
source_name.is_none_or(|source_name| source_name != name),
"group not flattened for {name}"
);
}
PubGrubDependency {
package,
version,
url,
}
}
PubGrubPackageInner::Root(_) => unreachable!("root package in dependencies"),
PubGrubPackageInner::Python(_) => {
unreachable!("python package in dependencies")
}
}
})
PubGrubPackageInner::Root(_) => unreachable!("root package in dependencies"),
PubGrubPackageInner::Python(_) => {
unreachable!("python package in dependencies")
}
}
})
}
}

Expand Down
186 changes: 93 additions & 93 deletions crates/uv-resolver/src/resolver/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1605,7 +1605,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
);

requirements
.iter()
.flat_map(|requirement| {
PubGrubDependency::from_requirement(
&self.conflicts,
Expand Down Expand Up @@ -1704,7 +1703,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
);

requirements
.iter()
.flat_map(|requirement| {
PubGrubDependency::from_requirement(
&self.conflicts,
Expand Down Expand Up @@ -1800,116 +1798,118 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
name: Option<&PackageName>,
env: &'a ResolverEnvironment,
python_requirement: &'a PythonRequirement,
) -> Vec<Cow<'a, Requirement>> {
// Start with the requirements for the current extra of the package (for an extra
// requirement) or the non-extra (regular) dependencies (if extra is None), plus
// the constraints for the current package.
let regular_and_dev_dependencies = if let Some(dev) = dev {
Either::Left(dev_dependencies.get(dev).into_iter().flatten())
} else {
Either::Right(dependencies.iter())
};
) -> impl Iterator<Item = Cow<'a, Requirement>> {
let python_marker = python_requirement.to_marker_tree();
let mut requirements = self
.requirements_for_extra(
regular_and_dev_dependencies,

if let Some(dev) = dev {
// Dependency groups can include the project itself, so no need to flatten recursive
// dependencies.
Either::Left(Either::Left(self.requirements_for_extra(
dev_dependencies.get(dev).into_iter().flatten(),
extra,
env,
python_marker,
python_requirement,
)
.collect::<Vec<_>>();

// Dependency groups can include the project itself, so no need to flatten recursive
// dependencies.
if dev.is_some() {
return requirements;
}

// Check if there are recursive self inclusions; if so, we need to go into the expensive
// branch.
if !requirements
)))
} else if !dependencies
.iter()
.any(|req| name == Some(&req.name) && !req.extras.is_empty())
{
return requirements;
}

// Transitively process all extras that are recursively included, starting with the current
// extra.
let mut seen = FxHashSet::<(ExtraName, MarkerTree)>::default();
let mut queue: VecDeque<_> = requirements
.iter()
.filter(|req| name == Some(&req.name))
.flat_map(|req| req.extras.iter().cloned().map(|extra| (extra, req.marker)))
.collect();
while let Some((extra, marker)) = queue.pop_front() {
if !seen.insert((extra.clone(), marker)) {
continue;
}
for requirement in self.requirements_for_extra(
dependencies,
Some(&extra),
// If the project doesn't define any recursive dependencies, take the fast path.
Either::Left(Either::Right(self.requirements_for_extra(
dependencies.iter(),
extra,
env,
python_marker,
python_requirement,
) {
let requirement = match requirement {
Cow::Owned(mut requirement) => {
requirement.marker.and(marker);
requirement
}
Cow::Borrowed(requirement) => {
let mut marker = marker;
marker.and(requirement.marker);
Requirement {
name: requirement.name.clone(),
extras: requirement.extras.clone(),
groups: requirement.groups.clone(),
source: requirement.source.clone(),
origin: requirement.origin.clone(),
marker: marker.simplify_extras(slice::from_ref(&extra)),
)))
} else {
let mut requirements = self
.requirements_for_extra(
dependencies.iter(),
extra,
env,
python_marker,
python_requirement,
)
.collect::<Vec<_>>();

// Transitively process all extras that are recursively included, starting with the current
// extra.
let mut seen = FxHashSet::<(ExtraName, MarkerTree)>::default();
let mut queue: VecDeque<_> = requirements
.iter()
.filter(|req| name == Some(&req.name))
.flat_map(|req| req.extras.iter().cloned().map(|extra| (extra, req.marker)))
.collect();
while let Some((extra, marker)) = queue.pop_front() {
if !seen.insert((extra.clone(), marker)) {
continue;
}
for requirement in self.requirements_for_extra(
dependencies,
Some(&extra),
env,
python_marker,
python_requirement,
) {
let requirement = match requirement {
Cow::Owned(mut requirement) => {
requirement.marker.and(marker);
requirement
}
Cow::Borrowed(requirement) => {
let mut marker = marker;
marker.and(requirement.marker);
Requirement {
name: requirement.name.clone(),
extras: requirement.extras.clone(),
groups: requirement.groups.clone(),
source: requirement.source.clone(),
origin: requirement.origin.clone(),
marker: marker.simplify_extras(slice::from_ref(&extra)),
}
}
};
if name == Some(&requirement.name) {
// Add each transitively included extra.
queue.extend(
requirement
.extras
.iter()
.cloned()
.map(|extra| (extra, requirement.marker)),
);
} else {
// Add the requirements for that extra.
requirements.push(Cow::Owned(requirement));
}
};
if name == Some(&requirement.name) {
// Add each transitively included extra.
queue.extend(
requirement
.extras
.iter()
.cloned()
.map(|extra| (extra, requirement.marker)),
);
} else {
// Add the requirements for that extra.
requirements.push(Cow::Owned(requirement));
}
}
}

// Retain any self-constraints for that extra, e.g., if `project[foo]` includes
// `project[bar]>1.0`, as a dependency, we need to propagate `project>1.0`, in addition to
// transitively expanding `project[bar]`.
let mut self_constraints = vec![];
for req in &requirements {
if name == Some(&req.name) && !req.source.is_empty() {
self_constraints.push(Requirement {
name: req.name.clone(),
extras: vec![],
groups: req.groups.clone(),
source: req.source.clone(),
origin: req.origin.clone(),
marker: req.marker,
});
// Retain any self-constraints for that extra, e.g., if `project[foo]` includes
// `project[bar]>1.0`, as a dependency, we need to propagate `project>1.0`, in addition to
// transitively expanding `project[bar]`.
let mut self_constraints = vec![];
for req in &requirements {
if name == Some(&req.name) && !req.source.is_empty() {
self_constraints.push(Requirement {
name: req.name.clone(),
extras: vec![],
groups: req.groups.clone(),
source: req.source.clone(),
origin: req.origin.clone(),
marker: req.marker,
});
}
}
}

// Drop all the self-requirements now that we flattened them out.
requirements.retain(|req| name != Some(&req.name) || req.extras.is_empty());
requirements.extend(self_constraints.into_iter().map(Cow::Owned));
// Drop all the self-requirements now that we flattened them out.
requirements.retain(|req| name != Some(&req.name) || req.extras.is_empty());
requirements.extend(self_constraints.into_iter().map(Cow::Owned));

requirements
Either::Right(requirements.into_iter())
}
}

/// The set of the regular and dev dependencies, filtered by Python version,
Expand Down

0 comments on commit 62f8a48

Please sign in to comment.