From e6b40fc2e553b0d6c3a645218d16348bbbc21c18 Mon Sep 17 00:00:00 2001
From: konstin <konstin@mailbox.org>
Date: Tue, 3 Sep 2024 09:37:05 +0200
Subject: [PATCH] Move unreachable wheels check

Prep for fixing #6512. No functional changes.
---
 crates/uv-resolver/src/lock/mod.rs | 66 ++++++++++++++++--------------
 1 file changed, 36 insertions(+), 30 deletions(-)

diff --git a/crates/uv-resolver/src/lock/mod.rs b/crates/uv-resolver/src/lock/mod.rs
index 887e90772ef04..fcd4ae99103e9 100644
--- a/crates/uv-resolver/src/lock/mod.rs
+++ b/crates/uv-resolver/src/lock/mod.rs
@@ -83,35 +83,40 @@ impl Lock {
     /// Initialize a [`Lock`] from a [`ResolutionGraph`].
     pub fn from_resolution_graph(graph: &ResolutionGraph, root: &Path) -> Result<Self, LockError> {
         let mut locked_dists = BTreeMap::new();
+        let requires_python = graph.requires_python.clone();
 
         // Lock all base packages.
         for node_index in graph.petgraph.node_indices() {
             let ResolutionGraphNode::Dist(dist) = &graph.petgraph[node_index] else {
                 continue;
             };
-            if dist.is_base() {
-                let fork_markers = graph
-                    .fork_markers(dist.name(), &dist.version, dist.dist.version_or_url().url())
-                    .cloned()
-                    .unwrap_or_default();
-                let mut locked_dist = Package::from_annotated_dist(dist, fork_markers, root)?;
+            if !dist.is_base() {
+                continue;
+            }
+            let fork_markers = graph
+                .fork_markers(dist.name(), &dist.version, dist.dist.version_or_url().url())
+                .cloned()
+                .unwrap_or_default();
+            let mut package = Package::from_annotated_dist(dist, fork_markers, root)?;
 
-                // Add all dependencies
-                for edge in graph.petgraph.edges(node_index) {
-                    let ResolutionGraphNode::Dist(dependency_dist) = &graph.petgraph[edge.target()]
-                    else {
-                        continue;
-                    };
-                    let marker = edge.weight().clone();
-                    locked_dist.add_dependency(dependency_dist, marker, root)?;
-                }
-                let id = locked_dist.id.clone();
-                if let Some(locked_dist) = locked_dists.insert(id, locked_dist) {
-                    return Err(LockErrorKind::DuplicatePackage {
-                        id: locked_dist.id.clone(),
-                    }
-                    .into());
+            Self::remove_unreachable_wheels(&requires_python, &mut package);
+
+            // Add all dependencies
+            for edge in graph.petgraph.edges(node_index) {
+                let ResolutionGraphNode::Dist(dependency_dist) = &graph.petgraph[edge.target()]
+                else {
+                    continue;
+                };
+                let marker = edge.weight().clone();
+                package.add_dependency(dependency_dist, marker, root)?;
+            }
+
+            let id = package.id.clone();
+            if let Some(locked_dist) = locked_dists.insert(id, package) {
+                return Err(LockErrorKind::DuplicatePackage {
+                    id: locked_dist.id.clone(),
                 }
+                .into());
             }
         }
 
@@ -164,7 +169,6 @@ impl Lock {
         }
 
         let packages = locked_dists.into_values().collect();
-        let requires_python = graph.requires_python.clone();
         let options = ResolverOptions {
             resolution_mode: graph.options.resolution_mode,
             prerelease_mode: graph.options.prerelease_mode,
@@ -182,6 +186,16 @@ impl Lock {
         Ok(lock)
     }
 
+    fn remove_unreachable_wheels(requires_python: &Option<RequiresPython>, package: &mut Package) {
+        // Remove wheels that don't match `requires-python` and can't be selected for
+        // installation.
+        if let Some(requires_python) = &requires_python {
+            package
+                .wheels
+                .retain(|wheel| requires_python.matches_wheel_tag(&wheel.filename));
+        }
+    }
+
     /// Initialize a [`Lock`] from a list of [`Package`] entries.
     fn new(
         version: u32,
@@ -238,14 +252,6 @@ impl Lock {
                     }
                 }
             }
-
-            // Remove wheels that don't match `requires-python` and can't be selected for
-            // installation.
-            if let Some(requires_python) = &requires_python {
-                package
-                    .wheels
-                    .retain(|wheel| requires_python.matches_wheel_tag(&wheel.filename));
-            }
         }
         packages.sort_by(|dist1, dist2| dist1.id.cmp(&dist2.id));