smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(pi-ast): prune subtrees that cannot hold a block boundary

Browse Source 
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<Vec<u32>> 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.
This commit is contained in:
Sunil Srivatsa
2026-08-17 12:57:59 -07:00
parent 010eda7834
commit 9169ac5c52
1 changed files with 327 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pi-ast/src/block.rs
+327
View File
@@ -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<u32>) {
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<u32>,
) {
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<Vec<u32>> {
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<Vec<(u32, u32)>> {
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<Vec<u32>>`
/// 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<usize>) -> Vec<PathBuf> {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../..");
let mut files: Vec<PathBuf> = 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");
}
}