diff --git a/crates/socket-patch-core/src/patch/redirect/pdm.rs b/crates/socket-patch-core/src/patch/redirect/pdm.rs index 3cb7054c4..28ce88a05 100644 --- a/crates/socket-patch-core/src/patch/redirect/pdm.rs +++ b/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"); diff --git a/crates/socket-patch-core/src/patch/redirect/poetry.rs b/crates/socket-patch-core/src/patch/redirect/poetry.rs index c875d9f20..b7ce6e344 100644 --- a/crates/socket-patch-core/src/patch/redirect/poetry.rs +++ b/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::().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::().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> = 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,6 +506,8 @@ 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!( @@ -499,6 +515,16 @@ mod equivalence_tests { 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}"); } } } diff --git a/crates/socket-patch-core/src/utils/lock_fragments.rs b/crates/socket-patch-core/src/utils/lock_fragments.rs index 354b08baa..d0e234670 100644 --- a/crates/socket-patch-core/src/utils/lock_fragments.rs +++ b/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 = 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 = 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 = 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>) { + 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, 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> { + /// 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> = 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

( 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

( }); } #[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, &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

( 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| 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, &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| 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

( 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"] { diff --git a/crates/socket-patch-core/src/utils/poetry_lock.rs b/crates/socket-patch-core/src/utils/poetry_lock.rs index c7280c770..8c6487b9f 100644 --- a/crates/socket-patch-core/src/utils/poetry_lock.rs +++ b/crates/socket-patch-core/src/utils/poetry_lock.rs @@ -568,17 +568,40 @@ pub struct PoetryLockDep<'a> { /// them when the batch can vouch for the step-by-step result, else step by /// step. pub fn rewrite_poetry_lock_all(text: &str, deps: &[PoetryLockDep]) -> LockBatch { - rewrite_poetry_lock_batch(text, deps).unwrap_or_else(|| rewrite_poetry_lock_steps(text, deps)) + rewrite_poetry_lock_all_with_version(text, deps).0 +} + +/// [`rewrite_poetry_lock_all`], plus the lock's `lock_version` when the +/// batch read it from its own parse (`None` when it went step by step, or +/// the lock has no supported version), so a caller needn't parse the lock +/// again for it. No package rewrite touches `[metadata] lock-version`, so +/// it is the rewritten lock's version too. +pub fn rewrite_poetry_lock_all_with_version( + text: &str, + deps: &[PoetryLockDep], +) -> (LockBatch, Option) { + let mut version = None; + match rewrite_poetry_lock_batch(text, deps, &mut version) { + Some(batch) => (batch, version), + None => (rewrite_poetry_lock_steps(text, deps), None), + } } /// [`rewrite_poetry_lock_all`] over one parse and one render, or `None` (see -/// [`rewrite_batch`]). -fn rewrite_poetry_lock_batch(text: &str, deps: &[PoetryLockDep]) -> Option { +/// [`rewrite_batch`]); `version` gets the parse's `lock_version`. +fn rewrite_poetry_lock_batch( + text: &str, + deps: &[PoetryLockDep], + version: &mut Option, +) -> Option { let names: Vec<&str> = deps.iter().map(|dep| dep.name).collect(); rewrite_batch( text, &names, |index, lock| { + if version.is_none() { + *version = lock_version_of(lock).ok().map(str::to_string); + } let dep = &deps[index]; let sha256 = checked_sha256( dep.name, @@ -1635,10 +1658,20 @@ mod batch_equivalence_tests { { cases += 1; let steps = rewrite_poetry_lock_steps(&text, &deps); - let batch = rewrite_poetry_lock_batch(&text, &deps); + let mut read_version = None; + let batch = rewrite_poetry_lock_batch(&text, &deps, &mut read_version); if let Some(batch) = &batch { batched += 1; assert_eq!(batch, &steps, "{what}"); + // The version the batch read is the output's. + if read_version.is_some() { + let output = batch.text.parse::().unwrap(); + assert_eq!( + read_version.as_deref(), + lock_version(&output).ok(), + "{what}" + ); + } } assert_eq!(rewrite_poetry_lock_all(&text, &deps), steps, "{what}"); // Poetry 0.12 refuses every URL source: nothing to