Fix Poetry/PDM scans re-parsing the lock (#1473, #1624)

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

@@ -5,7 +5,6 @@
 use std::collections::BTreeMap;
 
 use serde_json::Value;
-use toml_edit::DocumentMut;
 
 use super::{DepOverride, FileEdit, RewriteResult, RewriteWarning};
 use crate::crawlers::python_crawler::canonicalize_pypi_name;
@@ -29,9 +28,8 @@ fn pre_1_4_writer(lock_text: &str, version: Option<&str>) -> Option<&'static str
     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()?;
+            // The batch already parsed (and shared) these bytes.
+            parsed = crate::utils::lock_fragments::parse_lock(lock_text).ok()?;
             lock_version(&parsed).ok()?
         }
     };

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

@@ -9,9 +9,12 @@
 //! refuses) stays in `poetry_lock` / `pdm_lock`; everything around it lives
 //! here, once.
 
+use std::sync::Arc;
+
 use toml_edit::Table;
 
 use crate::utils::line_endings::majority_terminator;
+use crate::vendor::parse_memo::ParseMemo;
 
 #[cfg(test)]
 thread_local! {
@@ -33,6 +36,54 @@ pub(crate) fn tick(counter: &'static std::thread::LocalKey<std::cell::Cell<usize
     counter.with(|count| count.set(count.get() + 1));
 }
 
+/// The read-only parses of the `poetry.lock` / `pdm.lock` texts this run
+/// last read or wrote (see `vendor::parse_memo`). A hosted scan reads one
+/// lock from four places (the lock inventory, VEX discovery before and
+/// after the rewrite, the redirect's own probe) and the rewrite renders it
+/// once more; each whole-lock parse is ~10% of the scan on a 400-package
+/// lock (#1473, #1624), so those readers go through here and a hosted scan
+/// parses each of the two texts once. Two slots: the lock as read and as
+/// rewritten; a miss (another lock read in between) costs only the parse
+/// it would have cost without the memo. Intermediate texts of the
+/// one-dep-at-a-time rewrite stay out of it, as nobody reads them again.
+#[cfg(not(test))]
+static LOCK_DOCS: ParseMemo<toml_edit::Document<String>, 2> = ParseMemo::new();
+
+#[cfg(test)]
+thread_local! {
+    /// Per thread under test, so one test's lock never satisfies another's
+    /// parse and the [`PARSES`] counts stay exact.
+    static LOCK_DOCS: ParseMemo<toml_edit::Document<String>, 2> = const { ParseMemo::new() };
+}
+
+fn with_lock_docs<R>(f: impl FnOnce(&ParseMemo<toml_edit::Document<String>, 2>) -> R) -> R {
+    #[cfg(not(test))]
+    return f(&LOCK_DOCS);
+    #[cfg(test)]
+    return LOCK_DOCS.with(f);
+}
+
+/// The read-only parse of a `poetry.lock` / `pdm.lock` `text`, shared with
+/// every other reader of the same bytes this run. A parse error is not
+/// cached, so each caller gets the error a fresh parse raises.
+pub(crate) fn parse_lock(
+    text: &str,
+) -> Result<Arc<toml_edit::Document<String>>, toml_edit::TomlError> {
+    with_lock_docs(|memo| {
+        memo.parse(text.as_bytes(), || {
+            #[cfg(test)]
+            tick(&PARSES);
+            toml_edit::Document::parse(text.to_owned())
+        })
+    })
+}
+
+/// Hand the parse of a lock text this run produced to the next reader of
+/// those bytes (VEX discovery reads the written lock back).
+fn share_lock(doc: toml_edit::Document<String>) -> Arc<toml_edit::Document<String>> {
+    with_lock_docs(|memo| memo.store(doc.raw().as_bytes().to_vec(), doc))
+}
+
 /// Takes `name`'s fragments of `text` from its (spanned) parse.
 pub(crate) type FragmentsIn =
     fn(&toml_edit::Document<String>, &str, &str) -> Result<Vec<String>, String>;
@@ -44,39 +95,52 @@ pub(crate) type FragmentsIn =
 /// fragments) when the splice reproduced that output byte for byte.
 /// `DocumentMut`'s own parser is exactly `Document::parse(..).into_mut()`, so
 /// every result is the fresh parse's.
+/// Backed by [`parse_lock`], so a text another reader already parsed costs
+/// nothing here either.
 #[derive(Default)]
 pub struct LockParse {
-    doc: Option<toml_edit::Document<String>>,
+    doc: Option<Arc<toml_edit::Document<String>>>,
 }
 
 impl LockParse {
     /// The parse of `text`, reusing the held one when it is of these bytes.
     pub fn parsed(&mut self, text: &str) -> Result<&Table, toml_edit::TomlError> {
         if self.doc.as_ref().is_none_or(|doc| doc.raw() != text) {
             self.doc = None;
-            self.doc = Some(toml_edit::Document::parse(text.to_owned())?);
+            self.doc = Some(parse_lock(text)?);
         }
         Ok(self.doc.as_ref().expect("parsed just above").as_table())
     }
 
-    /// The held parse of `text` (taken out), or a fresh one; `what` names the
-    /// lock in the parse error.
+    /// The held parse of `text` (taken out), a copy of the shared one, or a
+    /// fresh private parse; `what` names the lock in the parse error. A
+    /// parse someone else still holds is copied, which costs a fraction of
+    /// parsing it again; a text nobody parsed yet is parsed without sharing
+    /// it, as the caller is about to mutate it into the next text.
     pub(crate) fn take(
         &mut self,
         text: &str,
         what: &str,
     ) -> Result<toml_edit::Document<String>, String> {
-        match self.doc.take() {
-            Some(doc) if doc.raw() == text => Ok(doc),
-            _ => toml_edit::Document::parse(text.to_owned())
-                .map_err(|e| format!("invalid {what} lock: {e}")),
-        }
+        let doc = match self.doc.take() {
+            Some(doc) if doc.raw() == text => doc,
+            _ => match with_lock_docs(|memo| memo.get(text.as_bytes())) {
+                Some(shared) => shared,
+                None => {
+                    #[cfg(test)]
+                    tick(&PARSES);
+                    return toml_edit::Document::parse(text.to_owned())
+                        .map_err(|e| format!("invalid {what} lock: {e}"));
+                }
+            },
+        };
+        Ok(Arc::try_unwrap(doc).unwrap_or_else(|shared| (*shared).clone()))
     }
 
     /// Hand back a parse of the text it was taken for, unchanged (a refusal
     /// before any mutation).
     pub(crate) fn restore(&mut self, doc: toml_edit::Document<String>) {
-        self.doc = Some(doc);
+        self.doc = Some(Arc::new(doc));
     }
 }
 
@@ -370,10 +434,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 original = parse_lock(text).ok()?;
+    let mut lock = (*original).clone().into_mut();
     let mut steps = Vec::with_capacity(names.len());
     let mut rewritten: Vec<(usize, Vec<String>)> = Vec::new();
     let mut packages = std::collections::BTreeSet::new();
@@ -476,6 +538,9 @@ pub(crate) fn rewrite_batch<P>(
         .flat_map(|(_, edits)| edits.iter().map(|(_, new)| new.as_str()))
         .collect();
     sole_occurrences(&result, &news)?;
+    // `result` is the rendering `after_doc` parsed: whoever reads the
+    // written lock back reads this parse.
+    share_lock(after_doc);
     for (index, edits) in edits_of {
         if !edits.is_empty() {
             steps[index] = LockStep::Rewritten(edits);

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

@@ -396,7 +396,8 @@ pub(super) fn python_lock_inventory(text: &str) -> Option<Vec<LockfileEntry>> {
 /// can be chosen there). A lock that is not TOML contributes nothing.
 async fn inventory_poetry_lock(view: &ProjectView<'_>) -> Option<Vec<LockfileEntry>> {
     let text = view.read_text("poetry.lock").await.ok()?;
-    let document = toml_edit::Document::parse(text.as_str()).ok()?;
+    // The parse VEX discovery and the redirect share (#1473, #1624).
+    let document = crate::utils::lock_fragments::parse_lock(&text).ok()?;
     let pure_wheel_sha = |files: Vec<&dyn TableLike>| {
         files.into_iter().find_map(|entry| {
             let file = entry.get("file")?.as_str()?;
@@ -589,7 +590,8 @@ async fn inventory_pipfile_lock(view: &ProjectView<'_>) -> Option<Vec<LockfileEn
 /// does not probe, so a lock-only vendored path would not survive a re-scan.
 async fn inventory_pdm_lock(view: &ProjectView<'_>) -> Option<Vec<LockfileEntry>> {
     let text = view.read_text("pdm.lock").await.ok()?;
-    let document = toml_edit::Document::parse(text.as_str()).ok()?;
+    // The parse VEX discovery and the redirect share (#1473, #1624).
+    let document = crate::utils::lock_fragments::parse_lock(&text).ok()?;
     let out: Vec<LockfileEntry> = toml_package_coords(&document)
         .into_iter()
         .map(|(name, version, purl, _)| LockfileEntry {

crates/socket-patch-core/src/vex/discover/mod.rs

@@ -2496,7 +2496,18 @@ pub(crate) fn toml_or_diag<S: AsRef<str>>(
     style: TomlDiag,
     out: &mut Discovery,
 ) -> Option<toml_edit::Document<S>> {
-    match toml_edit::Document::parse(text) {
+    toml_result_or_diag(file, toml_edit::Document::parse(text), style, out)
+}
+
+/// [`toml_or_diag`] for a parse the caller already took (a shared one, as
+/// `lock_fragments::parse_lock` hands out).
+pub(crate) fn toml_result_or_diag<T>(
+    file: &str,
+    parsed: Result<T, toml_edit::TomlError>,
+    style: TomlDiag,
+    out: &mut Discovery,
+) -> Option<T> {
+    match parsed {
         Ok(doc) => Some(doc),
         Err(e) => {
             let reason = match style {

crates/socket-patch-core/src/vex/discover/pypi_locks.rs

@@ -101,8 +101,9 @@ use std::collections::BTreeSet;
 use toml_edit::{DocumentMut, Item, Table, TableLike, Value};
 
 use super::{
-    mentions_vendor_dir, pypi_purl, toml_or_diag, vendored_leaf_purl, DiscoverCtx, Discovery,
-    LocateOpts, PatchedRef, TomlDiag, Wired, DIAG_LOCKFILE_UNPARSEABLE, DIAG_REF_INVALID,
+    mentions_vendor_dir, pypi_purl, toml_or_diag, toml_result_or_diag, vendored_leaf_purl,
+    DiscoverCtx, Discovery, LocateOpts, PatchedRef, TomlDiag, Wired, DIAG_LOCKFILE_UNPARSEABLE,
+    DIAG_REF_INVALID,
 };
 use crate::crawlers::python_crawler::canonicalize_pypi_name;
 use crate::utils::digest::{sha256_hex, sha256_prefixed};
@@ -795,7 +796,10 @@ async fn extract_poetry(ctx: &DiscoverCtx<'_>, out: &mut Discovery) {
     let Some(text) = ctx.read_text(file, out).await else {
         return;
     };
-    let Some(doc) = toml_or_diag(file, &text, TomlDiag::Raw, out) else {
+    // The shared parse: the lock inventory and the redirect read these
+    // same bytes (#1473, #1624).
+    let parsed = crate::utils::lock_fragments::parse_lock(&text);
+    let Some(doc) = toml_result_or_diag(file, parsed, TomlDiag::Raw, out) else {
         return;
     };
     let format = crate::utils::poetry_lock::lock_version(&doc)
@@ -931,7 +935,10 @@ async fn extract_pdm(ctx: &DiscoverCtx<'_>, out: &mut Discovery) {
     let Some(text) = ctx.read_text(file, out).await else {
         return;
     };
-    let Some(doc) = toml_or_diag(file, &text, TomlDiag::Raw, out) else {
+    // The shared parse: the lock inventory and the redirect read these
+    // same bytes (#1473, #1624).
+    let parsed = crate::utils::lock_fragments::parse_lock(&text);
+    let Some(doc) = toml_result_or_diag(file, parsed, TomlDiag::Raw, out) else {
         return;
     };
     let lock_version = doc
@@ -2033,4 +2040,102 @@ mod tests {
         assert_eq!(elsewhere, ["pkg:pypi/click@8.1.7", "pkg:pypi/idna@3.6"]);
         assert!(out.diagnostics.is_empty(), "{:#?}", out.diagnostics);
     }
+
+    // ── one parse per lock text (#1473, #1624) ───────────────────────────
+
+    /// What a hosted scan does to one lock: the lock inventory reads it,
+    /// VEX discovery reads it, the batch rewrite rewrites it, and VEX
+    /// discovery reads the written lock back. Each distinct text costs one
+    /// whole-lock parse: the original, and the rewrite's rendering.
+    async fn scan_parses(lock: &str, file: &str, rewritten: &str) -> usize {
+        use crate::utils::lock_fragments::PARSES;
+        let p = Project::new();
+        p.write(file, lock);
+        PARSES.with(|parses| parses.set(0));
+        let inventory = crate::vendor::lock_inventory::inventory_project(p.root()).await;
+        assert!(
+            inventory.iter().any(|entry| entry.name == "urllib3"),
+            "{file}: {inventory:?}"
+        );
+        assert_refs(&run(&p).await, &[]);
+        let written = rewritten_lock(file, lock);
+        assert_eq!(written, rewritten, "{file}: the batch output");
+        p.write(file, &written);
+        assert_refs(&run(&p).await, &[(URLLIB3, UUID_A, WiringMode::Hosted)]);
+        PARSES.with(|parses| parses.get())
+    }
+
+    fn rewritten_lock(file: &str, lock: &str) -> String {
+        use crate::utils::lock_fragments::{LockBatch, LockStep};
+        let url = hosted_url("pypi", "urllib3", "1.26.18", UUID_A, URLLIB3_WHEEL);
+        let LockBatch { text, steps } = if file == super::POETRY_LOCK {
+            crate::utils::poetry_lock::rewrite_poetry_lock_all(
+                lock,
+                &[crate::utils::poetry_lock::PoetryLockDep {
+                    name: "urllib3",
+                    version: "1.26.18",
+                    source_type: "url",
+                    source_url: &url,
+                    filename: URLLIB3_WHEEL,
+                    sha256: SHA,
+                }],
+            )
+        } else {
+            crate::utils::pdm_lock::rewrite_pdm_lock_all(
+                lock,
+                &[crate::utils::pdm_lock::PdmLockDep {
+                    name: "urllib3",
+                    version: "1.26.18",
+                    source: ("url", &url),
+                    filename: URLLIB3_WHEEL,
+                    sha256: SHA,
+                }],
+            )
+        };
+        assert!(
+            matches!(steps[..], [LockStep::Rewritten(_)]),
+            "{file}: {steps:?}"
+        );
+        text
+    }
+
+    /// #1473: a hosted Poetry scan parsed `poetry.lock` five times.
+    #[tokio::test]
+    async fn poetry_scan_parses_each_lock_text_once() {
+        for version in ["1.2.2", "1.8.5", "2.1.4"] {
+            let native = poetry_fixture(version);
+            // The expected output, rewritten on another thread so its
+            // parses neither count here nor seed this thread's shared parse.
+            let rewritten = {
+                let native = native.clone();
+                std::thread::spawn(move || rewritten_lock(super::POETRY_LOCK, &native))
+                    .join()
+                    .unwrap()
+            };
+            assert_eq!(
+                scan_parses(&native, super::POETRY_LOCK, &rewritten).await,
+                2,
+                "poetry {version}"
+            );
+        }
+    }
+
+    /// #1624: a hosted PDM scan parsed `pdm.lock` six times.
+    #[tokio::test]
+    async fn pdm_scan_parses_each_lock_text_once() {
+        for version in ["2.11.2", "2.17.3", "2.29.2"] {
+            let native = pdm_fixture(version);
+            let rewritten = {
+                let native = native.clone();
+                std::thread::spawn(move || rewritten_lock(super::PDM_LOCK, &native))
+                    .join()
+                    .unwrap()
+            };
+            assert_eq!(
+                scan_parses(&native, super::PDM_LOCK, &rewritten).await,
+                2,
+                "pdm {version}"
+            );
+        }
+    }
 }