Fix slow hosted Poetry lock rewrite (#1473)

crates/socket-patch-core/src/patch/redirect/pdm.rs

@@ -480,7 +480,7 @@ mod parse_reuse_equivalence_tests {
     /// re-parse, not one per package (#762).
     #[test]
     fn many_patches_render_the_lock_once() {
-        use crate::utils::lock_fragments::RENDERS;
+        use crate::utils::lock_fragments::{RENDERS, SCANS};
         let url = |name: &str| {
             format!("https://patch.socket.dev/patch/pypi/{name}/a/{name}-1.26.18-py3-none-any.whl")
         };
@@ -512,6 +512,7 @@ mod parse_reuse_equivalence_tests {
                     .collect();
                 let files = BTreeMap::from([("pdm.lock".to_string(), lock)]);
                 RENDERS.with(|renders| renders.set(0));
+                SCANS.with(|scans| scans.set(0));
                 let mut got = RewriteResult::default();
                 rewrite(&files, &deps, &mut got);
                 if got.confirmed_pdm_uuids.len() != 12 {
@@ -523,6 +524,9 @@ mod parse_reuse_equivalence_tests {
                     1,
                     "{fixture} crlf={crlf}"
                 );
+                // One pass over the original and one over the result find
+                // and count every dep's fragments (#1473).
+                assert_eq!(SCANS.with(|scans| scans.get()), 2, "{fixture} crlf={crlf}");
             }
         }
         assert!(checked >= 16, "only {checked} locks landed every dep");

crates/socket-patch-core/src/patch/redirect/poetry.rs

@@ -11,8 +11,8 @@ use super::{DepOverride, FileEdit, RewriteResult, RewriteWarning};
 use crate::crawlers::python_crawler::canonicalize_pypi_name;
 use crate::utils::poetry_lock::{
     generated_by_version, lock_version, lock_wiring, pyproject_source_constraints,
-    rewrite_poetry_lock_all, source_constraint_advisory, source_constraint_breaks_writer,
-    source_constraint_refusal, LockBatch, LockStep, PoetryLockDep,
+    rewrite_poetry_lock_all_with_version, source_constraint_advisory,
+    source_constraint_breaks_writer, source_constraint_refusal, LockBatch, LockStep, PoetryLockDep,
 };
 
 /// Whether the lock was written by a Poetry release older than 1.4. Those
@@ -21,15 +21,25 @@ use crate::utils::poetry_lock::{
 /// serving the upstream bytes after the redirect. Formats `0`/`1.0`/`1.1` are
 /// only written by such releases; lock `2.0` is written by 1.3 through 1.8, so
 /// the `@generated by Poetry X.Y.Z` header (present from 1.4) decides there.
-fn pre_1_4_writer(lock_text: &str) -> Option<&'static str> {
-    let lock = lock_text.parse::<DocumentMut>().ok()?;
-    match lock_version(&lock) {
-        Ok("0") => Some("0"),
-        Ok("1.0") => Some("1.0"),
-        Ok("1.1") => Some("1.1"),
-        Ok("2.0") if !matches!(generated_by_version(lock_text), Some(v) if v >= (1, 4)) => {
-            Some("2.0")
+///
+/// `version` is the lock's `lock_version` when the caller already has it;
+/// otherwise the lock is parsed for it.
+fn pre_1_4_writer(lock_text: &str, version: Option<&str>) -> Option<&'static str> {
+    let parsed;
+    let version = match version {
+        Some(version) => version,
+        None => {
+            #[cfg(test)]
+            crate::utils::lock_fragments::tick(&crate::utils::lock_fragments::PARSES);
+            parsed = lock_text.parse::<DocumentMut>().ok()?;
+            lock_version(&parsed).ok()?
         }
+    };
+    match version {
+        "0" => Some("0"),
+        "1.0" => Some("1.0"),
+        "1.1" => Some("1.1"),
+        "2.0" if !matches!(generated_by_version(lock_text), Some(v) if v >= (1, 4)) => Some("2.0"),
         _ => None,
     }
 }
@@ -116,13 +126,17 @@ pub(super) fn rewrite_poetry(
             })
             .collect();
         // Every dep is rewritten over one parse and one render of the lock.
-        let LockBatch {
-            text: content,
-            steps,
-        } = rewrite_poetry_lock_all(original, &deps);
+        let (
+            LockBatch {
+                text: content,
+                steps,
+            },
+            batch_lock_version,
+        ) = rewrite_poetry_lock_all_with_version(original, &deps);
         let mut stale_warned = false;
         // `pre_1_4_writer` reads only the first line and `[metadata]
-        // lock-version`, which no package rewrite touches: judged once.
+        // lock-version`, which no package rewrite touches: judged once, from
+        // the batch's parse when it has one.
         let mut writer_format: Option<Option<&'static str>> = None;
         for (&(dep, _), step) in usable.iter().zip(steps) {
             match step {
@@ -152,9 +166,9 @@ pub(super) fn rewrite_poetry(
                         });
                     }
                     if !stale_warned {
-                        if let Some(format) =
-                            *writer_format.get_or_insert_with(|| pre_1_4_writer(&content))
-                        {
+                        if let Some(format) = *writer_format.get_or_insert_with(|| {
+                            pre_1_4_writer(&content, batch_lock_version.as_deref())
+                        }) {
                             stale_warned = true;
                             // Poetry 1.0 installs url sources through pip, which reads
                             // the hash from the URL fragment; pip 22.3–23.0 take the
@@ -471,7 +485,7 @@ mod equivalence_tests {
     /// re-parse, not one per package (#760).
     #[test]
     fn many_patches_render_the_lock_once() {
-        use crate::utils::lock_fragments::RENDERS;
+        use crate::utils::lock_fragments::{PARSES, RENDERS, SCANS};
         for version in VERSIONS.iter().filter(|version| !version.starts_with("0.")) {
             for crlf in [false, true] {
                 let mut lock = grown(version, 11);
@@ -492,13 +506,25 @@ mod equivalence_tests {
                 let mut deps = vec![dep("urllib3", "1.26.18", Some(SHA), 0)];
                 deps.extend((0..11).map(|i| dep(&format!("pkg{i}"), "1.26.18", Some(SHA), i + 1)));
                 RENDERS.with(|renders| renders.set(0));
+                PARSES.with(|parses| parses.set(0));
+                SCANS.with(|scans| scans.set(0));
                 let got = run(rewrite_poetry, &files, &deps);
                 assert_eq!(got.confirmed_python_lock_uuids.len(), 12, "{version}");
                 assert_eq!(
                     RENDERS.with(|renders| renders.get()),
                     1,
                     "{version} crlf={crlf}"
                 );
+                // The batch parses the lock and its rendering, and the
+                // stale-install check reads the batch's parse instead of its
+                // own; the fragments are found and counted in one pass over
+                // the original and one over the result (#1473).
+                assert_eq!(
+                    PARSES.with(|parses| parses.get()),
+                    2,
+                    "{version} crlf={crlf}"
+                );
+                assert_eq!(SCANS.with(|scans| scans.get()), 2, "{version} crlf={crlf}");
             }
         }
     }

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

@@ -18,6 +18,19 @@ thread_local! {
     /// Whole-lock renders this thread's rewrites took, each followed by a
     /// full re-parse: what a hosted rewrite of N deps must not pay N times.
     pub(crate) static RENDERS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
+    /// Whole-lock parses at the instrumented sites: [`rewrite_batch`]'s
+    /// parse of the original and of its rendering, and the Poetry
+    /// redirect's stale-install check when it has to parse for itself.
+    pub(crate) static PARSES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
+    /// Whole-text fragment searches [`rewrite_batch`] took to find or count
+    /// its changed fragments: what a batch of N deps must not pay N times.
+    pub(crate) static SCANS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
+}
+
+/// Count one tick of a test counter ([`RENDERS`], [`PARSES`], [`SCANS`]).
+#[cfg(test)]
+pub(crate) fn tick(counter: &'static std::thread::LocalKey<std::cell::Cell<usize>>) {
+    counter.with(|count| count.set(count.get() + 1));
 }
 
 /// Takes `name`'s fragments of `text` from its (spanned) parse.
@@ -220,7 +233,7 @@ pub(crate) fn finish<'a>(
     fragments_in: FragmentsIn,
 ) -> Result<FragmentRewrite<'a>, String> {
     #[cfg(test)]
-    RENDERS.with(|renders| renders.set(renders.get() + 1));
+    tick(&RENDERS);
     let before = before?;
     let rendered = crate::utils::python_lock::preserve_line_endings(text, rendered);
     let after_doc = toml_edit::Document::parse(rendered).map_err(|e| e.to_string())?;
@@ -271,6 +284,45 @@ pub(crate) fn finish<'a>(
     })
 }
 
+/// The start of each of `needles` in `haystack`, found in one pass; `None`
+/// unless every needle occurs exactly once.
+///
+/// The pass looks for each needle's first [`PREFIX`] bytes (a small
+/// automaton, where the whole fragments would make a large one) and checks
+/// the full needle at every place its prefix occurs. Every occurrence starts
+/// with its prefix, so none is missed. They are counted overlapping, so a
+/// needle that overlaps itself may count more than its non-overlapping
+/// `matches` would; that only ever sends the batch to the step-by-step
+/// rewrite, never accepts one that path would refuse.
+fn sole_occurrences(haystack: &str, needles: &[&str]) -> Option<Vec<usize>> {
+    /// Enough of a lock fragment to name its package.
+    const PREFIX: usize = 64;
+    if needles.is_empty() {
+        return Some(Vec::new());
+    }
+    if needles.iter().any(|needle| needle.is_empty()) {
+        return None;
+    }
+    #[cfg(test)]
+    tick(&SCANS);
+    let prefixes = needles
+        .iter()
+        .map(|needle| &needle.as_bytes()[..needle.len().min(PREFIX)]);
+    let finder = aho_corasick::AhoCorasick::new(prefixes).ok()?;
+    let haystack = haystack.as_bytes();
+    let mut starts: Vec<Option<usize>> = vec![None; needles.len()];
+    for found in finder.find_overlapping_iter(haystack) {
+        let index = found.pattern().as_usize();
+        if !haystack[found.start()..].starts_with(needles[index].as_bytes()) {
+            continue;
+        }
+        if starts[index].replace(found.start()).is_some() {
+            return None;
+        }
+    }
+    starts.into_iter().collect()
+}
+
 /// What one dep's rewrite did to a lock, whether it ran in a
 /// [`rewrite_batch`] or step by step.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -318,6 +370,8 @@ pub(crate) fn rewrite_batch<P>(
     if text.contains("\r\n") && text.replace("\r\n", "").contains('\n') {
         return None;
     }
+    #[cfg(test)]
+    tick(&PARSES);
     let original = toml_edit::Document::parse(text.to_owned()).ok()?;
     let mut lock = original.clone().into_mut();
     let mut steps = Vec::with_capacity(names.len());
@@ -347,14 +401,13 @@ pub(crate) fn rewrite_batch<P>(
         });
     }
     #[cfg(test)]
-    RENDERS.with(|renders| renders.set(renders.get() + 1));
+    tick(&RENDERS);
     let rendered = crate::utils::python_lock::preserve_line_endings(text, lock.to_string());
+    #[cfg(test)]
+    tick(&PARSES);
     let after_doc = toml_edit::Document::parse(rendered).ok()?;
     let rendered = after_doc.raw();
     let file_terminator = majority_terminator(text);
-    // Each changed fragment as the byte range of `text` it replaces (trimmed
-    // to what differs) and the replacement.
-    let mut splices: Vec<(std::ops::Range<usize>, &str)> = Vec::new();
     let mut edits_of: Vec<(usize, Vec<(String, String)>)> = Vec::new();
     for (index, before) in &rewritten {
         let after = fragments_in(&after_doc, rendered, names[*index]).ok()?;
@@ -364,40 +417,42 @@ pub(crate) fn rewrite_batch<P>(
         let mut edits = Vec::new();
         for (old, new) in before.iter().zip(after) {
             let new = respell(old, &new, file_terminator);
-            if *old == new {
-                continue;
-            }
-            if text.matches(old.as_str()).count() != 1 {
-                return None;
+            if *old != new {
+                edits.push((old.clone(), new));
             }
-            edits.push((old.clone(), new));
         }
         edits_of.push((*index, edits));
     }
-    for (_, edits) in &edits_of {
-        for (old, new) in edits {
-            let at = text.find(old.as_str())?;
-            // Only the bytes that differ are replaced, so the units' shared
-            // boundaries (a unit's fragment ends at the next unit's header)
-            // never overlap.
-            let common =
-                |a: &mut dyn Iterator<Item = (u8, u8)>| a.take_while(|(x, y)| x == y).count();
-            let mut prefix = common(&mut old.bytes().zip(new.bytes()));
-            while !old.is_char_boundary(prefix) || !new.is_char_boundary(prefix) {
-                prefix -= 1;
-            }
-            let mut suffix =
-                common(&mut old[prefix..].bytes().rev().zip(new[prefix..].bytes().rev()));
-            while !old.is_char_boundary(old.len() - suffix)
-                || !new.is_char_boundary(new.len() - suffix)
-            {
-                suffix -= 1;
-            }
-            splices.push((
-                at + prefix..at + old.len() - suffix,
-                &new[prefix..new.len() - suffix],
-            ));
+    // Where each changed fragment sits in `text`, all found in one pass; a
+    // fragment that is not there exactly once ends the batch, as it refuses
+    // the step-by-step rewrite.
+    let olds: Vec<&str> = edits_of
+        .iter()
+        .flat_map(|(_, edits)| edits.iter().map(|(old, _)| old.as_str()))
+        .collect();
+    let starts = sole_occurrences(text, &olds)?;
+    // Each changed fragment as the byte range of `text` it replaces (trimmed
+    // to what differs) and the replacement.
+    let mut splices: Vec<(std::ops::Range<usize>, &str)> = Vec::new();
+    let edits = edits_of.iter().flat_map(|(_, edits)| edits);
+    for ((old, new), at) in edits.zip(starts) {
+        // Only the bytes that differ are replaced, so the units' shared
+        // boundaries (a unit's fragment ends at the next unit's header)
+        // never overlap.
+        let common = |a: &mut dyn Iterator<Item = (u8, u8)>| a.take_while(|(x, y)| x == y).count();
+        let mut prefix = common(&mut old.bytes().zip(new.bytes()));
+        while !old.is_char_boundary(prefix) || !new.is_char_boundary(prefix) {
+            prefix -= 1;
         }
+        let mut suffix = common(&mut old[prefix..].bytes().rev().zip(new[prefix..].bytes().rev()));
+        while !old.is_char_boundary(old.len() - suffix) || !new.is_char_boundary(new.len() - suffix)
+        {
+            suffix -= 1;
+        }
+        splices.push((
+            at + prefix..at + old.len() - suffix,
+            &new[prefix..new.len() - suffix],
+        ));
     }
     splices.sort_by_key(|(range, _)| (range.start, range.end));
     let mut result = String::with_capacity(rendered.len());
@@ -414,13 +469,14 @@ pub(crate) fn rewrite_batch<P>(
     if result != rendered {
         return None;
     }
+    // Every replacement must be unique in the result too, so its rollback
+    // can't hit the wrong place.
+    let news: Vec<&str> = edits_of
+        .iter()
+        .flat_map(|(_, edits)| edits.iter().map(|(_, new)| new.as_str()))
+        .collect();
+    sole_occurrences(&result, &news)?;
     for (index, edits) in edits_of {
-        if edits
-            .iter()
-            .any(|(_, new)| result.matches(new.as_str()).count() != 1)
-        {
-            return None;
-        }
         if !edits.is_empty() {
             steps[index] = LockStep::Rewritten(edits);
         }
@@ -450,6 +506,26 @@ mod tests {
         Ok(vec![text[span].to_string()])
     }
 
+    #[test]
+    fn sole_occurrences_finds_each_needle_once_or_refuses() {
+        // Long needles sharing their first 64+ bytes are told apart by the
+        // full comparison.
+        let shared = "x".repeat(80);
+        let (a, b) = (format!("{shared}A"), format!("{shared}B"));
+        let haystack = format!("[{a}]\n[{b}]\n");
+        assert_eq!(
+            sole_occurrences(&haystack, &[&b, &a, "]\n["]),
+            Some(vec![haystack.find(&b).unwrap(), 1, a.len() + 1])
+        );
+        assert_eq!(sole_occurrences(&haystack, &[]), Some(vec![]));
+        // Missing, repeated, empty and self-overlapping needles all refuse.
+        assert_eq!(sole_occurrences(&haystack, &[&a, "missing"]), None);
+        assert_eq!(sole_occurrences(&haystack, &[&shared]), None);
+        assert_eq!(sole_occurrences(&haystack, &["]\n"]), None);
+        assert_eq!(sole_occurrences(&haystack, &[""]), None);
+        assert_eq!(sole_occurrences("aaa", &["aa"]), None);
+    }
+
     #[test]
     fn pair_fragments_refuses_a_shape_change() {
         for what in ["Poetry", "PDM"] {