Share prepared analysis ownership for queries

crates/no-mistakes/src/codebase/analysis_session.rs

@@ -25,12 +25,22 @@ pub struct AnalysisSession {
     datasets: DashMap<PathBuf, Arc<DatasetCell>>,
     supplemental_sources: Arc<SourceStore>,
     resolver_caches: DashMap<ResolverCacheScopeKey, Arc<ResolverResultCache>>,
+    registry_extension_reports: DashMap<RegistryExtensionKey, RegistryExtensionCell>,
     parse_attempts: Option<DashMap<PathBuf, u64>>,
 }
 
 type AnalysisDataset = crate::codebase::analysis_dataset::AnalysisDataset;
 type DatasetCell = OnceLock<Arc<AnalysisDataset>>;
 type SourceReadResult = Result<Arc<str>, SourceReadError>;
+type RegistryExtensionResult =
+    Result<Arc<crate::registry_extension_query::RegistryExtensionReport>, Arc<str>>;
+type RegistryExtensionCell = Arc<OnceLock<RegistryExtensionResult>>;
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+struct RegistryExtensionKey {
+    root: PathBuf,
+    path: PathBuf,
+}
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct SourceReadError {
@@ -64,6 +74,7 @@ impl AnalysisSession {
             datasets: DashMap::new(),
             supplemental_sources,
             resolver_caches: DashMap::new(),
+            registry_extension_reports: DashMap::new(),
             parse_attempts: collect_keyed_work.then(DashMap::new),
         })
     }
@@ -191,6 +202,38 @@ impl AnalysisSession {
             .map(|(root, cell)| self.dataset_from_cell(&root, &cell).sources_for(&root))
             .unwrap_or_else(|| Arc::clone(&self.supplemental_sources))
     }
+
+    /// Memoize a request-owned registry-extension report projection for one
+    /// root/file pair. This is not a canonical TS fact because it is the
+    /// query's rendered report, but it owns no OXC data and can therefore be
+    /// reused without parsing the same source again.
+    pub(crate) fn registry_extension_report(
+        &self,
+        root: &Path,
+        path: &Path,
+        build: impl FnOnce() -> anyhow::Result<crate::registry_extension_query::RegistryExtensionReport>,
+    ) -> anyhow::Result<crate::registry_extension_query::RegistryExtensionReport> {
+        let key = RegistryExtensionKey {
+            root: normalize_path(root),
+            path: normalize_path(path),
+        };
+        let cell = match self.registry_extension_reports.entry(key) {
+            Entry::Occupied(entry) => Arc::clone(entry.get()),
+            Entry::Vacant(entry) => {
+                let cell = Arc::new(OnceLock::new());
+                entry.insert(Arc::clone(&cell));
+                cell
+            }
+        };
+        cell.get_or_init(|| {
+            build()
+                .map(Arc::new)
+                .map_err(|error| Arc::<str>::from(format!("{error:#}")))
+        })
+        .clone()
+        .map(|report| (*report).clone())
+        .map_err(|error| anyhow::anyhow!(error))
+    }
 }
 
 #[cfg(test)]

crates/no-mistakes/src/codebase/ts_source/tests/gitignore.rs

@@ -68,10 +68,6 @@ fn pass5a_visible_adapters_preserve_fixture_backed_discovery_output() {
 #[test]
 fn pass5a_public_wrappers_create_one_request_snapshot() {
     let cases = [
-        (
-            include_str!("../../../fetches/pipeline/run.rs"),
-            "pub(crate) fn run_with_base_root(",
-        ),
         (
             include_str!("../../../queue/graph.rs"),
             "pub fn analyze_project(",
@@ -92,7 +88,6 @@ fn pass5a_public_wrappers_create_one_request_snapshot() {
             include_str!("../../../react_traits/pipeline/check.rs"),
             "pub fn check_enabled(",
         ),
-        (include_str!("../../../data_pw_query.rs"), "pub fn run("),
         (include_str!("../../../ci.rs"), "pub fn impact_report("),
         (include_str!("../../../ci.rs"), "pub fn env_report("),
         (
@@ -126,6 +121,55 @@ fn pass5a_public_wrappers_create_one_request_snapshot() {
     }
 }
 
+#[test]
+fn pass5a_session_owned_queries_do_not_restart_discovery_or_source_reads() {
+    let data_pw = include_str!("../../../data_pw_query.rs");
+    let data_pw_run = function_body(data_pw, "pub fn run(");
+    let data_pw_prepared = function_body(data_pw, "pub(crate) fn run_with_session(");
+    assert_eq!(data_pw_run.matches("AnalysisSession::new").count(), 1);
+    assert!(!data_pw_run.contains("VisiblePathSnapshot::new"));
+    for forbidden in [
+        "VisiblePathSnapshot::new",
+        "load_v2_config(",
+        "discover_visible_paths(",
+    ] {
+        assert!(
+            !data_pw_prepared.contains(forbidden),
+            "data-pw prepared run: {forbidden}"
+        );
+    }
+    let data_pw_scan = include_str!("../../../data_pw_query/scan.rs");
+    assert!(data_pw_scan.contains("session.read_source(path)"));
+    assert!(!data_pw_scan.contains("std::fs::read_to_string"));
+
+    let registry = include_str!("../../../registry_extension_query.rs");
+    let registry_run = function_body(registry, "pub fn run(");
+    let registry_prepared = function_body(registry, "pub(crate) fn run_with_session(");
+    assert_eq!(registry_run.matches("AnalysisSession::new").count(), 1);
+    assert!(!registry_run.contains("std::fs::read_to_string"));
+    assert!(registry_prepared.contains(".read_source(&path)"));
+    assert!(registry_prepared.contains(".with_program"));
+
+    let fetches = include_str!("../../../fetches/pipeline/run.rs");
+    let fetch_run = function_body(fetches, "pub(crate) fn run_with_base_root(");
+    let fetch_prepared = function_body(fetches, "pub(crate) fn run_with_base_root_and_session(");
+    assert_eq!(fetch_run.matches("AnalysisSession::new").count(), 1);
+    for forbidden in [
+        "VisiblePathSnapshot::new",
+        "load_v2_config(",
+        "discover_visible_paths(",
+    ] {
+        assert!(
+            !fetch_prepared.contains(forbidden),
+            "fetches prepared run: {forbidden}"
+        );
+    }
+    let fetch_facts = include_str!("../../../fetch/file_facts.rs");
+    assert!(fetch_facts.contains(".read_source(&abs_path)"));
+    assert!(fetch_facts.contains(".with_program"));
+    assert!(!fetch_facts.contains("std::fs::read_to_string"));
+}
+
 #[test]
 fn pass5a_prepared_bodies_do_not_restart_discovery_or_config_loading() {
     let cases = [

crates/no-mistakes/src/data_pw_query.rs

@@ -22,8 +22,9 @@ use globset::{GlobBuilder, GlobSet, GlobSetBuilder};
 use rayon::prelude::*;
 use serde::Serialize;
 
+use crate::codebase::analysis_session::AnalysisSession;
 use crate::codebase::ts_source::relative_slash_path;
-use crate::config::v2::{load_v2_config_from_visible, ConfigView};
+use crate::config::v2::ConfigView;
 use crate::playwright::selectors::compile_selector_attribute_value_regex;
 
 const SOURCE_EXTENSIONS: &[&str] = &["tsx", "ts", "jsx", "js", "mts", "cts", "mjs", "cjs"];
@@ -115,15 +116,37 @@ pub fn run(
     attribute_override: &[String],
     scan_override: &[String],
     include: &DataPwInclude,
+) -> Result<DataPwReport> {
+    let session = AnalysisSession::new(crate::diagnostics::current());
+    run_with_session(
+        &session,
+        root,
+        config_path,
+        value,
+        attribute_override,
+        scan_override,
+        include,
+    )
+}
+
+/// Run the query using the caller-owned request analysis session.
+pub(crate) fn run_with_session(
+    session: &AnalysisSession,
+    root: &Path,
+    config_path: Option<&Path>,
+    value: &str,
+    attribute_override: &[String],
+    scan_override: &[String],
+    include: &DataPwInclude,
 ) -> Result<DataPwReport> {
     // VisiblePathSnapshot returns lexically normalized paths. Normalize the
     // matching boundary once too so roots containing `.`/`..` still produce
     // relative report paths and pass the visibility filter.
     let root = crate::codebase::ts_source::normalize_discovery_path(root);
-    let snapshot = crate::codebase::ts_source::VisiblePathSnapshot::new(&root);
+    let snapshot = session.visible_paths(&root);
     let visible_paths = snapshot.paths_for(&root);
     crate::invocation::check_timeout()?;
-    let config = load_v2_config_from_visible(&root, config_path, &visible_paths)?;
+    let config = session.config(&root, config_path)?;
     let view = ConfigView::new(&config);
 
     let attributes: Vec<String> = if attribute_override.is_empty() {
@@ -177,7 +200,7 @@ pub fn run(
     };
 
     let files = discover_files_from_visible_paths(&root, &visible_paths, view.skip_directories());
-    let hits = scan_files(&files, &root, &scan)?;
+    let hits = scan_files(&files, &root, &scan, session)?;
 
     let mut source: Vec<DataPwHit> = Vec::new();
     let mut test: Vec<DataPwHit> = Vec::new();

crates/no-mistakes/src/data_pw_query/scan.rs

@@ -32,9 +32,8 @@ fn discover_files_from_visible_paths(
                 .is_ok_and(|rel| !rel_path_under_skip_dir(rel, extra_skip))
         })
         .filter(|path| has_source_extension(path))
-        // `WalkDir`'s default (non-link-following) file type never reports a
-        // symlink as a file, so match that here rather than `Path::is_file`,
-        // which follows the link.
+        // Keep the scanner's non-link-following discovery semantics. Source
+        // text itself is still read only through the request SourceStore.
         .filter(|path| std::fs::symlink_metadata(path).is_ok_and(|metadata| metadata.is_file()))
         .cloned()
         .collect()
@@ -87,22 +86,24 @@ fn scan_files(
     files: &[PathBuf],
     root: &Path,
     scan: &ScanConfig<'_>,
+    session: &crate::codebase::analysis_session::AnalysisSession,
 ) -> Result<Vec<(FileKind, DataPwHit)>> {
-    scan_files_with_timeout_check(files, root, scan, &crate::invocation::check_timeout)
+    scan_files_with_timeout_check(files, root, scan, session, &crate::invocation::check_timeout)
 }
 
 fn scan_files_with_timeout_check(
     files: &[PathBuf],
     root: &Path,
     scan: &ScanConfig<'_>,
+    session: &crate::codebase::analysis_session::AnalysisSession,
     check_timeout: &(impl Fn() -> Result<()> + Sync),
 ) -> Result<Vec<(FileKind, DataPwHit)>> {
     files
         .par_iter()
         .try_fold(Vec::new, |mut hits, path| -> Result<_> {
             check_timeout()?;
             let rel = relative_slash_path(root, path);
-            hits.extend(scan_file(path, &rel, scan, check_timeout)?);
+            hits.extend(scan_file(path, &rel, scan, session, check_timeout)?);
             Ok(hits)
         })
         .try_reduce(Vec::new, |mut left, mut right| -> Result<_> {
@@ -115,6 +116,7 @@ fn scan_file(
     path: &Path,
     rel: &str,
     scan: &ScanConfig,
+    session: &crate::codebase::analysis_session::AnalysisSession,
     check_timeout: &(impl Fn() -> Result<()> + Sync),
 ) -> Result<Vec<(FileKind, DataPwHit)>> {
     check_timeout()?;
@@ -141,7 +143,7 @@ fn scan_file(
     } else {
         FileKind::Source
     };
-    let Ok(source) = std::fs::read_to_string(path) else {
+    let Ok(source) = session.read_source(path) else {
         return Ok(Vec::new());
     };
     let mut hits = Vec::new();

crates/no-mistakes/src/data_pw_query/tests.rs

@@ -142,6 +142,26 @@ fn finds_source_and_test_usages() {
     assert_eq!(test[0].attribute, "data-pw");
 }
 
+#[test]
+fn prepared_run_reuses_one_visible_inventory_and_source_store() {
+    let observer = crate::diagnostics::InvocationObserver::new(true);
+    let session = crate::codebase::analysis_session::AnalysisSession::new(Some(observer));
+    let include = DataPwInclude::default();
+
+    let first =
+        run_with_session(&session, &fixture(), None, "search-bar", &[], &[], &include).unwrap();
+    let second =
+        run_with_session(&session, &fixture(), None, "search-bar", &[], &[], &include).unwrap();
+
+    assert_eq!(first, second);
+    let work = session.work_snapshot();
+    assert!(!work.source_reads.is_empty());
+    assert!(
+        work.source_reads.values().all(|count| *count == 1),
+        "source reads must be memoized by the request session: {work:?}"
+    );
+}
+
 #[test]
 fn expired_deadline_returns_timeout_instead_of_a_partial_report() {
     let _deadline = crate::invocation::install_test_deadline(std::time::Duration::ZERO).unwrap();
@@ -268,6 +288,7 @@ fn is_skip_dir_honors_defaults_and_config() {
 
 #[test]
 fn scan_file_ignores_unreadable_path() {
+    let session = crate::codebase::analysis_session::AnalysisSession::disabled();
     let regex = compile_selector_attribute_value_regex(&["data-pw".to_string()]).unwrap();
     let globs = build_globset(&[]).unwrap();
     let scan = ScanConfig {
@@ -283,6 +304,7 @@ fn scan_file_ignores_unreadable_path() {
         Path::new("/no/such/file.tsx"),
         "x.tsx",
         &scan,
+        &session,
         &crate::invocation::check_timeout,
     )
     .unwrap();
@@ -291,6 +313,7 @@ fn scan_file_ignores_unreadable_path() {
 
 #[test]
 fn parallel_scan_propagates_a_deadline_that_expires_inside_a_worker() {
+    let session = crate::codebase::analysis_session::AnalysisSession::disabled();
     let root = fixture();
     let regex = compile_selector_attribute_value_regex(&["data-pw".to_string()]).unwrap();
     let globs = build_globset(&[]).unwrap();
@@ -321,6 +344,7 @@ fn parallel_scan_propagates_a_deadline_that_expires_inside_a_worker() {
         &[root.join("app/search.tsx")],
         &root,
         &scan,
+        &session,
         &expire_mid_scan,
     )
     .unwrap_err();

crates/no-mistakes/src/fetch/file_analysis.rs

@@ -55,6 +55,7 @@ pub(crate) fn analyze_file_from_visible(
 }
 
 pub(crate) struct VisibleFileAnalysis<'a> {
+    pub session: &'a crate::codebase::analysis_session::AnalysisSession,
     pub root: &'a Path,
     pub visited: &'a mut HashSet<(PathBuf, bool, bool)>,
     pub fetches: &'a mut Vec<FetchOccurrence>,
@@ -69,6 +70,7 @@ pub(crate) fn analyze_file_from_visible_with_facts(
     context: &mut VisibleFileAnalysis<'_>,
 ) -> Result<bool> {
     let VisibleFileAnalysis {
+        session,
         root,
         visited,
         fetches,
@@ -94,7 +96,13 @@ pub(crate) fn analyze_file_from_visible_with_facts(
         return Ok(cached_fetches.is_client);
     }
 
-    let facts = parsed_files.load(&abs_path, root, &mut cache.imports, visible_files)?;
+    let facts = parsed_files.load_with_session(
+        session,
+        &abs_path,
+        root,
+        &mut cache.imports,
+        visible_files,
+    )?;
     let is_client = !inherited_is_route_handler
         && !facts.has_use_server_directive
         && (inherited_is_client || facts.has_use_client_directive);
@@ -112,6 +120,7 @@ pub(crate) fn analyze_file_from_visible_with_facts(
             &import,
             (is_client, inherited_is_route_handler),
             &mut VisibleFileAnalysis {
+                session,
                 root,
                 visited,
                 fetches: &mut file_fetches,

crates/no-mistakes/src/fetch/file_facts.rs

@@ -25,8 +25,9 @@ pub struct ParsedFileCache {
 }
 
 impl ParsedFileCache {
-    pub(crate) fn load(
+    pub(crate) fn load_with_session(
         &mut self,
+        session: &crate::codebase::analysis_session::AnalysisSession,
         path: &Path,
         root: &Path,
         import_cache: &mut HashMap<PathBuf, Vec<PathBuf>>,
@@ -38,8 +39,10 @@ impl ParsedFileCache {
         }
 
         let result = (|| {
-            let source = std::fs::read_to_string(&abs_path)?;
-            crate::ast::with_program(&abs_path, &source, |program, _| {
+            let source = session
+                .read_source(&abs_path)
+                .map_err(|error| anyhow::anyhow!(error.to_string()))?;
+            session.with_program(&abs_path, &source, |program, _| {
                 ParsedFileFacts::from_program(
                     &abs_path,
                     root,