From 57b9372baca922058587f406a06e75f2db67833f Mon Sep 17 00:00:00 2001
From: Jacob Finkelman <jfinkelm@amazon.com>
Date: Wed, 25 Oct 2023 14:25:06 +0000
Subject: [PATCH] move yank handling up

---
 src/cargo/sources/registry/index.rs | 32 +++---------------
 src/cargo/sources/registry/mod.rs   | 50 +++++++++++++----------------
 2 files changed, 27 insertions(+), 55 deletions(-)

diff --git a/src/cargo/sources/registry/index.rs b/src/cargo/sources/registry/index.rs
index c235ecc46cf..35b59c01ee9 100644
--- a/src/cargo/sources/registry/index.rs
+++ b/src/cargo/sources/registry/index.rs
@@ -99,7 +99,7 @@ use semver::Version;
 use serde::Deserialize;
 use std::borrow::Cow;
 use std::collections::BTreeMap;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::fs;
 use std::io::ErrorKind;
 use std::path::Path;
@@ -573,8 +573,7 @@ impl<'cfg> RegistryIndex<'cfg> {
         name: InternedString,
         req: &OptVersionReq,
         load: &mut dyn RegistryData,
-        yanked_whitelist: &HashSet<PackageId>,
-        f: &mut dyn FnMut(Summary),
+        f: &mut dyn FnMut(IndexSummary),
     ) -> Poll<CargoResult<()>> {
         if self.config.offline() {
             // This should only return `Poll::Ready(Ok(()))` if there is at least 1 match.
@@ -591,31 +590,15 @@ impl<'cfg> RegistryIndex<'cfg> {
             let callback = &mut |s: IndexSummary| {
                 if !s.is_offline() {
                     called = true;
-                    f(s.into_summary());
+                    f(s);
                 }
             };
-            ready!(self.query_inner_with_online(
-                name,
-                req,
-                load,
-                yanked_whitelist,
-                callback,
-                false
-            )?);
+            ready!(self.query_inner_with_online(name, req, load, callback, false)?);
             if called {
                 return Poll::Ready(Ok(()));
             }
         }
-        self.query_inner_with_online(
-            name,
-            req,
-            load,
-            yanked_whitelist,
-            &mut |s| {
-                f(s.into_summary());
-            },
-            true,
-        )
+        self.query_inner_with_online(name, req, load, f, true)
     }
 
     /// Inner implementation of [`Self::query_inner`]. Returns the number of
@@ -627,7 +610,6 @@ impl<'cfg> RegistryIndex<'cfg> {
         name: InternedString,
         req: &OptVersionReq,
         load: &mut dyn RegistryData,
-        yanked_whitelist: &HashSet<PackageId>,
         f: &mut dyn FnMut(IndexSummary),
         online: bool,
     ) -> Poll<CargoResult<()>> {
@@ -649,10 +631,6 @@ impl<'cfg> RegistryIndex<'cfg> {
                     IndexSummary::Offline(s.as_summary().clone())
                 }
             })
-            // Next filter out all yanked packages. Some yanked packages may
-            // leak through if they're in a whitelist (aka if they were
-            // previously in `Cargo.lock`
-            .filter(|s| !s.is_yanked() || yanked_whitelist.contains(&s.package_id()))
             .for_each(f);
         Poll::Ready(Ok(()))
     }
diff --git a/src/cargo/sources/registry/mod.rs b/src/cargo/sources/registry/mod.rs
index 273f6ffb5c4..7ee461edd80 100644
--- a/src/cargo/sources/registry/mod.rs
+++ b/src/cargo/sources/registry/mod.rs
@@ -730,18 +730,15 @@ impl<'cfg> Source for RegistrySource<'cfg> {
         if kind == QueryKind::Exact && req.is_locked() && !self.ops.is_updated() {
             debug!("attempting query without update");
             let mut called = false;
-            ready!(self.index.query_inner(
-                dep.package_name(),
-                &req,
-                &mut *self.ops,
-                &self.yanked_whitelist,
-                &mut |s| {
-                    if dep.matches(&s) {
+            ready!(self
+                .index
+                .query_inner(dep.package_name(), &req, &mut *self.ops, &mut |s| {
+                    if dep.matches(s.as_summary()) {
+                        // We are looking for a package from a lock file so we do not care about yank
                         called = true;
-                        f(s);
+                        f(s.into_summary());
                     }
-                },
-            ))?;
+                },))?;
             if called {
                 Poll::Ready(Ok(()))
             } else {
@@ -751,22 +748,23 @@ impl<'cfg> Source for RegistrySource<'cfg> {
             }
         } else {
             let mut called = false;
-            ready!(self.index.query_inner(
-                dep.package_name(),
-                &req,
-                &mut *self.ops,
-                &self.yanked_whitelist,
-                &mut |s| {
+            ready!(self
+                .index
+                .query_inner(dep.package_name(), &req, &mut *self.ops, &mut |s| {
                     let matched = match kind {
-                        QueryKind::Exact => dep.matches(&s),
+                        QueryKind::Exact => dep.matches(s.as_summary()),
                         QueryKind::Fuzzy => true,
                     };
-                    if matched {
-                        f(s);
+                    // Next filter out all yanked packages. Some yanked packages may
+                    // leak through if they're in a whitelist (aka if they were
+                    // previously in `Cargo.lock`
+                    if matched
+                        && (!s.is_yanked() || self.yanked_whitelist.contains(&s.package_id()))
+                    {
+                        f(s.into_summary());
                         called = true;
                     }
-                }
-            ))?;
+                }))?;
             if called {
                 return Poll::Ready(Ok(()));
             }
@@ -788,13 +786,9 @@ impl<'cfg> Source for RegistrySource<'cfg> {
                     }
                     any_pending |= self
                         .index
-                        .query_inner(
-                            name_permutation,
-                            &req,
-                            &mut *self.ops,
-                            &self.yanked_whitelist,
-                            f,
-                        )?
+                        .query_inner(name_permutation, &req, &mut *self.ops, &mut |s| {
+                            f(s.into_summary());
+                        })?
                         .is_pending();
                 }
             }