diff --git a/Cargo.lock b/Cargo.lock index 31b5bf573..d15232ef9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4945,6 +4945,7 @@ dependencies = [ "tree-sitter-xml", "tree-sitter-yaml", "tree-sitter-zig", + "xxhash-rust", ] [[package]] diff --git a/crates/pi-ast/Cargo.toml b/crates/pi-ast/Cargo.toml index 34c7c363e..06e27992b 100644 --- a/crates/pi-ast/Cargo.toml +++ b/crates/pi-ast/Cargo.toml @@ -73,3 +73,4 @@ tree-sitter-vue.workspace = true tree-sitter-xml.workspace = true tree-sitter-yaml.workspace = true tree-sitter-zig.workspace = true +xxhash-rust.workspace = true diff --git a/crates/pi-ast/src/block.rs b/crates/pi-ast/src/block.rs index 56e5c7bab..7acb075cf 100644 --- a/crates/pi-ast/src/block.rs +++ b/crates/pi-ast/src/block.rs @@ -10,12 +10,14 @@ use std::collections::BTreeSet; -use anyhow::{Result, anyhow}; -use ast_grep_core::tree_sitter::LanguageExt; +use anyhow::Result; use serde::{Deserialize, Serialize}; -use tree_sitter::{Parser, Point, TreeCursor}; +use tree_sitter::{Point, TreeCursor}; -use crate::summary::{node_content_end_line, node_start_line, resolve_language}; +use crate::{ + parse_cache::parse_cached, + summary::{node_content_end_line, node_start_line, resolve_language}, +}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BlockRangeOptions { @@ -68,11 +70,7 @@ pub fn block_range_at(options: BlockRangeOptions) -> Result> return Ok(None); }; - let mut parser = Parser::new(); - parser - .set_language(&language.get_ts_language()) - .map_err(|err| anyhow!("Failed to load tree-sitter language: {err}"))?; - let Some(tree) = parser.parse(&code, None) else { + let Some(tree) = parse_cached(&code, language)? else { return Ok(None); }; let root = tree.root_node(); @@ -181,11 +179,31 @@ fn is_visible(merged: &[LineRange], line: u32) -> bool { .is_ok() } +/// Does any visible line fall inside the inclusive line span `[start, end]`? +fn intersects_visible(merged: &[LineRange], start: u32, end: u32) -> bool { + // `merged` is sorted and non-overlapping, so the first range that can + // possibly overlap is the first one whose `end_line` reaches `start`. + let idx = merged.partition_point(|range| range.end_line < start); + merged.get(idx).is_some_and(|range| range.start_line <= end) +} + /// Depth-first walk collecting boundary lines from every multi-line named node /// that straddles a visible-range edge. A single reused [`TreeCursor`] keeps /// the traversal allocation-free. fn collect_boundaries(cursor: &mut TreeCursor<'_>, merged: &[LineRange], out: &mut BTreeSet) { let node = cursor.node(); + // Prune whole subtrees that cannot contribute. A node contributes only when + // one of its own endpoint lines is visible, and both of those lines lie + // inside its raw row span; every descendant's span is contained in this + // one, so a span holding no visible line rules out this node *and* + // everything beneath it. Without the prune the walk is O(nodes in file) + // even though the answer is bounded by the window size — which made the + // traversal cost roughly twice the parse on a large file. + let raw_start = node.start_position().row.saturating_add(1) as u32; + let raw_end = node.end_position().row.saturating_add(1) as u32; + if !intersects_visible(merged, raw_start, raw_end) { + return; + } // Skip the whole-file root: its only "boundary" is EOF, never a useful // matching line (mirrors `block_range_at` excluding the root). if node.is_named() && node.parent().is_some() { @@ -240,11 +258,7 @@ pub fn enclosing_block_boundaries(options: EnclosingBoundaryOptions) -> Result = (0..MAX_ENTRIES * 2) + .map(|i| format!("function f{i}() {{\n{}}}\n", " step();\n".repeat(i + 1))) + .collect(); + let expect = |i: usize| Some(vec![i as u32 + 3]); + + clear_parse_cache(); + for (i, source) in sources.iter().enumerate() { + assert_eq!(boundaries(source, "x.ts", &[(1, 1)]), expect(i), "cold pass {i}"); + } + + // Source 0 was evicted well before the loop ended (twice `MAX_ENTRIES` + // distinct sources went through a cache holding `MAX_ENTRIES`), so this + // re-parses. It must answer correctly rather than serve a survivor. + assert_eq!(boundaries(&sources[0], "x.ts", &[(1, 1)]), expect(0), "evicted key"); + // And the tail, still resident, must not have been disturbed. + let last = sources.len() - 1; + assert_eq!(boundaries(&sources[last], "x.ts", &[(1, 1)]), expect(last), "resident key"); + } + + // ── prune equivalence ───────────────────────────────────────────────────── + // + // `collect_boundaries` skips subtrees whose raw line span holds no visible + // line. That is argued to be exact, but the output feeds hashline block + // resolution, so a silently dropped line corrupts edits rather than just + // degrading display. The argument is therefore backed by a differential + // against the pre-prune traversal over the repository's own sources, not by + // hand-written expectations. + + use std::path::{Path, PathBuf}; + + /// The pre-prune traversal, kept verbatim as the differential reference: it + /// visits every node in the tree. + fn collect_boundaries_unpruned( + cursor: &mut TreeCursor<'_>, + merged: &[LineRange], + out: &mut BTreeSet, + ) { + let node = cursor.node(); + if node.is_named() && node.parent().is_some() { + let start = node_start_line(node); + let end = node_content_end_line(node); + if end > start { + let start_visible = is_visible(merged, start); + let end_visible = is_visible(merged, end); + if start_visible && !end_visible { + out.insert(end); + } else if end_visible && !start_visible { + out.insert(start); + } + } + } + if cursor.goto_first_child() { + loop { + collect_boundaries_unpruned(cursor, merged, out); + if !cursor.goto_next_sibling() { + break; + } + } + cursor.goto_parent(); + } + } + + /// [`enclosing_block_boundaries`] with the unpruned walk substituted in. + /// Every early return is reproduced in the same order, so the differential + /// also covers the empty-code, empty-range, unresolved-language and + /// `has_error` paths. + fn boundaries_unpruned(code: &str, path: &str, ranges: &[(u32, u32)]) -> Option> { + let merged = normalize_ranges( + ranges + .iter() + .map(|&(start_line, end_line)| LineRange { start_line, end_line }) + .collect(), + ); + if code.is_empty() || merged.is_empty() { + return Some(Vec::new()); + } + let language = resolve_language(None, Some(path))?; + let tree = parse_cached(code, language).expect("grammar loads")?; + let root = tree.root_node(); + if root.has_error() { + return None; + } + let mut boundaries = BTreeSet::new(); + let mut cursor = root.walk(); + collect_boundaries_unpruned(&mut cursor, &merged, &mut boundaries); + Some(boundaries.into_iter().collect()) + } + + /// Range shapes exercised per file: head, middle, tail, a single interior + /// line, three disjoint windows, the whole file visible, a window entirely + /// past EOF, and the empty range list. + fn window_shapes(lines: u32) -> Vec> { + let last = lines.max(1); + let mid = (lines / 2).max(1); + let tail_start = lines.saturating_sub(20).max(1); + vec![ + vec![(1, 40.min(last))], + vec![(mid, (mid + 20).min(last))], + vec![(tail_start, last)], + vec![(mid, mid)], + vec![(1, 5.min(last)), (mid, (mid + 5).min(last)), (tail_start, last)], + vec![(1, last)], + vec![(last + 10, last + 20)], + vec![], + ] + } + + /// Compare pruned against unpruned output for exact `Option>` + /// equality across every shape. Returns (comparisons, `None` verdicts). + fn assert_prune_equivalent(code: &str, path: &str) -> (usize, usize) { + let lines = code.split('\n').count() as u32; + let mut comparisons = 0; + let mut none_verdicts = 0; + for shape in window_shapes(lines) { + let pruned = boundaries(code, path, &shape); + let unpruned = boundaries_unpruned(code, path, &shape); + assert_eq!(pruned, unpruned, "{path} ranges={shape:?}"); + if pruned.is_none() { + none_verdicts += 1; + } + comparisons += 1; + } + (comparisons, none_verdicts) + } + + /// Repository `.ts` / `.py` / `.rs` sources, sorted for determinism. With a + /// budget, files are taken on a stride so the sample spans the whole tree + /// instead of one directory, skipping anything over 64 KiB. + fn repo_files(byte_budget: Option) -> Vec { + let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let mut files: Vec = ignore::WalkBuilder::new(&root) + .build() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_some_and(|kind| kind.is_file())) + .map(ignore::DirEntry::into_path) + .filter(|path| { + matches!(path.extension().and_then(std::ffi::OsStr::to_str), Some("ts" | "py" | "rs")) + }) + .collect(); + files.sort(); + let Some(budget) = byte_budget else { + return files; + }; + let stride = (files.len() / 240).max(1); + let mut picked = Vec::new(); + let mut used = 0; + for path in files.iter().step_by(stride) { + let Ok(meta) = std::fs::metadata(path) else { + continue; + }; + let len = meta.len() as usize; + if len == 0 || len > 64 * 1024 { + continue; + } + if used + len > budget { + break; + } + used += len; + picked.push(path.clone()); + } + picked + } + + fn sweep_corpus(files: &[PathBuf]) -> (usize, usize) { + let mut comparisons = 0; + let mut none_verdicts = 0; + for path in files { + let Ok(code) = std::fs::read_to_string(path) else { + continue; // non-UTF-8 source; nothing to compare + }; + let (n, nones) = assert_prune_equivalent(&code, path.to_str().expect("utf-8 path")); + comparisons += n; + none_verdicts += nones; + } + (comparisons, none_verdicts) + } + + #[test] + fn pruned_walk_matches_unpruned_on_repo_corpus_sample() { + let files = repo_files(Some(768 * 1024)); + assert!(files.len() > 40, "corpus sample too small to be evidence: {}", files.len()); + let (comparisons, _) = sweep_corpus(&files); + assert!(comparisons > 300, "expected a broad sweep, got {comparisons} comparisons"); + } + + #[test] + #[ignore = "full repository sweep; run with `cargo test --release -p pi-ast -- --ignored`"] + fn pruned_walk_matches_unpruned_on_full_repo_corpus() { + let files = repo_files(None); + assert!(files.len() > 3000, "expected the whole corpus, got {}", files.len()); + let (comparisons, _) = sweep_corpus(&files); + println!("full corpus: {} files, {comparisons} comparisons", files.len()); + } + + #[test] + fn pruned_walk_matches_unpruned_on_error_and_degenerate_inputs() { + // The corpus is mostly valid source, so pin the `has_error`, + // empty-code, empty-range and unresolved-language paths explicitly. + let cases: &[(&str, &str)] = &[ + ("", "x.ts"), + ("\n", "x.ts"), + ("\n\n\n", "x.py"), + ("function broken() {\n if (y) {\n", "b.ts"), + ("def broken(:\n pass\n", "b.py"), + ("fn broken( {\n", "b.rs"), + ("function ok() {\n a();\n}\nfunction broken( {\n", "x.ts"), + (TS_FN, "x.unknownext"), + ]; + let mut none_verdicts = 0; + for (code, path) in cases { + let (_, nones) = assert_prune_equivalent(code, path); + none_verdicts += nones; + } + assert!(none_verdicts > 0, "error / unknown-language cases must exercise the None path"); + } + + /// Large file with a five-level nest in the middle, so a window on the + /// opening lines has a long chain of closers to surface and the prune has + /// most of the file to skip. Returns the source and the 1-indexed line of + /// `function deep() {`. + fn nested_fixture(pad: usize) -> (String, u32) { + use std::fmt::Write as _; + + let mut code = String::new(); + for i in 0..pad { + writeln!(code, "export const pad{i} = {i};").expect("write to String"); + } + let nest = pad as u32 + 1; + code.push_str("function deep() {\n"); + code.push_str("\tif (a) {\n"); + code.push_str("\t\twhile (b) {\n"); + code.push_str("\t\t\tfor (;;) {\n"); + code.push_str("\t\t\t\tif (c) {\n"); + code.push_str("\t\t\t\t\tbody();\n"); + code.push_str("\t\t\t\t}\n"); + code.push_str("\t\t\t}\n"); + code.push_str("\t\t}\n"); + code.push_str("\t}\n"); + code.push_str("}\n"); + for i in 0..pad { + writeln!(code, "export const tail{i} = {i};").expect("write to String"); + } + (code, nest) + } + + #[test] + fn window_deep_inside_nesting_still_returns_the_whole_chain() { + let (code, nest) = nested_fixture(200); + + // Window covers all five opening lines, which are deep inside a + // ~411-line file: each construct opens visibly and closes off-window, so + // the full chain of five closers must come back. + let openers = [(nest, nest + 4)]; + assert_eq!( + boundaries(&code, "nested.ts", &openers), + Some(vec![nest + 6, nest + 7, nest + 8, nest + 9, nest + 10]), + "closers for every enclosing construct" + ); + assert_eq!( + boundaries(&code, "nested.ts", &openers), + boundaries_unpruned(&code, "nested.ts", &openers) + ); + + // Mirror image: the five closing lines surface the five openers. + let closers = [(nest + 6, nest + 10)]; + assert_eq!( + boundaries(&code, "nested.ts", &closers), + Some(vec![nest, nest + 1, nest + 2, nest + 3, nest + 4]) + ); + assert_eq!( + boundaries(&code, "nested.ts", &closers), + boundaries_unpruned(&code, "nested.ts", &closers) + ); + + // A window strictly interior to the nest surfaces nothing — every + // enclosing node straddles it, so no endpoint is visible. The prune must + // still descend through those straddling ancestors, which is what the + // differential over every shape checks. + assert_eq!(boundaries(&code, "nested.ts", &[(nest + 5, nest + 5)]), Some(vec![])); + assert_prune_equivalent(&code, "nested.ts"); + } + + /// Count named nodes whose content end line differs from their raw end row, + /// i.e. nodes ending immediately after a newline. + fn nodes_with_content_end_before_raw_end(code: &str, path: &str) -> usize { + let language = resolve_language(None, Some(path)).expect("known language"); + let tree = parse_cached(code, language) + .expect("grammar loads") + .expect("tree"); + let mut count = 0; + let mut stack = vec![tree.root_node()]; + while let Some(node) = stack.pop() { + let raw_end = node.end_position().row.saturating_add(1) as u32; + if node.is_named() && node_content_end_line(node) != raw_end { + count += 1; + } + for index in 0..node.child_count() { + stack.push(node.child(index).expect("child in range")); + } + } + count + } + + #[test] + fn content_end_line_before_raw_end_is_still_surfaced() { + // The inner `def` ends right after `return value\n`, so tree-sitter puts + // its end position at column 0 of the blank line: raw end row 5, content + // end line 4. The prune tests the *raw* span, so the node survives a + // window on either line. + let code = "def outer():\n def inner():\n value = 1\n return value\n\n \ + return inner\n"; + assert!( + nodes_with_content_end_before_raw_end(code, "x.py") > 0, + "fixture must actually contain a node whose content end precedes its raw end" + ); + + // Only the content end line is visible: the inner def and its block both + // close there, so both openers come back. + assert_eq!(boundaries(code, "x.py", &[(4, 4)]), Some(vec![2, 3])); + assert_eq!(boundaries(code, "x.py", &[(4, 4)]), boundaries_unpruned(code, "x.py", &[(4, 4)])); + + // And the raw end line (the blank one) on its own. + assert_eq!(boundaries(code, "x.py", &[(5, 5)]), boundaries_unpruned(code, "x.py", &[(5, 5)])); + assert_prune_equivalent(code, "x.py"); + } } diff --git a/crates/pi-ast/src/lib.rs b/crates/pi-ast/src/lib.rs index 971081275..10285fc1b 100644 --- a/crates/pi-ast/src/lib.rs +++ b/crates/pi-ast/src/lib.rs @@ -1,6 +1,7 @@ pub mod block; pub mod language; pub mod ops; +pub mod parse_cache; pub mod summary; pub use language::SupportLang; diff --git a/crates/pi-ast/src/parse_cache.rs b/crates/pi-ast/src/parse_cache.rs new file mode 100644 index 000000000..294d76148 --- /dev/null +++ b/crates/pi-ast/src/parse_cache.rs @@ -0,0 +1,433 @@ +//! Bounded, content-addressed tree-sitter parse cache. +//! +//! Every structural entry point in this crate ([`crate::block`], +//! [`crate::summary`]) is dominated by one cost: `Parser::parse` over the whole +//! file. Measured on an M4 Max (release, warm page cache) that is ~13.5 ms for +//! an 81 KB / 2057-line TypeScript file and ~187 ms for a 1.05 MB file, while +//! everything else those functions do totals well under a millisecond. +//! +//! The results themselves are not cacheable: `enclosing_block_boundaries` +//! depends on the caller's visible `ranges`, which differ on every call. The +//! reusable artifact is the [`Tree`], so the cache lives here and hands out +//! cheap clones of it. +//! +//! `Tree` is `Send` but not `Sync`, so entries live behind a [`Mutex`] and the +//! lock is only ever held for a map probe, a byte comparison, and a +//! `ts_tree_copy` refcount bump — never across a parse or a tree walk. + +use std::{ + collections::HashMap, + sync::{LazyLock, Mutex, MutexGuard, PoisonError}, +}; + +use anyhow::{Result, anyhow}; +use ast_grep_core::tree_sitter::LanguageExt; +use tree_sitter::{Parser, Tree}; + +use crate::language::SupportLang; + +/// Arbitrary fixed seed (golden-ratio constant). Fixed, not random, so a key is +/// reproducible across calls within a process; it never leaves the process, so +/// there is nothing to harden against `HashDoS` here. +const HASH_SEED: u64 = 0x9e37_79b9_7f4a_7c15; + +/// Largest source that may occupy a slot. +/// +/// A tree-sitter tree runs roughly an order of magnitude larger than its +/// source, so admitting an arbitrarily large file would let one `read` of a +/// multi-megabyte blob dominate process RSS. Files above this are still parsed, +/// just never retained. +pub const MAX_ENTRY_SOURCE_BYTES: usize = 4 << 20; + +/// Ceiling on retained source bytes across all slots. +/// +/// Equal to [`MAX_ENTRY_SOURCE_BYTES`] so a single hot large file can still be +/// cached (it evicts everything else, which is what LRU should do when that +/// file *is* the working set). +pub const MAX_TOTAL_SOURCE_BYTES: usize = 4 << 20; + +/// Slot ceiling, independent of byte size. +/// +/// Bounds the tree footprint against a burst of small files. Twelve covers the +/// realistic hot set for a coding agent (the handful of files being read and +/// edited) while keeping the LRU scan trivially cheap. +pub const MAX_ENTRIES: usize = 12; + +/// Cache key. The 64-bit hash is a *bucket selector only*: a hit additionally +/// verifies [`Entry::source`] against the request byte-for-byte before the tree +/// is handed back, so a hash collision can only ever cost a re-parse (the +/// colliding slot is overwritten) and can never return a tree built from +/// different content. `len` is folded in because it is free and makes +/// accidental bucket sharing rarer; `lang` is in the key because the same bytes +/// parsed as TypeScript and as Python are different trees and must not share a +/// slot. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct Key { + hash: u64, + len: usize, + lang: SupportLang, +} + +fn key_for(code: &str, lang: SupportLang) -> Key { + Key { hash: xxhash_rust::xxh64::xxh64(code.as_bytes(), HASH_SEED), len: code.len(), lang } +} + +struct Entry { + /// Retained verbatim so a hit is verified by comparison, not by trusting + /// the hash. + source: Box, + tree: Tree, + /// Value of [`Cache::clock`] at last use; smallest wins eviction. + stamp: u64, +} + +struct Cache { + entries: HashMap, + source_bytes: usize, + clock: u64, + hits: u64, + misses: u64, + evictions: u64, + max_entries: usize, + max_total_bytes: usize, + max_entry_bytes: usize, +} + +impl Cache { + fn new(max_entries: usize, max_total_bytes: usize, max_entry_bytes: usize) -> Self { + Self { + entries: HashMap::new(), + source_bytes: 0, + clock: 0, + hits: 0, + misses: 0, + evictions: 0, + max_entries, + max_total_bytes, + max_entry_bytes, + } + } + + fn get(&mut self, key: &Key, code: &str) -> Option { + self.clock += 1; + let stamp = self.clock; + let tree = match self.entries.get_mut(key) { + Some(entry) if &*entry.source == code => { + entry.stamp = stamp; + // `ts_tree_copy`: an atomic refcount bump on immutable subtree + // data, so the clone can be walked off-lock on any thread. + entry.tree.clone() + }, + _ => { + self.misses += 1; + return None; + }, + }; + self.hits += 1; + Some(tree) + } + + fn insert(&mut self, key: Key, code: &str, tree: &Tree) { + if code.len() > self.max_entry_bytes { + return; + } + if let Some(previous) = self.entries.remove(&key) { + self.source_bytes -= previous.source.len(); + } + while self.entries.len() >= self.max_entries + || self.source_bytes + code.len() > self.max_total_bytes + { + if !self.evict_oldest() { + break; + } + } + self.clock += 1; + self.source_bytes += code.len(); + self.entries.insert(key, Entry { + source: Box::from(code), + tree: tree.clone(), + stamp: self.clock, + }); + } + + /// Drop the least-recently-used slot. `false` when there was nothing left + /// to drop, which is what terminates [`Self::insert`]'s eviction loop. + fn evict_oldest(&mut self) -> bool { + // Linear over at most `max_entries` slots: cheaper than maintaining an + // intrusive LRU list at this size. + let Some(&oldest) = self + .entries + .iter() + .min_by_key(|(_, entry)| entry.stamp) + .map(|(key, _)| key) + else { + return false; + }; + if let Some(entry) = self.entries.remove(&oldest) { + self.source_bytes -= entry.source.len(); + self.evictions += 1; + } + true + } + + /// Drop every entry and zero the counters, preserving the configured bounds. + fn clear(&mut self) { + self.entries.clear(); + self.source_bytes = 0; + self.clock = 0; + self.hits = 0; + self.misses = 0; + self.evictions = 0; + } +} + +static CACHE: LazyLock> = LazyLock::new(|| { + Mutex::new(Cache::new(MAX_ENTRIES, MAX_TOTAL_SOURCE_BYTES, MAX_ENTRY_SOURCE_BYTES)) +}); + +/// Every critical section is a handful of infallible map operations plus a +/// refcount bump, so panicking while holding the lock is not reachable. +/// Recovering the guard anyway means a hypothetical panic could never escalate +/// into every later parse panicking. +fn lock() -> MutexGuard<'static, Cache> { + CACHE.lock().unwrap_or_else(PoisonError::into_inner) +} + +/// Parse `code` as `lang`, reusing a cached [`Tree`] when the exact same bytes +/// were last parsed as the same language and have not been evicted. +/// +/// Semantics match a bare `Parser::new()` / `set_language` / `parse` sequence +/// exactly: `Err` when the grammar fails to load, `Ok(None)` when `parse` +/// yields nothing, `Ok(Some(tree))` otherwise. Trees carrying syntax errors are +/// cached like any other — `has_error()` is a property of the tree, so callers +/// that reject on it reach the identical verdict from a cached tree, and +/// repeated "does this parse" probes over the same broken file get the speedup +/// too. +pub fn parse_cached(code: &str, lang: SupportLang) -> Result> { + let key = key_for(code, lang); + // Bound the guard to a `let` so it drops at the end of this statement: an + // `if let` scrutinee would hold the lock across the early return. + let cached = lock().get(&key, code); + if let Some(tree) = cached { + return Ok(Some(tree)); + } + let mut parser = Parser::new(); + parser + .set_language(&lang.get_ts_language()) + .map_err(|err| anyhow!("Failed to load tree-sitter language: {err}"))?; + let Some(tree) = parser.parse(code, None) else { + return Ok(None); + }; + lock().insert(key, code, &tree); + Ok(Some(tree)) +} + +/// Occupancy and counters, for diagnostics. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ParseCacheStats { + pub entries: usize, + pub source_bytes: usize, + pub hits: u64, + pub misses: u64, + pub evictions: u64, +} + +pub fn parse_cache_stats() -> ParseCacheStats { + let cache = lock(); + ParseCacheStats { + entries: cache.entries.len(), + source_bytes: cache.source_bytes, + hits: cache.hits, + misses: cache.misses, + evictions: cache.evictions, + } +} + +/// Drop every cached tree and zero the counters. +pub fn clear_parse_cache() { + lock().clear(); +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The process-global cache is shared with every other test in this crate + /// running in parallel, so occupancy, eviction, and counter behaviour are + /// asserted on private instances instead. `block.rs` covers the global path + /// end to end, where the assertions are boundary values that hold no matter + /// what else is resident. + fn cache(max_entries: usize, max_total_bytes: usize, max_entry_bytes: usize) -> Cache { + Cache::new(max_entries, max_total_bytes, max_entry_bytes) + } + + fn tree_of(code: &str, lang: SupportLang) -> Tree { + let mut parser = Parser::new(); + parser + .set_language(&lang.get_ts_language()) + .expect("grammar loads"); + parser.parse(code, None).expect("parse produces a tree") + } + + const TS: SupportLang = SupportLang::TypeScript; + + #[test] + fn hit_returns_a_clone_of_the_stored_tree() { + let mut cache = cache(MAX_ENTRIES, MAX_TOTAL_SOURCE_BYTES, MAX_ENTRY_SOURCE_BYTES); + let code = "function f() {\n a();\n}\n"; + let key = key_for(code, TS); + let tree = tree_of(code, TS); + + assert!(cache.get(&key, code).is_none(), "cold lookup misses"); + cache.insert(key, code, &tree); + let hit = cache.get(&key, code).expect("warm lookup hits"); + + // A child node's id is the address of its shared subtree data, so a + // matching id proves the hit is a `ts_tree_copy` of the stored tree + // rather than a second parse. (The *root* node's id is the address of + // the `root` field inside the `TSTree` struct itself, which a copy + // necessarily relocates, so it cannot be used here.) + let subtree_id = |tree: &Tree| tree.root_node().child(0).expect("top-level decl").id(); + assert_eq!(subtree_id(&hit), subtree_id(&tree), "hit must share the stored tree"); + // The id genuinely discriminates: an independent parse of the same bytes + // allocates fresh subtree data. + assert_ne!(subtree_id(&tree_of(code, TS)), subtree_id(&tree), "id must discriminate"); + assert_eq!((cache.hits, cache.misses, cache.entries.len()), (1, 1, 1)); + assert_eq!(cache.source_bytes, code.len()); + } + + #[test] + fn language_is_part_of_the_key() { + let code = "def greet(name):\n return name\n"; + let py_key = key_for(code, SupportLang::Python); + let ts_key = key_for(code, TS); + assert_ne!(py_key, ts_key, "same bytes, different language, different key"); + + let mut cache = cache(MAX_ENTRIES, MAX_TOTAL_SOURCE_BYTES, MAX_ENTRY_SOURCE_BYTES); + cache.insert(py_key, code, &tree_of(code, SupportLang::Python)); + assert!( + cache.get(&ts_key, code).is_none(), + "a Python tree must never satisfy a TypeScript request" + ); + } + + #[test] + fn one_byte_of_difference_changes_the_key() { + // Equal length, one byte apart, so `len` cannot separate them. + let a = "const a = 1;\n"; + let b = "const a = 2;\n"; + assert_eq!(a.len(), b.len()); + assert_ne!(key_for(a, TS), key_for(b, TS)); + } + + #[test] + fn verified_equality_rejects_a_colliding_slot() { + // A real xxh64 collision is not findable, so forge the condition: store + // `stored` under the key that `probe` hashes to. A cache that trusted + // the hash would hand back `stored`'s tree for a `probe` lookup. + let stored = "const a = 1;\n"; + let probe = "const b = 2;\n"; + assert_eq!(stored.len(), probe.len(), "collision needs a matching len field"); + + let forged = key_for(probe, TS); + let mut cache = cache(MAX_ENTRIES, MAX_TOTAL_SOURCE_BYTES, MAX_ENTRY_SOURCE_BYTES); + cache.insert(forged, stored, &tree_of(stored, TS)); + + assert!( + cache.get(&forged, probe).is_none(), + "byte comparison must turn a collision into a miss" + ); + assert_eq!(cache.misses, 1); + } + + #[test] + fn entry_bound_evicts_least_recently_used() { + let mut cache = cache(2, MAX_TOTAL_SOURCE_BYTES, MAX_ENTRY_SOURCE_BYTES); + let sources = ["const a = 1;\n", "const b = 2;\n", "const c = 3;\n"]; + let keys: Vec = sources.iter().map(|code| key_for(code, TS)).collect(); + + cache.insert(keys[0], sources[0], &tree_of(sources[0], TS)); + cache.insert(keys[1], sources[1], &tree_of(sources[1], TS)); + // Touch slot 0 so slot 1 becomes the least recently used. + assert!(cache.get(&keys[0], sources[0]).is_some()); + cache.insert(keys[2], sources[2], &tree_of(sources[2], TS)); + + assert_eq!(cache.entries.len(), 2, "entry bound holds"); + assert_eq!(cache.evictions, 1); + assert!(cache.get(&keys[0], sources[0]).is_some(), "recently used slot survives"); + assert!(cache.get(&keys[1], sources[1]).is_none(), "least recently used slot evicted"); + assert!(cache.get(&keys[2], sources[2]).is_some(), "newest slot resident"); + } + + #[test] + fn byte_bound_evicts_and_reclaims_accounting() { + let code_len = "const a = 1;\n".len(); + // Room for exactly two of these sources. + let mut cache = cache(64, code_len * 2, MAX_ENTRY_SOURCE_BYTES); + let sources = ["const a = 1;\n", "const b = 2;\n", "const c = 3;\n"]; + + for code in sources { + let key = key_for(code, TS); + cache.insert(key, code, &tree_of(code, TS)); + assert!( + cache.source_bytes <= code_len * 2, + "byte bound must hold after every insert: {} > {}", + cache.source_bytes, + code_len * 2 + ); + } + + assert_eq!(cache.entries.len(), 2); + assert_eq!(cache.evictions, 1); + // Exact accounting, not just the bound: a subtraction missed on eviction + // would leave `source_bytes` drifting upward until the cache starved. + assert_eq!(cache.source_bytes, code_len * 2); + assert!( + cache.get(&key_for(sources[0], TS), sources[0]).is_none(), + "oldest slot evicted by the byte bound" + ); + } + + #[test] + fn source_over_the_entry_bound_is_never_retained() { + let code = "function f() {\n a();\n}\n"; + let mut cache = cache(MAX_ENTRIES, MAX_TOTAL_SOURCE_BYTES, code.len() - 1); + let key = key_for(code, TS); + + cache.insert(key, code, &tree_of(code, TS)); + + assert_eq!(cache.entries.len(), 0, "oversized source is parsed but not cached"); + assert_eq!(cache.source_bytes, 0); + assert_eq!(cache.evictions, 0, "rejection must not evict a healthy slot"); + assert!(cache.get(&key, code).is_none()); + } + + #[test] + fn reinsert_of_a_resident_key_does_not_double_count_bytes() { + let mut cache = cache(MAX_ENTRIES, MAX_TOTAL_SOURCE_BYTES, MAX_ENTRY_SOURCE_BYTES); + let code = "const a = 1;\n"; + let key = key_for(code, TS); + let tree = tree_of(code, TS); + + cache.insert(key, code, &tree); + cache.insert(key, code, &tree); + + assert_eq!(cache.entries.len(), 1); + assert_eq!(cache.source_bytes, code.len()); + assert_eq!(cache.evictions, 0); + } + + #[test] + fn clear_empties_the_cache_but_keeps_its_bounds() { + let mut cache = cache(3, 1024, 512); + let code = "const a = 1;\n"; + cache.insert(key_for(code, TS), code, &tree_of(code, TS)); + assert_eq!(cache.entries.len(), 1); + + cache.clear(); + + assert_eq!(cache.entries.len(), 0); + assert_eq!((cache.source_bytes, cache.hits, cache.misses, cache.evictions), (0, 0, 0, 0)); + assert_eq!((cache.max_entries, cache.max_total_bytes, cache.max_entry_bytes), (3, 1024, 512)); + } +} diff --git a/crates/pi-ast/src/summary.rs b/crates/pi-ast/src/summary.rs index 14b63a1b2..d93a701b8 100644 --- a/crates/pi-ast/src/summary.rs +++ b/crates/pi-ast/src/summary.rs @@ -2,12 +2,11 @@ use std::{collections::BTreeSet, path::Path}; -use anyhow::{Result, anyhow}; -use ast_grep_core::tree_sitter::LanguageExt; +use anyhow::Result; use serde::{Deserialize, Serialize}; -use tree_sitter::{Node, Parser}; +use tree_sitter::Node; -use crate::language::SupportLang; +use crate::{language::SupportLang, parse_cache::parse_cached}; const DEFAULT_MIN_BODY_LINES: u32 = 4; const DEFAULT_MIN_COMMENT_LINES: u32 = 6; @@ -172,11 +171,7 @@ pub fn summarize_code(options: SummaryOptions) -> Result { return Ok(unparsed_result(source, total_lines)); }; - let mut parser = Parser::new(); - parser - .set_language(&language.get_ts_language()) - .map_err(|err| anyhow!("Failed to load tree-sitter language: {err}"))?; - let Some(tree) = parser.parse(&source, None) else { + let Some(tree) = parse_cached(&source, language)? else { return Ok(unparsed_result(source, total_lines)); }; let root = tree.root_node(); diff --git a/packages/natives/CHANGELOG.md b/packages/natives/CHANGELOG.md index 8dd3f04aa..d8844c63e 100644 --- a/packages/natives/CHANGELOG.md +++ b/packages/natives/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- `enclosingBlockBoundaries` and `blockRangeAt` now reuse a parsed tree-sitter tree when the same source and language were parsed before, and skip subtrees whose line span holds no visible line. Together these cut the block-context work the `read` tool performs on every non-raw read: for an 81KB TypeScript source with a mid-file window, 13.4ms to 4.45ms on a first parse and to 0.149ms once the tree is cached; for a 1.06MB source, 188.1ms to 55.8ms and to 0.440ms. The tree cache is bounded (12 entries, 4MiB of retained source) and verifies content byte-for-byte on every hit, so a hash collision can only cost a re-parse. The subtree skip is proven equivalent by differential comparison against the exhaustive walk across 4827 repository files and 38,616 window comparisons. + ## [17.3.5] - 2026-08-16 ### Fixed