Flatten requirements eagerly in get_dependencies (#4430)

Downstack PR: #4515 Upstack PR: #4481

Consider these two cases:

A:
```
werkzeug==2.0.0
werkzeug @ 960bb4017c/Werkzeug-2.0.0-py3-none-any.whl
```

B:
```toml
dependencies = [
  "iniconfig == 1.1.1 ; python_version < '3.12'",
  "iniconfig @ git+https://github.com/pytest-dev/iniconfig@93f5930e668c0d1ddf4597e38dd0dea4e2665e7a ; python_version >= '3.12'",
]
```

In the first case, `werkzeug==2.0.0` should be overridden by the url. In
the second case `iniconfig == 1.1.1` is in a different fork and must
remain a registry distribution.

That means the conversion from `Requirement` to `PubGrubPackage` is
dependent on the other requirements of the package. We can either look
into the other packages immediately, or we can move the forking before
the conversion to `PubGrubDependencies` instead of after. Either version
requires a flat list of `Requirement`s to use. This refactoring gives us
this list.

I'll add support for both of the above cases in the forking urls branch
before merging this PR. I also have to move constraints over to this.

crates/uv-configuration/src/constraints.rs

@@ -16,7 +16,11 @@ impl Constraints {
             constraints
                 .entry(requirement.name.clone())
                 .or_default()
-                .push(requirement);
+                .push(Requirement {
+                    // We add and apply constraints independent of their extras.
+                    extras: vec![],
+                    ..requirement
+                });
         }
         Self(constraints)
     }

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

@@ -1,26 +1,19 @@
-use std::collections::BTreeMap;
-
-use either::Either;
 use itertools::Itertools;
 use pubgrub::range::Range;
-use rustc_hash::FxHashSet;
+use tracing::warn;
-use tracing::{trace, warn};
 
 use distribution_types::Verbatim;
 use pep440_rs::Version;
-use pep508_rs::{MarkerEnvironment, MarkerTree};
 use pypi_types::{
     ParsedArchiveUrl, ParsedDirectoryUrl, ParsedGitUrl, ParsedPathUrl, ParsedUrl, Requirement,
     RequirementSource,
 };
-use uv_configuration::{Constraints, Overrides};
 use uv_git::GitResolver;
-use uv_normalize::{ExtraName, GroupName, PackageName};
+use uv_normalize::{ExtraName, PackageName};
 
-use crate::pubgrub::specifier::PubGrubSpecifier;
 use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner};
 use crate::resolver::{Locals, Urls};
-use crate::ResolveError;
+use crate::{PubGrubSpecifier, ResolveError};
 
 #[derive(Debug)]
 pub struct PubGrubDependencies(Vec<(PubGrubPackage, Range<Version>)>);
@@ -29,41 +22,51 @@ impl PubGrubDependencies {
     /// Generate a set of PubGrub dependencies from a set of requirements.
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn from_requirements(
-        requirements: &[Requirement],
+        flattened_requirements: &[&Requirement],
-        dev_dependencies: &BTreeMap<GroupName, Vec<Requirement>>,
-        constraints: &Constraints,
-        overrides: &Overrides,
         source_name: Option<&PackageName>,
-        source_extra: Option<&ExtraName>,
-        source_dev: Option<&GroupName>,
         urls: &Urls,
         locals: &Locals,
         git: &GitResolver,
-        env: Option<&MarkerEnvironment>,
-        requires_python: Option<&MarkerTree>,
-        fork_markers: &MarkerTree,
     ) -> Result<Self, ResolveError> {
-        let mut dependencies = Vec::default();
+        let mut dependencies = Vec::new();
-        let mut seen = FxHashSet::default();
+        for requirement in flattened_requirements {
+            // Add the package, plus any extra variants.
+            for result in std::iter::once(PubGrubRequirement::from_requirement(
+                requirement,
+                None,
+                urls,
+                locals,
+                git,
+            ))
+            .chain(requirement.extras.clone().into_iter().map(|extra| {
+                PubGrubRequirement::from_requirement(requirement, Some(extra), urls, locals, git)
+            })) {
+                let PubGrubRequirement { package, version } = result?;
 
-        add_requirements(
+                match &*package {
-            requirements,
+                    PubGrubPackageInner::Package { name, .. } => {
-            dev_dependencies,
+                        // Detect self-dependencies.
-            constraints,
+                        if source_name.is_some_and(|source_name| source_name == name) {
-            overrides,
+                            warn!("{name} has a dependency on itself");
-            source_name,
+                            continue;
-            source_extra,
+                        }
-            source_dev,
-            urls,
-            locals,
-            git,
-            env,
-            requires_python,
-            fork_markers,
-            &mut dependencies,
-            &mut seen,
-        )?;
 
+                        dependencies.push((package.clone(), version.clone()));
+                    }
+                    PubGrubPackageInner::Marker { .. } => {
+                        dependencies.push((package.clone(), version.clone()));
+                    }
+                    PubGrubPackageInner::Extra { name, .. } => {
+                        debug_assert!(
+                            !source_name.is_some_and(|source_name| source_name == name),
+                            "extras not flattened for {name}"
+                        );
+                        dependencies.push((package.clone(), version.clone()));
+                    }
+                    _ => {}
+                }
+            }
+        }
         Ok(Self(dependencies))
     }
 
@@ -78,183 +81,6 @@ impl PubGrubDependencies {
     }
 }
 
-/// Add a set of requirements to a list of dependencies.
-#[allow(clippy::too_many_arguments)]
-fn add_requirements(
-    requirements: &[Requirement],
-    dev_dependencies: &BTreeMap<GroupName, Vec<Requirement>>,
-    constraints: &Constraints,
-    overrides: &Overrides,
-    source_name: Option<&PackageName>,
-    source_extra: Option<&ExtraName>,
-    source_dev: Option<&GroupName>,
-    urls: &Urls,
-    locals: &Locals,
-    git: &GitResolver,
-    env: Option<&MarkerEnvironment>,
-    requires_python: Option<&MarkerTree>,
-    fork_markers: &MarkerTree,
-    dependencies: &mut Vec<(PubGrubPackage, Range<Version>)>,
-    seen: &mut FxHashSet<ExtraName>,
-) -> Result<(), ResolveError> {
-    // Iterate over all declared requirements.
-    for requirement in overrides.apply(if let Some(source_dev) = source_dev {
-        Either::Left(dev_dependencies.get(source_dev).into_iter().flatten())
-    } else {
-        Either::Right(requirements.iter())
-    }) {
-        // If the requirement would not be selected with any Python version
-        // supported by the root, skip it.
-        if !satisfies_requires_python(requires_python, requirement) {
-            trace!(
-                "skipping {requirement} because of Requires-Python {requires_python}",
-                // OK because this filter only applies when there is a present
-                // Requires-Python specifier.
-                requires_python = requires_python.unwrap()
-            );
-            continue;
-        }
-        // If we're in universal mode, `fork_markers` might correspond to a
-        // non-trivial marker expression that provoked the resolver to fork.
-        // In that case, we should ignore any dependency that cannot possibly
-        // satisfy the markers that provoked the fork.
-        if !possible_to_satisfy_markers(fork_markers, requirement) {
-            trace!("skipping {requirement} because of context resolver markers {fork_markers}");
-            continue;
-        }
-        // If the requirement isn't relevant for the current platform, skip it.
-        match source_extra {
-            Some(source_extra) => {
-                // Only include requirements that are relevant for the current extra.
-                if requirement.evaluate_markers(env, &[]) {
-                    continue;
-                }
-                if !requirement.evaluate_markers(env, std::slice::from_ref(source_extra)) {
-                    continue;
-                }
-            }
-            None => {
-                if !requirement.evaluate_markers(env, &[]) {
-                    continue;
-                }
-            }
-        }
-
-        // Add the package, plus any extra variants.
-        for result in std::iter::once(PubGrubRequirement::from_requirement(
-            requirement,
-            None,
-            urls,
-            locals,
-            git,
-        ))
-        .chain(requirement.extras.clone().into_iter().map(|extra| {
-            PubGrubRequirement::from_requirement(requirement, Some(extra), urls, locals, git)
-        })) {
-            let PubGrubRequirement { package, version } = result?;
-
-            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");
-                        continue;
-                    }
-
-                    dependencies.push((package.clone(), version.clone()));
-                }
-                PubGrubPackageInner::Marker { .. } => {
-                    dependencies.push((package.clone(), version.clone()));
-                }
-                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) {
-                        if seen.insert(extra.clone()) {
-                            add_requirements(
-                                requirements,
-                                dev_dependencies,
-                                constraints,
-                                overrides,
-                                source_name,
-                                Some(extra),
-                                None,
-                                urls,
-                                locals,
-                                git,
-                                env,
-                                requires_python,
-                                fork_markers,
-                                dependencies,
-                                seen,
-                            )?;
-                        }
-                    } else {
-                        dependencies.push((package.clone(), version.clone()));
-                    }
-                }
-                _ => {}
-            }
-
-            // If the requirement was constrained, add those constraints.
-            for constraint in constraints.get(&requirement.name).into_iter().flatten() {
-                // If the requirement would not be selected with any Python
-                // version supported by the root, skip it.
-                if !satisfies_requires_python(requires_python, constraint) {
-                    trace!(
-                        "skipping {requirement} because of Requires-Python {requires_python}",
-                        // OK because this filter only applies when there is a present
-                        // Requires-Python specifier.
-                        requires_python = requires_python.unwrap()
-                    );
-                    continue;
-                }
-                // If we're in universal mode, `fork_markers` might correspond
-                // to a non-trivial marker expression that provoked the
-                // resolver to fork. In that case, we should ignore any
-                // dependency that cannot possibly satisfy the markers that
-                // provoked the fork.
-                if !possible_to_satisfy_markers(fork_markers, constraint) {
-                    trace!(
-                        "skipping {requirement} because of context resolver markers {fork_markers}"
-                    );
-                    continue;
-                }
-                // If the requirement isn't relevant for the current platform, skip it.
-                match source_extra {
-                    Some(source_extra) => {
-                        if !constraint.evaluate_markers(env, std::slice::from_ref(source_extra)) {
-                            continue;
-                        }
-                    }
-                    None => {
-                        if !constraint.evaluate_markers(env, &[]) {
-                            continue;
-                        }
-                    }
-                }
-
-                // Add the package.
-                let PubGrubRequirement { package, version } =
-                    PubGrubRequirement::from_constraint(constraint, urls, locals, git)?;
-
-                // Ignore self-dependencies.
-                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");
-                        continue;
-                    }
-                }
-
-                dependencies.push((package.clone(), version.clone()));
-            }
-        }
-    }
-
-    Ok(())
-}
-
 /// Convert a [`PubGrubDependencies`] to a [`DependencyConstraints`].
 impl From<PubGrubDependencies> for Vec<(PubGrubPackage, Range<Version>)> {
     fn from(dependencies: PubGrubDependencies) -> Self {
@@ -455,39 +281,4 @@ impl PubGrubRequirement {
             }
         }
     }
-
-    /// Convert a constraint to a PubGrub-compatible package and range.
-    pub(crate) fn from_constraint(
-        constraint: &Requirement,
-        urls: &Urls,
-        locals: &Locals,
-        git: &GitResolver,
-    ) -> Result<Self, ResolveError> {
-        Self::from_requirement(constraint, None, urls, locals, git)
-    }
-}
-
-/// Returns true if and only if the given requirement's marker expression has a
-/// possible true value given the `requires_python` specifier given.
-///
-/// While this is always called, a `requires_python` is only non-None when in
-/// universal resolution mode. In non-universal mode, `requires_python` is
-/// `None` and this always returns `true`.
-fn satisfies_requires_python(
-    requires_python: Option<&MarkerTree>,
-    requirement: &Requirement,
-) -> bool {
-    let Some(requires_python) = requires_python else {
-        return true;
-    };
-    possible_to_satisfy_markers(requires_python, requirement)
-}
-
-/// Returns true if and only if the given requirement's marker expression has a
-/// possible true value given the `markers` expression given.
-fn possible_to_satisfy_markers(markers: &MarkerTree, requirement: &Requirement) -> bool {
-    let Some(marker) = requirement.marker.as_ref() else {
-        return true;
-    };
-    !crate::marker::is_disjoint(markers, marker)
 }

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

@@ -1,13 +1,14 @@
 //! Given a set of requirements, find a set of compatible packages.
 
 use std::borrow::Cow;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, VecDeque};
 use std::fmt::{Display, Formatter};
 use std::ops::Bound;
 use std::sync::Arc;
-use std::thread;
+use std::{iter, thread};
 
 use dashmap::DashMap;
+use either::Either;
 use futures::{FutureExt, StreamExt, TryFutureExt};
 use itertools::Itertools;
 use pubgrub::error::PubGrubError;
@@ -938,27 +939,26 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
     ) -> Result<Dependencies, ResolveError> {
         let (dependencies, name) = match &**package {
             PubGrubPackageInner::Root(_) => {
-                // Add the root requirements.
+                let no_dev_deps = BTreeMap::default();
-                let dependencies = PubGrubDependencies::from_requirements(
+                let requirements = self.flatten_requirements(
                     &self.requirements,
-                    &BTreeMap::default(),
+                    &no_dev_deps,
-                    &self.constraints,
-                    &self.overrides,
                     None,
                     None,
                     None,
+                    markers,
+                );
+
+                let dependencies = PubGrubDependencies::from_requirements(
+                    &requirements,
+                    None,
                     &self.urls,
                     &self.locals,
                     &self.git,
-                    self.markers.as_ref(),
-                    self.requires_python.as_ref(),
-                    markers,
                 )?;
+
                 (dependencies, None)
             }
-
-            PubGrubPackageInner::Python(_) => return Ok(Dependencies::Available(Vec::default())),
-
             PubGrubPackageInner::Package {
                 name,
                 extra,
@@ -1072,43 +1072,45 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     }
                 };
 
-                let mut dependencies = PubGrubDependencies::from_requirements(
+                let requirements = self.flatten_requirements(
                     &metadata.requires_dist,
                     &metadata.dev_dependencies,
-                    &self.constraints,
-                    &self.overrides,
-                    Some(name),
                     extra.as_ref(),
                     dev.as_ref(),
+                    Some(name),
+                    markers,
+                );
+
+                let mut dependencies = PubGrubDependencies::from_requirements(
+                    &requirements,
+                    Some(name),
                     &self.urls,
                     &self.locals,
                     &self.git,
-                    self.markers.as_ref(),
-                    self.requires_python.as_ref(),
-                    markers,
                 )?;
-
                 // If a package has metadata for an enabled dependency group,
                 // add a dependency from it to the same package with the group
                 // enabled.
                 if extra.is_none() && dev.is_none() {
-                    for dev in &self.dev {
+                    for group in &self.dev {
-                        if metadata.dev_dependencies.contains_key(dev) {
+                        if !metadata.dev_dependencies.contains_key(group) {
-                            dependencies.push(
+                            continue;
-                                PubGrubPackage::from(PubGrubPackageInner::Dev {
-                                    name: name.clone(),
-                                    dev: dev.clone(),
-                                    marker: marker.clone(),
-                                    url: url.clone(),
-                                }),
-                                Range::singleton(version.clone()),
-                            );
                         }
+                        dependencies.push(
+                            PubGrubPackage::from(PubGrubPackageInner::Dev {
+                                name: name.clone(),
+                                dev: group.clone(),
+                                marker: marker.clone(),
+                                url: url.clone(),
+                            }),
+                            Range::singleton(version.clone()),
+                        );
                     }
                 }
 
                 (dependencies, Some(name))
             }
+            PubGrubPackageInner::Python(_) => return Ok(Dependencies::Available(Vec::default())),
 
             // Add a dependency on both the marker and base package.
             PubGrubPackageInner::Marker { name, marker, url } => {
@@ -1195,6 +1197,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
             if let Some(name) = name {
                 debug!("Adding transitive dependency for {name}=={version}: {dep_package}{dep_version}");
             } else {
+                // A dependency from the root package or requirements.txt.
                 debug!("Adding direct dependency: {dep_package}{dep_version}");
             }
 
@@ -1208,6 +1211,130 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
         Ok(Dependencies::Available(dependencies.into()))
     }
 
+    /// The regular and dev dependencies filtered by Python version and the markers of this fork,
+    /// plus the extras dependencies of the current package (e.g., `black` depending on
+    /// `black[colorama]`).
+    fn flatten_requirements<'a>(
+        &'a self,
+        dependencies: &'a [Requirement],
+        dev_dependencies: &'a BTreeMap<GroupName, Vec<Requirement>>,
+        extra: Option<&'a ExtraName>,
+        dev: Option<&'a GroupName>,
+        name: Option<&PackageName>,
+        markers: &'a MarkerTree,
+    ) -> Vec<&'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())
+        };
+        let mut requirements: Vec<&Requirement> = self
+            .requirements_for_extra(regular_and_dev_dependencies, extra, markers)
+            .collect();
+
+        // Check if there are recursive self inclusions and we need to go into the expensive branch.
+        if !requirements
+            .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::default();
+        let mut queue: VecDeque<_> = requirements
+            .iter()
+            .filter(|req| name == Some(&req.name))
+            .flat_map(|req| &req.extras)
+            .collect();
+        while let Some(extra) = queue.pop_front() {
+            if !seen.insert(extra) {
+                continue;
+            }
+            // Add each transitively included extra.
+            queue.extend(
+                self.requirements_for_extra(dependencies, Some(extra), markers)
+                    .filter(|req| name == Some(&req.name))
+                    .flat_map(|req| &req.extras),
+            );
+            // Add the requirements for that extra
+            requirements.extend(self.requirements_for_extra(dependencies, Some(extra), markers));
+        }
+        // Drop all the self-requirements now that we flattened them out.
+        requirements.retain(|req| name != Some(&req.name));
+
+        requirements
+    }
+
+    /// The set of the regular and dev dependencies, filtered by Python version,
+    /// the markers of this fork and the requested extra.
+    fn requirements_for_extra<'a>(
+        &'a self,
+        dependencies: impl IntoIterator<Item = &'a Requirement>,
+        extra: Option<&'a ExtraName>,
+        markers: &'a MarkerTree,
+    ) -> impl Iterator<Item = &'a Requirement> {
+        self.overrides
+            .apply(dependencies)
+            .flat_map(|requirement| {
+                iter::once(requirement).chain(
+                    // If the requirement was constrained, add those constraints.
+                    self.constraints
+                        .get(&requirement.name)
+                        .into_iter()
+                        .flatten(),
+                )
+            })
+            .filter(move |requirement| {
+                // If the requirement would not be selected with any Python version
+                // supported by the root, skip it.
+                if !satisfies_requires_python(self.requires_python.as_ref(), requirement) {
+                    trace!(
+                        "skipping {requirement} because of Requires-Python {requires_python}",
+                        // OK because this filter only applies when there is a present
+                        // Requires-Python specifier.
+                        requires_python = self.requires_python.as_ref().unwrap()
+                    );
+                    return false;
+                }
+
+                // If we're in universal mode, `fork_markers` might correspond to a
+                // non-trivial marker expression that provoked the resolver to fork.
+                // In that case, we should ignore any dependency that cannot possibly
+                // satisfy the markers that provoked the fork.
+                if !possible_to_satisfy_markers(markers, requirement) {
+                    trace!("skipping {requirement} because of context resolver markers {markers}");
+                    return false;
+                }
+
+                // If the requirement isn't relevant for the current platform, skip it.
+                match extra {
+                    Some(source_extra) => {
+                        // Only include requirements that are relevant for the current extra.
+                        if requirement.evaluate_markers(self.markers.as_ref(), &[]) {
+                            return false;
+                        }
+                        if !requirement.evaluate_markers(
+                            self.markers.as_ref(),
+                            std::slice::from_ref(source_extra),
+                        ) {
+                            return false;
+                        }
+                    }
+                    None => {
+                        if !requirement.evaluate_markers(self.markers.as_ref(), &[]) {
+                            return false;
+                        }
+                    }
+                }
+                true
+            })
+    }
+
     /// Fetch the metadata for a stream of packages and versions.
     async fn fetch<Provider: ResolverProvider>(
         self: Arc<Self>,
@@ -2290,3 +2417,28 @@ impl<'a> PossibleFork<'a> {
         false
     }
 }
+
+/// Returns true if and only if the given requirement's marker expression has a
+/// possible true value given the `requires_python` specifier given.
+///
+/// While this is always called, a `requires_python` is only non-None when in
+/// universal resolution mode. In non-universal mode, `requires_python` is
+/// `None` and this always returns `true`.
+fn satisfies_requires_python(
+    requires_python: Option<&MarkerTree>,
+    requirement: &Requirement,
+) -> bool {
+    let Some(requires_python) = requires_python else {
+        return true;
+    };
+    possible_to_satisfy_markers(requires_python, requirement)
+}
+
+/// Returns true if and only if the given requirement's marker expression has a
+/// possible true value given the `markers` expression given.
+fn possible_to_satisfy_markers(markers: &MarkerTree, requirement: &Requirement) -> bool {
+    let Some(marker) = requirement.marker.as_ref() else {
+        return true;
+    };
+    !crate::marker::is_disjoint(markers, marker)
+}