Coalesce per-file plugins into one salsa query per file (#308)

`per_file_plugin_ops` was keyed `(file, plugin_id)`, so N registered
per-file plugins meant N salsa queries (and N memo entries) per file —
the salsa bookkeeping (memo table, dependency edges, per-build
revalidation) grew O(files × plugins). Re-key it on `(file, set_id)`,
where `set_id` interns the ordered, deduped per-file plugin id list, and
run the whole set inside the one query body. Salsa work is now
O(files); the per-plugin loop is plain Rust, free of salsa.

* New `PER_FILE_PLUGIN_SETS` registry interns an ordered id list to a
  stable `set_id` (FxHasher over the list). Interning by the list keeps
  it sound across `Analysis`es with different plugin sets — a different
  list → a different `set_id` → separate memo entries — and identical
  sets (incl. across `re_materialize`) share entries. Same append-only,
  immutable-per-id contract as the configured/external registries, so
  the untracked lookup from inside the query is sound.
* `extract_per_file_plugin_set` dedups by id preserving registration
  order, exactly reproducing the dedup the old `(file, id)` key gave for
  free (two `main_block()`s, or identical `ServerConfig` configs that
  intern to one id, run once).
* Populate warms one query per file; assemble reads one slice per file.
  The shared dispatch is factored into `run_one_per_file`.

No support for the plugin list changing mid-session is added (per the
design): invalidation stays purely per-file via the file's tracked
inputs. Byte-identical output (A/B on the framework corpus with all 17
builtin plugins), full suite (745) green incl. the ServerConfig
cache-interning tests.

https://claude.ai/code/session_01KYNAs8Y9ctXsM1waLmXEaP

Co-authored-by: Claude <noreply@anthropic.com>

Author: lpetre

runtime/src/native_plugins.rs

@@ -943,41 +943,105 @@ fn configured_per_file_plugin(id: u32) -> Option<Arc<ConfiguredPerFile>> {
         .map(Arc::clone)
 }
 
-/// Salsa-tracked per-file plugin invocation. Keyed on ``(file, id)``;
-/// re-runs only when the file's tracked inputs (``file_to_nodes`` /
-/// ``parsed_module`` / ``line_index``) change. Returns the file-local ops
-/// the harness translates to global indices at apply time.
-#[salsa::tracked(returns(ref), heap_size = ruff_memory_usage::heap_size)]
-pub(crate) fn per_file_plugin_ops(
-    db: &dyn ProjectDb,
-    file: File,
+/// Process-global registry interning an *ordered* per-file plugin id list to a
+/// stable `set_id`. This is what lets [`per_file_plugin_ops`] be keyed
+/// `(file, set_id)` — **one** salsa query per file that runs the whole set —
+/// rather than `(file, plugin)`, one per plugin. Salsa's per-build memo work
+/// is then O(files), not O(files × plugins): the per-plugin loop lives inside
+/// the query body, costing nothing in salsa. Interning by the ordered list
+/// keeps it sound across `Analysis`es with different plugin sets (a different
+/// list → a different `set_id` → separate memo entries), and an identical set
+/// reconstructed across a `re_materialize` reuses its `set_id` (and so its
+/// cache entries). Append-only and immutable for a given `set_id`, so the
+/// untracked lookup from inside the query is sound — same contract as
+/// [`CONFIGURED_PER_FILE_PLUGINS`] / [`EXTERNAL_PER_FILE_PLUGINS`].
+static PER_FILE_PLUGIN_SETS: std::sync::OnceLock<std::sync::RwLock<PluginSetRegistry>> =
+    std::sync::OnceLock::new();
+
+#[derive(Default)]
+struct PluginSetRegistry {
+    sets: Vec<Arc<[PerFilePluginId]>>,
+    by_hash: FxHashMap<u64, u32>,
+}
+
+/// Register (intern) an ordered per-file plugin id list, returning its
+/// process-stable `set_id`. An identical list returns the existing id.
+pub(crate) fn register_per_file_set(ids: Vec<PerFilePluginId>) -> u32 {
+    let reg =
+        PER_FILE_PLUGIN_SETS.get_or_init(|| std::sync::RwLock::new(PluginSetRegistry::default()));
+    let mut hasher = FxHasher::default();
+    ids.hash(&mut hasher);
+    let hash = hasher.finish();
+    let mut guard = reg.write().expect("per-file plugin set registry poisoned");
+    if let Some(&id) = guard.by_hash.get(&hash) {
+        return id;
+    }
+    let set_id = guard.sets.len() as u32;
+    guard.sets.push(ids.into());
+    guard.by_hash.insert(hash, set_id);
+    set_id
+}
+
+/// Look up an interned per-file plugin id list by `set_id`.
+fn per_file_set(set_id: u32) -> Option<Arc<[PerFilePluginId]>> {
+    PER_FILE_PLUGIN_SETS
+        .get()?
+        .read()
+        .expect("per-file plugin set registry poisoned")
+        .sets
+        .get(set_id as usize)
+        .map(Arc::clone)
+}
+
+/// Run one per-file plugin (configless builtin, configured builtin, or
+/// external dylib) into `ops` — the single curated [`plugin_api::PerFilePlugin`]
+/// dispatch, shared by the whole-set loop in [`per_file_plugin_ops`].
+fn run_one_per_file(
     id: PerFilePluginId,
-) -> Vec<FileLocalOp> {
-    let file_ctx = FileContext::new(db, file);
-    let pctx = plugin_api::PluginFileCtx::new(&file_ctx);
-    let mut ops = plugin_api::FileOps::new();
-    // Every per-file plugin — configless builtin, configured builtin, or
-    // external dylib — runs through the one curated `PerFilePlugin` surface.
+    pctx: &plugin_api::PluginFileCtx<'_>,
+    ops: &mut plugin_api::FileOps,
+) {
     match id {
-        PerFilePluginId::Builtin(kind) => kind.plugin().run_on_file(&pctx, &mut ops),
+        PerFilePluginId::Builtin(kind) => kind.plugin().run_on_file(pctx, ops),
         PerFilePluginId::Configured(plugin_id) => {
             // Resolve the configured plugin from the registry; a stale id
             // (shouldn't happen) yields no ops.
             if let Some(cfg) = configured_per_file_plugin(plugin_id) {
                 let per_file: &dyn plugin_api::PerFilePlugin = cfg.as_ref();
-                per_file.run_on_file(&pctx, &mut ops);
+                per_file.run_on_file(pctx, ops);
             }
         }
         PerFilePluginId::External(plugin_id) => {
             // Resolve the plugin and its per-file capability from the
             // registry; a stale id (shouldn't happen) yields no ops.
             if let Some(plugin) = external_per_file_plugin(plugin_id) {
                 if let Some(per_file) = plugin.per_file() {
-                    per_file.run_on_file(&pctx, &mut ops);
+                    per_file.run_on_file(pctx, ops);
                 }
             }
         }
     }
+}
+
+/// Salsa-tracked per-file plugin invocation. Keyed on ``(file, set_id)`` — the
+/// whole registered per-file plugin *set* runs in one query, so re-runs are
+/// O(files) not O(files × plugins). Re-runs only when the file's tracked
+/// inputs (``file_to_nodes`` / ``parsed_module`` / ``line_index``) change.
+/// Returns every plugin's file-local ops concatenated in registration order
+/// (the order the harness translates + folds at apply time).
+#[salsa::tracked(returns(ref), heap_size = ruff_memory_usage::heap_size)]
+pub(crate) fn per_file_plugin_ops(db: &dyn ProjectDb, file: File, set_id: u32) -> Vec<FileLocalOp> {
+    let file_ctx = FileContext::new(db, file);
+    let pctx = plugin_api::PluginFileCtx::new(&file_ctx);
+    let mut ops = plugin_api::FileOps::new();
+    // Every per-file plugin in the set — configless builtin, configured
+    // builtin, or external dylib — runs through the one curated
+    // `PerFilePlugin` surface, in registration order.
+    if let Some(ids) = per_file_set(set_id) {
+        for &id in ids.iter() {
+            run_one_per_file(id, &pctx, &mut ops);
+        }
+    }
     ops.into_inner()
 }
 

runtime/src/project.rs

@@ -97,7 +97,7 @@ impl Project {
             false,
             None,
             &counters,
-            &[],
+            None,
             GraphBuilder::with_capacity(0),
             ResolveCache::default(),
             None,
@@ -376,7 +376,7 @@ pub(crate) fn build_project_graph(
     show_progress: bool,
     stack_size: Option<usize>,
     counters: &Arc<ProgressCounters>,
-    per_file_plugin_ids: &[crate::native_plugins::PerFilePluginId],
+    per_file_set_id: Option<u32>,
     carry: GraphBuilder,
     cache: ResolveCache,
     scope: Option<Vec<File>>,
@@ -552,7 +552,6 @@ pub(crate) fn build_project_graph(
         // commit message for the hazard.
         let dist_db: ProjectDatabase = db.clone();
         let files_ref: &[File] = &populate_files;
-        let per_file_ids_ref: &[crate::native_plugins::PerFilePluginId] = per_file_plugin_ids;
         let counters_ref = Arc::clone(counters);
         let run_populate = move || {
             use salsa::Database as _;
@@ -599,14 +598,17 @@ pub(crate) fn build_project_graph(
                                 // reads.
                                 let _ = file_extraction(local_db, file);
                                 // Per-file native plugins: warm the salsa-cached
-                                // `per_file_plugin_ops(file, id)` query on this
-                                // worker so the serial assembly fold below is a
-                                // pure cache read. `run_on_file` is GIL-free and
-                                // touches only this file, so it composes with the
-                                // GIL-released fan-out.
-                                for &id in per_file_ids_ref {
+                                // `per_file_plugin_ops(file, set_id)` query on
+                                // this worker so the assembly fold below is a
+                                // pure cache read. One query runs the whole
+                                // plugin set (keyed on `set_id`), so this is O(1)
+                                // in plugin count, not one warm per plugin.
+                                // `run_on_file` is GIL-free and touches only this
+                                // file, so it composes with the GIL-released
+                                // fan-out.
+                                if let Some(set_id) = per_file_set_id {
                                     let _ = crate::native_plugins::per_file_plugin_ops(
-                                        local_db, file, id,
+                                        local_db, file, set_id,
                                     );
                                 }
                                 // Every per-file AST consumer for this file has
@@ -688,7 +690,7 @@ pub(crate) fn build_project_graph(
         db,
         &project_files,
         &peer_pyi_to_py,
-        per_file_plugin_ids,
+        per_file_set_id,
         counters,
         &reload_log,
         carry,
@@ -879,7 +881,7 @@ fn assemble_graph<'db>(
     db: &'db ProjectDatabase,
     project_files: &[File],
     peer_pyi_to_py: &FxHashMap<File, File>,
-    per_file_plugin_ids: &[crate::native_plugins::PerFilePluginId],
+    per_file_set_id: Option<u32>,
     counters: &Arc<ProgressCounters>,
     reload_log: &ReloadLog,
     carry: GraphBuilder,
@@ -1049,23 +1051,16 @@ fn assemble_graph<'db>(
     for &file in project_files {
         spec_payloads.push(file_to_refspecs(db, file));
     }
-    // Per-file plugin ops, flattened across the registered per-file
-    // plugins (in registration order). Prefetched here so the per-part
-    // edge translation below can run db-free on rayon workers.
-    let plugin_ops: Vec<Vec<&'db crate::native_plugins::FileLocalOp>> = if per_file_plugin_ids
-        .is_empty()
-    {
-        vec![Vec::new(); n_files]
-    } else {
-        project_files
+    // Per-file plugin ops: one salsa query per file runs the whole plugin
+    // set (keyed on `set_id`) and returns every plugin's ops concatenated in
+    // registration order. Prefetched here (a cache read — warmed in populate)
+    // so the per-part edge translation below can run db-free on rayon workers.
+    let plugin_ops: Vec<&'db [crate::native_plugins::FileLocalOp]> = match per_file_set_id {
+        Some(set_id) => project_files
             .iter()
-            .map(|&file| {
-                per_file_plugin_ids
-                    .iter()
-                    .flat_map(|&id| crate::native_plugins::per_file_plugin_ops(db, file, id).iter())
-                    .collect()
-            })
-            .collect()
+            .map(|&file| crate::native_plugins::per_file_plugin_ops(db, file, set_id).as_slice())
+            .collect(),
+        None => vec![&[][..]; n_files],
     };
 
     // Per-file anchor-directory ids. ty's `resolve_module` is anchor-
@@ -2112,15 +2107,15 @@ fn assemble_graph<'db>(
     // derived state, rebuilt from every file's ops each pass.
     let mut topic_facts: FxHashMap<String, Vec<crate::native_plugins::plugin_api::Fact>> =
         FxHashMap::default();
-    if !per_file_plugin_ids.is_empty() {
+    if per_file_set_id.is_some() {
         use crate::native_plugins::FileLocalOp;
         for (pos, ops) in plugin_ops.iter().enumerate() {
             let mint = &mints[pos];
             let len = mint.payload.nodes.len();
             let to_global =
                 |local: u32| ((local as usize) < len).then_some(mint.base + local as usize);
             let mut file_path: Option<String> = None;
-            for op in ops {
+            for op in ops.iter() {
                 match op {
                     FileLocalOp::Edge { .. } => {}
                     FileLocalOp::Entrypoint { decl_local_idx } => {
@@ -2275,7 +2270,7 @@ fn assemble_graph<'db>(
                 {
                     use crate::native_plugins::FileLocalOp;
                     let len = mint.payload.nodes.len();
-                    for op in &plugin_ops_ref[pos] {
+                    for op in plugin_ops_ref[pos] {
                         if let FileLocalOp::Edge {
                             src_local_idx,
                             dst_local_idx,
@@ -2573,26 +2568,38 @@ fn run_job<T: Sync, R: Send>(
 /// post-build plugin pass. Project-wide plugins (including project-wide
 /// external dylibs) and non-native Python plugins are skipped here —
 /// they still run in [`collect_prepared_plugin_ops`].
-fn extract_per_file_plugin_ids(
-    py: Python<'_>,
-    plugins: &[PyObject],
-) -> Vec<crate::native_plugins::PerFilePluginId> {
+/// Collect the registered per-file plugins (in registration order) and intern
+/// the ordered id list to a single process-stable `set_id`. `None` when no
+/// per-file plugins are registered (the build skips the per-file pass). The
+/// whole set then rides one salsa query per file keyed on this `set_id` — see
+/// [`crate::native_plugins::per_file_plugin_ops`].
+fn extract_per_file_plugin_set(py: Python<'_>, plugins: &[PyObject]) -> Option<u32> {
     use crate::native_plugins::{NativePlugin, NativePluginKind, PerFilePluginId};
-    let mut ids = Vec::new();
+    let mut ids: Vec<PerFilePluginId> = Vec::new();
+    let mut seen: FxHashSet<PerFilePluginId> = FxHashSet::default();
+    // Dedup by id, preserving first-occurrence (registration) order. Two
+    // plugins with the same id — e.g. identical `ServerConfig` configs that
+    // intern to one `Configured` id, or two `main_block()`s — collapse to one
+    // run, exactly as the old `(file, id)` salsa key deduped them.
+    let mut push = |id: PerFilePluginId, ids: &mut Vec<PerFilePluginId>| {
+        if seen.insert(id) {
+            ids.push(id);
+        }
+    };
     for p in plugins {
         let Ok(native) = p.bind(py).downcast::<NativePlugin>() else {
             continue;
         };
         match &native.borrow().kind {
-            NativePluginKind::PerFile(id) => ids.push(*id),
+            NativePluginKind::PerFile(id) => push(*id, &mut ids),
             NativePluginKind::External {
                 per_file_id: Some(eid),
                 ..
-            } => ids.push(PerFilePluginId::External(*eid)),
+            } => push(PerFilePluginId::External(*eid), &mut ids),
             _ => {}
         }
     }
-    ids
+    (!ids.is_empty()).then(|| crate::native_plugins::register_per_file_set(ids))
 }
 
 /// Owned `(module, name, anchor-dir)` → resolved class-base memo.
@@ -3398,7 +3405,7 @@ impl ProjectContext {
         let counters = Arc::clone(&slf.borrow(py).progress);
         let build_result = {
             let mut this = slf.borrow_mut(py);
-            let per_file_ids = extract_per_file_plugin_ids(py, &this.plugins);
+            let per_file_set_id = extract_per_file_plugin_set(py, &this.plugins);
             // Hand the previous build's resolve cache + the
             // accumulated change scope to the new build. The cache
             // moves out (the old outputs keep an invalid Default);
@@ -3422,7 +3429,7 @@ impl ProjectContext {
                 show_progress,
                 stack_size,
                 &counters,
-                &per_file_ids,
+                per_file_set_id,
                 carry,
                 cache,
                 scope,