From 010eda78348a5b0d9e31fc662fd67c1932d3f0e3 Mon Sep 17 00:00:00 2001 From: Sunil Srivatsa Date: Mon, 17 Aug 2026 10:09:06 -0700 Subject: [PATCH 1/3] perf(pi-ast): cache parsed trees by content and language MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `enclosing_block_boundaries`, `block_range_at` and `summarize_code` each re-parsed the whole file on every call. The results are not cacheable — boundaries depend on the caller's visible ranges, which differ per call — but the `tree_sitter::Tree` is, so cache that instead and hand out `ts_tree_copy` clones. Keyed on (xxh64 of the source, source length, language). The hash is a bucket selector only: a hit re-verifies the stored source against the request byte-for-byte before returning the tree, so a collision costs a re-parse and can never yield a tree built from other content. Language is in the key because the same bytes parsed as TypeScript and as Python are different trees. Bounded at 12 slots and 4 MiB of retained source with LRU eviction; sources above 4 MiB are parsed but never retained. `Tree` is `Send` but not `Sync`, so entries sit behind a `Mutex` that is held only for a map probe, a byte compare and a refcount bump, never across a parse or walk. Error trees are cached like any other: `has_error()` is a property of the tree, so the callers' own checks reach an identical verdict from a cached tree, and repeated "does this parse" probes get the speedup too. Measured (M4 Max, bazel-built .node, median of 20, 1-40 visible): read.ts 81 KB 13.34 ms -> 8.86 ms on repeat; 1 MB synthetic 225.7 ms -> 138.3 ms. Parser::new + set_language measured at 0.30 us against a 3.91 ms parse, so no parser pooling. --- Cargo.lock | 1 + crates/pi-ast/Cargo.toml | 1 + crates/pi-ast/src/block.rs | 157 ++++++++++- crates/pi-ast/src/lib.rs | 1 + crates/pi-ast/src/parse_cache.rs | 433 +++++++++++++++++++++++++++++++ crates/pi-ast/src/summary.rs | 13 +- 6 files changed, 583 insertions(+), 23 deletions(-) create mode 100644 crates/pi-ast/src/parse_cache.rs diff --git a/Cargo.lock b/Cargo.lock index c6e330cac..89307d45f 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..036f91c57 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(); @@ -240,11 +238,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"); + } } 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(); From 9169ac5c521f766d5df42e96928b40ba70e5de22 Mon Sep 17 00:00:00 2001 From: Sunil Srivatsa Date: Mon, 17 Aug 2026 12:57:59 -0700 Subject: [PATCH 2/3] perf(pi-ast): prune subtrees that cannot hold a block boundary collect_boundaries walked every node in the file even though the answer is bounded by the visible window, which cost roughly twice the parse: on an 81KB source the walk was 8.95ms against a 4.32ms parse, and on 1MB it was 138.8ms against 87.6ms. A node contributes a boundary 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 its ancestor's, so a span holding no visible line rules out that node and everything beneath it. Skip such subtrees with a binary search over the merged visible ranges. The test is the raw span, not endpoint visibility: a node whose span merely straddles the window has both endpoints outside it yet can contain a child that opens exactly on a visible line. The raw span is also conservative relative to node_content_end_line, so the prune needs no reasoning about that newline adjustment. Equivalence is proven differentially rather than argued: the pre-prune walk is retained under cfg(test) and compared for exact Option> equality across 4827 .ts/.py/.rs files and 38,616 comparisons over eight window shapes, including whole-file-visible, past-EOF, disjoint ranges, empty range lists and files that fail to parse. Zero mismatches. root.has_error() is still evaluated on the whole tree before the walk, so pruning cannot change a None verdict. Measured on the built addon with a mid-file 40-line window, medians of 20, against the parse cache alone: 81KB 13.4ms -> 4.45ms cold and 9.04ms -> 0.149ms warm; 1.06MB 188.1ms -> 55.8ms cold and 131.7ms -> 0.440ms warm. --- crates/pi-ast/src/block.rs | 327 +++++++++++++++++++++++++++++++++++++ 1 file changed, 327 insertions(+) diff --git a/crates/pi-ast/src/block.rs b/crates/pi-ast/src/block.rs index 036f91c57..7acb075cf 100644 --- a/crates/pi-ast/src/block.rs +++ b/crates/pi-ast/src/block.rs @@ -179,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() { @@ -708,4 +728,311 @@ mod tests { 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"); + } } From 0e9f4dab52d2d07bf99e108149a22609351e376f Mon Sep 17 00:00:00 2001 From: Sunil Srivatsa Date: Mon, 17 Aug 2026 13:03:16 -0700 Subject: [PATCH 3/3] docs(changelog): record the tree cache and boundary prune --- packages/natives/CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) 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