Skip to content

Commit 62f8a48

Browse files
charliermarshcharliermarsh
authored
Remove extra collects when flattening requirements (#10837)
1 parent 47a31fc commit 62f8a48

File tree

2 files changed

+138
-135
lines changed

2 files changed

+138
-135
lines changed

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

Lines changed: 45 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use std::borrow::Cow;
12
use std::iter;
23

34
use either::Either;
@@ -26,7 +27,7 @@ pub(crate) struct PubGrubDependency {
2627
impl PubGrubDependency {
2728
pub(crate) fn from_requirement<'a>(
2829
conflicts: &Conflicts,
29-
requirement: &'a Requirement,
30+
requirement: Cow<'a, Requirement>,
3031
dev: Option<&'a GroupName>,
3132
source_name: Option<&'a PackageName>,
3233
) -> impl Iterator<Item = Self> + 'a {
@@ -89,58 +90,60 @@ impl PubGrubDependency {
8990
};
9091

9192
// Add the package, plus any extra variants.
92-
iter.map(|(extra, group)| PubGrubRequirement::from_requirement(requirement, extra, group))
93-
.map(move |requirement| {
94-
let PubGrubRequirement {
93+
iter.map(move |(extra, group)| {
94+
PubGrubRequirement::from_requirement(&requirement, extra, group)
95+
})
96+
.map(move |requirement| {
97+
let PubGrubRequirement {
98+
package,
99+
version,
100+
url,
101+
} = requirement;
102+
match &*package {
103+
PubGrubPackageInner::Package { .. } => PubGrubDependency {
104+
package,
105+
version,
106+
url,
107+
},
108+
PubGrubPackageInner::Marker { .. } => PubGrubDependency {
95109
package,
96110
version,
97111
url,
98-
} = requirement;
99-
match &*package {
100-
PubGrubPackageInner::Package { .. } => PubGrubDependency {
112+
},
113+
PubGrubPackageInner::Extra { name, .. } => {
114+
// Detect self-dependencies.
115+
if dev.is_none() {
116+
debug_assert!(
117+
source_name.is_none_or(|source_name| source_name != name),
118+
"extras not flattened for {name}"
119+
);
120+
}
121+
PubGrubDependency {
101122
package,
102123
version,
103124
url,
104-
},
105-
PubGrubPackageInner::Marker { .. } => PubGrubDependency {
125+
}
126+
}
127+
PubGrubPackageInner::Dev { name, .. } => {
128+
// Detect self-dependencies.
129+
if dev.is_none() {
130+
debug_assert!(
131+
source_name.is_none_or(|source_name| source_name != name),
132+
"group not flattened for {name}"
133+
);
134+
}
135+
PubGrubDependency {
106136
package,
107137
version,
108138
url,
109-
},
110-
PubGrubPackageInner::Extra { name, .. } => {
111-
// Detect self-dependencies.
112-
if dev.is_none() {
113-
debug_assert!(
114-
source_name.is_none_or(|source_name| source_name != name),
115-
"extras not flattened for {name}"
116-
);
117-
}
118-
PubGrubDependency {
119-
package,
120-
version,
121-
url,
122-
}
123-
}
124-
PubGrubPackageInner::Dev { name, .. } => {
125-
// Detect self-dependencies.
126-
if dev.is_none() {
127-
debug_assert!(
128-
source_name.is_none_or(|source_name| source_name != name),
129-
"group not flattened for {name}"
130-
);
131-
}
132-
PubGrubDependency {
133-
package,
134-
version,
135-
url,
136-
}
137-
}
138-
PubGrubPackageInner::Root(_) => unreachable!("root package in dependencies"),
139-
PubGrubPackageInner::Python(_) => {
140-
unreachable!("python package in dependencies")
141139
}
142140
}
143-
})
141+
PubGrubPackageInner::Root(_) => unreachable!("root package in dependencies"),
142+
PubGrubPackageInner::Python(_) => {
143+
unreachable!("python package in dependencies")
144+
}
145+
}
146+
})
144147
}
145148
}
146149

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

Lines changed: 93 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -1605,7 +1605,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
16051605
);
16061606

16071607
requirements
1608-
.iter()
16091608
.flat_map(|requirement| {
16101609
PubGrubDependency::from_requirement(
16111610
&self.conflicts,
@@ -1704,7 +1703,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
17041703
);
17051704

17061705
requirements
1707-
.iter()
17081706
.flat_map(|requirement| {
17091707
PubGrubDependency::from_requirement(
17101708
&self.conflicts,
@@ -1800,116 +1798,118 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
18001798
name: Option<&PackageName>,
18011799
env: &'a ResolverEnvironment,
18021800
python_requirement: &'a PythonRequirement,
1803-
) -> Vec<Cow<'a, Requirement>> {
1804-
// Start with the requirements for the current extra of the package (for an extra
1805-
// requirement) or the non-extra (regular) dependencies (if extra is None), plus
1806-
// the constraints for the current package.
1807-
let regular_and_dev_dependencies = if let Some(dev) = dev {
1808-
Either::Left(dev_dependencies.get(dev).into_iter().flatten())
1809-
} else {
1810-
Either::Right(dependencies.iter())
1811-
};
1801+
) -> impl Iterator<Item = Cow<'a, Requirement>> {
18121802
let python_marker = python_requirement.to_marker_tree();
1813-
let mut requirements = self
1814-
.requirements_for_extra(
1815-
regular_and_dev_dependencies,
1803+
1804+
if let Some(dev) = dev {
1805+
// Dependency groups can include the project itself, so no need to flatten recursive
1806+
// dependencies.
1807+
Either::Left(Either::Left(self.requirements_for_extra(
1808+
dev_dependencies.get(dev).into_iter().flatten(),
18161809
extra,
18171810
env,
18181811
python_marker,
18191812
python_requirement,
1820-
)
1821-
.collect::<Vec<_>>();
1822-
1823-
// Dependency groups can include the project itself, so no need to flatten recursive
1824-
// dependencies.
1825-
if dev.is_some() {
1826-
return requirements;
1827-
}
1828-
1829-
// Check if there are recursive self inclusions; if so, we need to go into the expensive
1830-
// branch.
1831-
if !requirements
1813+
)))
1814+
} else if !dependencies
18321815
.iter()
18331816
.any(|req| name == Some(&req.name) && !req.extras.is_empty())
18341817
{
1835-
return requirements;
1836-
}
1837-
1838-
// Transitively process all extras that are recursively included, starting with the current
1839-
// extra.
1840-
let mut seen = FxHashSet::<(ExtraName, MarkerTree)>::default();
1841-
let mut queue: VecDeque<_> = requirements
1842-
.iter()
1843-
.filter(|req| name == Some(&req.name))
1844-
.flat_map(|req| req.extras.iter().cloned().map(|extra| (extra, req.marker)))
1845-
.collect();
1846-
while let Some((extra, marker)) = queue.pop_front() {
1847-
if !seen.insert((extra.clone(), marker)) {
1848-
continue;
1849-
}
1850-
for requirement in self.requirements_for_extra(
1851-
dependencies,
1852-
Some(&extra),
1818+
// If the project doesn't define any recursive dependencies, take the fast path.
1819+
Either::Left(Either::Right(self.requirements_for_extra(
1820+
dependencies.iter(),
1821+
extra,
18531822
env,
18541823
python_marker,
18551824
python_requirement,
1856-
) {
1857-
let requirement = match requirement {
1858-
Cow::Owned(mut requirement) => {
1859-
requirement.marker.and(marker);
1860-
requirement
1861-
}
1862-
Cow::Borrowed(requirement) => {
1863-
let mut marker = marker;
1864-
marker.and(requirement.marker);
1865-
Requirement {
1866-
name: requirement.name.clone(),
1867-
extras: requirement.extras.clone(),
1868-
groups: requirement.groups.clone(),
1869-
source: requirement.source.clone(),
1870-
origin: requirement.origin.clone(),
1871-
marker: marker.simplify_extras(slice::from_ref(&extra)),
1825+
)))
1826+
} else {
1827+
let mut requirements = self
1828+
.requirements_for_extra(
1829+
dependencies.iter(),
1830+
extra,
1831+
env,
1832+
python_marker,
1833+
python_requirement,
1834+
)
1835+
.collect::<Vec<_>>();
1836+
1837+
// Transitively process all extras that are recursively included, starting with the current
1838+
// extra.
1839+
let mut seen = FxHashSet::<(ExtraName, MarkerTree)>::default();
1840+
let mut queue: VecDeque<_> = requirements
1841+
.iter()
1842+
.filter(|req| name == Some(&req.name))
1843+
.flat_map(|req| req.extras.iter().cloned().map(|extra| (extra, req.marker)))
1844+
.collect();
1845+
while let Some((extra, marker)) = queue.pop_front() {
1846+
if !seen.insert((extra.clone(), marker)) {
1847+
continue;
1848+
}
1849+
for requirement in self.requirements_for_extra(
1850+
dependencies,
1851+
Some(&extra),
1852+
env,
1853+
python_marker,
1854+
python_requirement,
1855+
) {
1856+
let requirement = match requirement {
1857+
Cow::Owned(mut requirement) => {
1858+
requirement.marker.and(marker);
1859+
requirement
18721860
}
1861+
Cow::Borrowed(requirement) => {
1862+
let mut marker = marker;
1863+
marker.and(requirement.marker);
1864+
Requirement {
1865+
name: requirement.name.clone(),
1866+
extras: requirement.extras.clone(),
1867+
groups: requirement.groups.clone(),
1868+
source: requirement.source.clone(),
1869+
origin: requirement.origin.clone(),
1870+
marker: marker.simplify_extras(slice::from_ref(&extra)),
1871+
}
1872+
}
1873+
};
1874+
if name == Some(&requirement.name) {
1875+
// Add each transitively included extra.
1876+
queue.extend(
1877+
requirement
1878+
.extras
1879+
.iter()
1880+
.cloned()
1881+
.map(|extra| (extra, requirement.marker)),
1882+
);
1883+
} else {
1884+
// Add the requirements for that extra.
1885+
requirements.push(Cow::Owned(requirement));
18731886
}
1874-
};
1875-
if name == Some(&requirement.name) {
1876-
// Add each transitively included extra.
1877-
queue.extend(
1878-
requirement
1879-
.extras
1880-
.iter()
1881-
.cloned()
1882-
.map(|extra| (extra, requirement.marker)),
1883-
);
1884-
} else {
1885-
// Add the requirements for that extra.
1886-
requirements.push(Cow::Owned(requirement));
18871887
}
18881888
}
1889-
}
18901889

1891-
// Retain any self-constraints for that extra, e.g., if `project[foo]` includes
1892-
// `project[bar]>1.0`, as a dependency, we need to propagate `project>1.0`, in addition to
1893-
// transitively expanding `project[bar]`.
1894-
let mut self_constraints = vec![];
1895-
for req in &requirements {
1896-
if name == Some(&req.name) && !req.source.is_empty() {
1897-
self_constraints.push(Requirement {
1898-
name: req.name.clone(),
1899-
extras: vec![],
1900-
groups: req.groups.clone(),
1901-
source: req.source.clone(),
1902-
origin: req.origin.clone(),
1903-
marker: req.marker,
1904-
});
1890+
// Retain any self-constraints for that extra, e.g., if `project[foo]` includes
1891+
// `project[bar]>1.0`, as a dependency, we need to propagate `project>1.0`, in addition to
1892+
// transitively expanding `project[bar]`.
1893+
let mut self_constraints = vec![];
1894+
for req in &requirements {
1895+
if name == Some(&req.name) && !req.source.is_empty() {
1896+
self_constraints.push(Requirement {
1897+
name: req.name.clone(),
1898+
extras: vec![],
1899+
groups: req.groups.clone(),
1900+
source: req.source.clone(),
1901+
origin: req.origin.clone(),
1902+
marker: req.marker,
1903+
});
1904+
}
19051905
}
1906-
}
19071906

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

1912-
requirements
1911+
Either::Right(requirements.into_iter())
1912+
}
19131913
}
19141914

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

0 commit comments

Comments
 (0)