Skip to content

Commit 6f44ff3

Browse files
charliermarshcharliermarsh
committed
Remove extra collects when flattening requirements
1 parent d952a97 commit 6f44ff3

File tree

2 files changed

+135
-135
lines changed

2 files changed

+135
-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: 90 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,115 @@ 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+
Either::Left(Either::Left(self.requirements_for_extra(
1806+
dev_dependencies.get(dev).into_iter().flatten(),
18161807
extra,
18171808
env,
18181809
python_marker,
18191810
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
1811+
)))
1812+
} else if !dependencies
18321813
.iter()
18331814
.any(|req| name == Some(&req.name) && !req.extras.is_empty())
18341815
{
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),
1816+
Either::Left(Either::Right(self.requirements_for_extra(
1817+
dependencies.iter(),
1818+
extra,
18531819
env,
18541820
python_marker,
18551821
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)),
1822+
)))
1823+
} else {
1824+
let mut requirements = self
1825+
.requirements_for_extra(
1826+
dependencies.iter(),
1827+
extra,
1828+
env,
1829+
python_marker,
1830+
python_requirement,
1831+
)
1832+
.collect::<Vec<_>>();
1833+
1834+
// Transitively process all extras that are recursively included, starting with the current
1835+
// extra.
1836+
let mut seen = FxHashSet::<(ExtraName, MarkerTree)>::default();
1837+
let mut queue: VecDeque<_> = requirements
1838+
.iter()
1839+
.filter(|req| name == Some(&req.name))
1840+
.flat_map(|req| req.extras.iter().cloned().map(|extra| (extra, req.marker)))
1841+
.collect();
1842+
while let Some((extra, marker)) = queue.pop_front() {
1843+
if !seen.insert((extra.clone(), marker)) {
1844+
continue;
1845+
}
1846+
for requirement in self.requirements_for_extra(
1847+
dependencies,
1848+
Some(&extra),
1849+
env,
1850+
python_marker,
1851+
python_requirement,
1852+
) {
1853+
let requirement = match requirement {
1854+
Cow::Owned(mut requirement) => {
1855+
requirement.marker.and(marker);
1856+
requirement
18721857
}
1858+
Cow::Borrowed(requirement) => {
1859+
let mut marker = marker;
1860+
marker.and(requirement.marker);
1861+
Requirement {
1862+
name: requirement.name.clone(),
1863+
extras: requirement.extras.clone(),
1864+
groups: requirement.groups.clone(),
1865+
source: requirement.source.clone(),
1866+
origin: requirement.origin.clone(),
1867+
marker: marker.simplify_extras(slice::from_ref(&extra)),
1868+
}
1869+
}
1870+
};
1871+
if name == Some(&requirement.name) {
1872+
// Add each transitively included extra.
1873+
queue.extend(
1874+
requirement
1875+
.extras
1876+
.iter()
1877+
.cloned()
1878+
.map(|extra| (extra, requirement.marker)),
1879+
);
1880+
} else {
1881+
// Add the requirements for that extra.
1882+
requirements.push(Cow::Owned(requirement));
18731883
}
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));
18871884
}
18881885
}
1889-
}
18901886

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

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));
1904+
// Drop all the self-requirements now that we flattened them out.
1905+
requirements.retain(|req| name != Some(&req.name) || req.extras.is_empty());
1906+
requirements.extend(self_constraints.into_iter().map(Cow::Owned));
19111907

1912-
requirements
1908+
Either::Right(requirements.into_iter())
1909+
}
19131910
}
19141911

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

0 commit comments

Comments
 (0)