Speed up bun and yarn-berry rescans (#1606, #1607)

crates/socket-patch-cli/src/commands/scan/mod.rs

@@ -104,12 +104,15 @@ const BATCH_BODY_BYTE_CAP: usize = 256 * 1024;
 /// PEP 440 spellings of the same release (#604), deduplicated in order.
 fn with_pypi_equivalents(purls: &[String], lockfile_only: &HashSet<PurlKey>) -> Vec<String> {
     let mut out = purls.to_vec();
-    let mut seen: HashSet<PurlKey> = purls.iter().map(|p| PurlKey::new(p)).collect();
+    // Keyed only once some purl has an equivalent: most scans have no
+    // PyPI pins, and keying every purl costs (#1606, #1607).
+    let mut seen: Option<HashSet<PurlKey>> = None;
     for purl in purls {
-        if !lockfile_only_contains(lockfile_only, purl) {
+        if !is_pypi_purl(purl) || !lockfile_only_contains(lockfile_only, purl) {
             continue;
         }
         for spelling in socket_patch_core::utils::purl_key::pypi_equivalent_purls(purl) {
+            let seen = seen.get_or_insert_with(|| purls.iter().map(|p| PurlKey::new(p)).collect());
             if seen.insert(PurlKey::new(&spelling)) {
                 out.push(spelling);
             }
@@ -118,6 +121,15 @@ fn with_pypi_equivalents(purls: &[String], lockfile_only: &HashSet<PurlKey>) ->
     out
 }
 
+/// Whether `purl` is a PyPI purl, in any case of its type (the spellings
+/// [`PurlKey`] folds together).
+fn is_pypi_purl(purl: &str) -> bool {
+    purl.trim_start()
+        .strip_prefix("pkg:")
+        .and_then(|rest| rest.get(..5))
+        .is_some_and(|ty| ty.eq_ignore_ascii_case("pypi/"))
+}
+
 /// Mark each API purl that names a lockfile-only PyPI pin under another
 /// PEP 440 spelling (`@1.16.0` for a lock's `@1.16`, #604) as lockfile-only
 /// too, so the `notInstalled` flag, the `[NOT INSTALLED]` marker and the
@@ -128,11 +140,14 @@ fn adopt_pypi_equivalents(
     scanned: &[String],
     lockfile_only: &mut HashSet<PurlKey>,
 ) {
-    let scanned_keys: HashSet<PurlKey> = scanned.iter().map(|p| PurlKey::new(p)).collect();
     let lock_only_pypi: Vec<&String> = scanned
         .iter()
-        .filter(|p| p.starts_with("pkg:pypi/") && lockfile_only_contains(lockfile_only, p))
+        .filter(|p| is_pypi_purl(p) && lockfile_only_contains(lockfile_only, p))
         .collect();
+    if lock_only_pypi.is_empty() {
+        return;
+    }
+    let scanned_keys: HashSet<PurlKey> = scanned.iter().map(|p| PurlKey::new(p)).collect();
     for pkg in packages {
         let key = PurlKey::new(&pkg.purl);
         if scanned_keys.contains(&key) {
@@ -3569,6 +3584,41 @@ async fn run_scan(
 mod tests {
     use super::*;
 
+    #[test]
+    fn pypi_equivalents_add_spellings_only_for_lockfile_only_pypi_pins() {
+        let purls: Vec<String> = [
+            "pkg:npm/left-pad@1.3.0",
+            "pkg:pypi/six@1.16",
+            "pkg:PyPI/attrs@23.1",
+            "pkg:pypi/idna@3.4",
+        ]
+        .map(String::from)
+        .to_vec();
+        let lockfile_only: HashSet<PurlKey> = ["pkg:pypi/six@1.16", "pkg:pypi/attrs@23.1"]
+            .into_iter()
+            .map(PurlKey::new)
+            .collect();
+        let out = with_pypi_equivalents(&purls, &lockfile_only);
+        assert_eq!(out[..purls.len()], purls[..]);
+        let added = &out[purls.len()..];
+        assert!(
+            added.contains(&"pkg:pypi/six@1.16.0".to_string()),
+            "{added:?}"
+        );
+        assert!(
+            added.contains(&"pkg:pypi/attrs@23.1.0".to_string()),
+            "{added:?}"
+        );
+        assert!(added
+            .iter()
+            .all(|p| !p.contains("idna") && !p.contains("npm")));
+        // No PyPI pin: the list goes out as it came in.
+        let npm_only = vec!["pkg:npm/left-pad@1.3.0".to_string()];
+        assert_eq!(with_pypi_equivalents(&npm_only, &lockfile_only), npm_only);
+        assert!(is_pypi_purl(" pkg:PYPI/six@1.16") && !is_pypi_purl("pkg:pypi"));
+        assert!(!is_pypi_purl("PKG:pypi/six@1.16"));
+    }
+
     #[test]
     fn project_dirs_resolve_directories_and_globs() {
         let tmp = tempfile::tempdir().unwrap();

crates/socket-patch-core/src/utils/purl_key.rs

@@ -49,6 +49,37 @@ use crate::vendor::nuget_feed::normalize_nuget_version;
 /// release spellings still differ here; compare with [`PurlKey`], which
 /// folds them.
 pub fn canonical_base_purl(purl: &str) -> String {
+    // Scans key thousands of crawled purls, nearly all already canonical:
+    // skip the rebuild for those (#1606, #1607).
+    if is_canonical_base(purl) {
+        return purl.to_string();
+    }
+    rebuild_canonical_base_purl(purl)
+}
+
+/// Whether [`canonical_base_purl`] would return `purl` unchanged: trimmed,
+/// no qualifiers, subpath or percent escapes, and a lowercase type whose
+/// names and versions keep their spelling.
+fn is_canonical_base(purl: &str) -> bool {
+    if purl.trim().len() != purl.len() || purl.contains(['?', '#', '%']) {
+        return false;
+    }
+    let Some((ty, _)) = purl
+        .strip_prefix("pkg:")
+        .and_then(|rest| rest.split_once('/'))
+    else {
+        return true;
+    };
+    !ty.bytes().any(|b| b.is_ascii_uppercase()) && !folds_spelling(ty)
+}
+
+/// Whether [`canonical_base_purl`] folds the name or version spelling of
+/// a (lowercase) purl type; every other type keeps its spelling.
+fn folds_spelling(ty: &str) -> bool {
+    matches!(ty, "pypi" | "nuget" | "composer")
+}
+
+fn rebuild_canonical_base_purl(purl: &str) -> String {
     let base = normalize_purl(strip_purl_qualifiers(purl.trim())).into_owned();
     let Some(rest) = base.strip_prefix("pkg:") else {
         return base;
@@ -65,7 +96,10 @@ pub fn canonical_base_purl(purl: &str) -> String {
         "pypi" => (canonicalize_pypi_name(name), version.map(str::to_string)),
         "nuget" => (name.to_lowercase(), version.map(str::to_lowercase)),
         "composer" => (name.to_lowercase(), version.map(str::to_string)),
-        _ => (name.to_string(), version.map(str::to_string)),
+        _ => {
+            debug_assert!(!folds_spelling(&ty), "{ty} folds but has no arm");
+            (name.to_string(), version.map(str::to_string))
+        }
     };
     match version {
         Some(version) => format!("pkg:{ty}/{name}@{version}"),
@@ -219,6 +253,39 @@ mod tests {
     use super::*;
     use crate::utils::composer_version::composer_versions_equivalent;
 
+    #[test]
+    fn canonical_fast_path_matches_the_rebuild() {
+        for purl in [
+            "pkg:npm/left-pad@1.3.0",
+            "pkg:npm/@scope/name@1.0.0",
+            "pkg:npm/%40scope/name@1.0.0",
+            "pkg:NPM/left-pad@1.3.0",
+            "pkg:Npm/left-pad@1.3.0",
+            " pkg:npm/left-pad@1.3.0",
+            "pkg:npm/left-pad@1.3.0\n",
+            "pkg:npm/left-pad@1.3.0?vcs_url=x",
+            "pkg:golang/github.com/Foo/Bar@v1.0.0#cmd/tool",
+            "pkg:pypi/Six@1.16",
+            "pkg:PyPI/six@1.16",
+            "pkg:nuget/Newtonsoft.Json@13.0.3",
+            "pkg:composer/Vendor/Name@v1.0",
+            "pkg:maven/org.Example/Lib@1.0",
+            "pkg:gem/nokogiri@1.16.5",
+            "pkg:cargo/serde@1.0.0",
+            "pkg:npm/no-version",
+            "pkg:npm/@scope",
+            "pkg:npm",
+            "not-a-purl",
+            "",
+        ] {
+            assert_eq!(
+                canonical_base_purl(purl),
+                rebuild_canonical_base_purl(purl),
+                "{purl:?}"
+            );
+        }
+    }
+
     #[test]
     fn canonical_base_purl_folds_only_insensitive_ecosystems() {
         assert_eq!(

crates/socket-patch-core/src/vendor/lock_inventory/yarn.rs

@@ -1,8 +1,10 @@
 //! `yarn.lock`, classic and berry: the entry models lockfile discovery
 //! shares ([`classic_entries`], [`berry_entries`]) and the registry views.
 
+use std::collections::{HashMap, HashSet};
 #[cfg(test)]
 use std::path::Path;
+use std::sync::OnceLock;
 
 use crate::formats::yarn::blocks::{
     berry_field, classic_field, live_blocks, scan_blocks, LockBlock,
@@ -74,6 +76,50 @@ pub(crate) struct BerryLock {
     pub(crate) cache_key: Option<String>,
     /// Every block except the exact `__metadata` one, in lock order.
     pub(crate) entries: Vec<YarnEntry>,
+    /// What [`berry_selector_routes_nothing`] asks of the live entries,
+    /// built on its first call.
+    selector_index: OnceLock<SelectorIndex>,
+}
+
+/// The live entries of a [`BerryLock`] as [`berry_selector_routes_nothing`]
+/// reads them, collected once per lock instead of once per selector
+/// (#1607).
+#[derive(Default)]
+struct SelectorIndex {
+    /// Every url a live entry resolves: its locator's reference, and an
+    /// older pin's `__archiveUrl=` binding as written and percent-decoded.
+    urls: HashSet<String>,
+    /// The descriptors live entries are keyed by, per package name, each
+    /// also without its `npm:` protocol.
+    descriptors: HashMap<String, HashSet<String>>,
+}
+
+impl SelectorIndex {
+    fn new(entries: &[YarnEntry]) -> Self {
+        let mut index = Self::default();
+        for entry in entries.iter().filter(|e| e.live) {
+            if let Some(locator) = entry.locator() {
+                index.urls.insert(locator.reference.to_string());
+                if let Some(archive) = locator.archive_url() {
+                    index.urls.insert(archive.to_string());
+                    index.urls.insert(
+                        crate::utils::purl::percent_decode_purl_component(archive).into_owned(),
+                    );
+                }
+            }
+            for pattern in &entry.patterns {
+                let Some((name, descriptor)) = split_pattern(pattern) else {
+                    continue;
+                };
+                let descriptors = index.descriptors.entry(name.to_string()).or_default();
+                descriptors.insert(descriptor.to_string());
+                if let Some(bare) = descriptor.strip_prefix("npm:") {
+                    descriptors.insert(bare.to_string());
+                }
+            }
+        }
+        index
+    }
 }
 
 /// A berry (yarn 2+) `yarn.lock` text — berry reuses the classic block
@@ -84,7 +130,11 @@ pub(crate) fn berry_entries(text: &str) -> BerryLock {
     let cache_key = crate::formats::yarn::berry_gates::cache_key(&blocks).map(str::to_string);
     let mut entries = yarn_entries(blocks, split_berry_key_patterns);
     entries.retain(|e| e.block.key != "__metadata");
-    BerryLock { cache_key, entries }
+    BerryLock {
+        cache_key,
+        entries,
+        selector_index: OnceLock::new(),
+    }
 }
 
 /// Whether a root `package.json` `resolutions` selector routing to the
@@ -108,35 +158,17 @@ pub(crate) fn berry_selector_routes_nothing(lock: &BerryLock, selector: &str, ur
         .map(|i| &selector[i..])
         .unwrap_or(target);
     let range = split_pattern(last).map(|(_, range)| range);
-    for entry in lock.entries.iter().filter(|e| e.live) {
-        if let Some(locator) = entry.locator() {
-            if locator.reference == url
-                || locator.archive_url().is_some_and(|archive| {
-                    archive == url
-                        || crate::utils::purl::percent_decode_purl_component(archive) == url
-                })
-            {
-                return false;
-            }
-        }
-        for pattern in &entry.patterns {
-            let Some((name, descriptor)) = split_pattern(pattern) else {
-                continue;
-            };
-            if name != target {
-                continue;
-            }
-            match range {
-                None => return false,
-                Some(want) => {
-                    if descriptor == want || descriptor.strip_prefix("npm:") == Some(want) {
-                        return false;
-                    }
-                }
-            }
-        }
+    let index = lock
+        .selector_index
+        .get_or_init(|| SelectorIndex::new(&lock.entries));
+    if index.urls.contains(url) {
+        return false;
+    }
+    match (index.descriptors.get(target), range) {
+        (None, _) => true,
+        (Some(_), None) => false,
+        (Some(descriptors), Some(want)) => !descriptors.contains(want),
     }
-    true
 }
 
 /// A berry `checksum:` value as the pin yarn enforces: `<cacheKey>/<hex>`,
@@ -288,3 +320,56 @@ fn berry_registry_view(text: &str) -> Vec<LockfileEntry> {
     }
     out
 }
+
+#[cfg(test)]
+mod selector_tests {
+    use super::*;
+
+    const HOSTED: &str = "https://patch.socket.dev/npm/left-pad/-/left-pad-1.3.0.tgz";
+
+    fn lock() -> BerryLock {
+        berry_entries(
+            "__metadata:\n  version: 8\n  cacheKey: 10c0\n\n\
+             \"lodash@npm:^4.17.0\":\n  version: 4.17.21\n  \
+             resolution: \"lodash@npm:4.17.21\"\n\n\
+             \"left-pad@npm:^1.3.0\":\n  version: 1.3.0\n  \
+             resolution: \"left-pad@npm:1.3.0::__archiveUrl=https%3A%2F%2Fpatch.socket.dev%2Fnpm%2Fleft-pad%2F-%2Fleft-pad-1.3.0.tgz\"\n\n\
+             \"chalk@https://patch.socket.dev/npm/chalk/-/chalk-5.0.0.tgz\":\n  version: 5.0.0\n  \
+             resolution: \"chalk@https://patch.socket.dev/npm/chalk/-/chalk-5.0.0.tgz\"\n",
+        )
+    }
+
+    #[test]
+    fn one_lock_answers_every_selector_from_its_index() {
+        let lock = lock();
+        // The url a live entry resolves, as a locator or a decoded
+        // `__archiveUrl=` binding, is still wired.
+        assert!(!berry_selector_routes_nothing(&lock, "left-pad", HOSTED));
+        let chalk = "https://patch.socket.dev/npm/chalk/-/chalk-5.0.0.tgz";
+        assert!(!berry_selector_routes_nothing(&lock, "nobody", chalk));
+        // A range-less selector matches any descriptor of its package.
+        let gone = "https://patch.socket.dev/npm/lodash/-/lodash-4.17.20.tgz";
+        assert!(!berry_selector_routes_nothing(&lock, "lodash", gone));
+        assert!(!berry_selector_routes_nothing(&lock, "**/lodash", gone));
+        // `name@range` matches that descriptor, with or without `npm:`.
+        assert!(!berry_selector_routes_nothing(
+            &lock,
+            "lodash@^4.17.0",
+            gone
+        ));
+        assert!(!berry_selector_routes_nothing(
+            &lock,
+            "lodash@npm:^4.17.0",
+            gone
+        ));
+        // Leftovers: another range, or a package the lock no longer has.
+        assert!(berry_selector_routes_nothing(&lock, "lodash@^3.0.0", gone));
+        let removed = "https://patch.socket.dev/npm/is-odd/-/is-odd-3.0.1.tgz";
+        assert!(berry_selector_routes_nothing(&lock, "is-odd", removed));
+        assert!(berry_selector_routes_nothing(
+            &lock,
+            "parent/is-odd@^3.0.0",
+            removed
+        ));
+    }
+}