Skip to content

perf(compiler2): tracked file_ast query — lower CST→AST once per file - #4043

Closed
hellovai wants to merge 1 commit into
canaryfrom
perf/reland-file-ast
Closed

perf(compiler2): tracked file_ast query — lower CST→AST once per file#4043
hellovai wants to merge 1 commit into
canaryfrom
perf/reland-file-ast

Conversation

@hellovai

@hellovai hellovai commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

What

On canary the CST→AST lowering step (baml_compiler2_ast::lower_file) is a plain, untracked function that six different consumers each re-run from scratch for the same file:

  • baml_compiler2_hir::file_semantic_index
  • baml_compiler2_ppir::file_semantic_index
  • baml_compiler2_ppir::ppir_expansion_items
  • the two project-wide expansion-map collectors (collect_block_attrs / collect_alias_bodies)
  • the baml_lsp2_actions check pass

Repeated CST traversal was ~31% of cold-compile CPU on the test corpus in the original audit. This PR memoizes lowering and removes two other redundant rebuilds.

Changes

  1. baml_compiler2_hir::file_ast(db, file) — a new #[salsa::tracked] query that lowers the CST once per file and shares items + lowering diagnostics + env_var_refs. All six consumers now read it instead of re-lowering. (EnvVarRef gains PartialEq, Eq so FileAst can use PartialEq for Salsa early-cutoff.)
  2. PPIR file_semantic_index delegates to HIR when a file has no expansion items. The post-expansion index is byte-for-byte identical to the pre-expansion one for such files (same AST items, same builder, same file range), so rebuilding it was pure wasted work. The merged path is preserved as file_semantic_index_expanded for files that actually have *$stream companions. The delegation also unifies scope identity for expansion-free files, which cuts infer_scope_types executions (15,590 → 13,331).
  3. PPIR function_body is now #[salsa::tracked] returning Arc<FunctionBody> (mirroring HIR's function_body). MIR lowering fetches the callee body at every direct-call site; the untracked version cloned the entire ExprBody arena each time.

Each changed site carries a short comment explaining what was being recomputed.

Before / after (cold compile)

Measured with the tools_compile_profile harness from #4038 (not committed here), disk cache disabled via BAML_NO_BYTECODE_CACHE=1, on crates/baml_tests/baml_src (77 files, 25 212 lines). Numbers are medians of 5 fresh-database cold runs, before/after binaries run interleaved on a quiet machine (two full rounds, both consistent; quietest round shown):

phase before (canary 2660b8b) after delta
check 1.118 s 0.752 s −33 %
emit 1.361 s 1.181 s −13 %
total 2.480 s 1.934 s −22 %

Load-independent evidence, same direction:

  • instructions retired (3 cold runs, /usr/bin/time -l): 93.7 G → 76.4 G (−18.5 %), stable across 3 interleaved rounds
  • file_ast executes 131× (exactly once per file) instead of lowering inline ~6× per file
  • file_semantic_index_expanded runs only 78× (files that really have *$stream expansions); the other 53 files reuse HIR's index
  • infer_scope_types: 15,590 → 13,331 executions

Tests

  • cargo test --workspace: all suites pass (~190 binaries). Initial failures were environmental only and reproduced on clean canary too: disk-full during the run, the Python SDK venv resolving to Python 3.10 (fixed with UV_PYTHON=3.12), and the Node fixtures needing sdk_tests/crates/typescript_node/setup.sh run once; after that sdk_test_python_pydantic2 (14/14) and sdk_test_typescript_node (15/15) pass.
  • Zero snapshot diffs (no .snap.new files) — outputs are byte-identical, as required for this track.
  • cargo clippy --workspace --all-targets --all-features -- -D warnings: clean.
  • cargo fmt check: clean.

Provenance

Re-derived against current canary; reference implementation: #4016 (perf/compiler2-cold-compile), audit items 3, 9, 13.

On canary the CST→AST lowering (baml_compiler2_ast::lower_file) is a plain,
untracked function that six consumers each re-ran from scratch: HIR's
file_semantic_index, PPIR's file_semantic_index, ppir_expansion_items, the
two project-wide expansion-map collectors (collect_block_attrs /
collect_alias_bodies), and the lsp2 check pass. Repeated CST traversal was
~31% of cold-compile CPU on the test corpus.

Changes:
- Add a salsa-tracked baml_compiler2_hir::file_ast(db, file) query that
  performs lowering once per file and shares items + lowering diagnostics +
  env var refs. Point all six consumers at it.
- PPIR file_semantic_index now delegates to HIR's file_semantic_index when a
  file has no expansion items: the post-expansion index is byte-for-byte the
  pre-expansion one for such files, so rebuilding it was wasted work (this
  also dedups scope identity, cutting infer_scope_types executions).
- Make PPIR function_body a tracked query returning Arc<FunctionBody> so MIR's
  per-call-site body fetches don't re-clone the ExprBody arena each time.

Re-derived against current canary; reference: PR #4016 (perf/compiler2-cold-compile).

Co-authored-by: Cursor <cursoragent@cursor.com>
@cursor

cursor Bot commented Jul 15, 2026

Copy link
Copy Markdown

Bugbot is not enabled for your account, so this pull request was not reviewed.

Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs.

@vercel

vercel Bot commented Jul 15, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
beps Ready Ready Preview, Comment Jul 15, 2026 8:56pm
promptfiddle Ready Ready Preview, Comment Jul 15, 2026 8:56pm
promptfiddle2 Ready Ready Preview, Comment Jul 15, 2026 8:56pm

Request Review

@coderabbitai

coderabbitai Bot commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

The change centralizes per-file CST-to-AST lowering in a Salsa-tracked file_ast query. HIR, PPIR, and LSP code reuse its items, diagnostics, and environment-variable references; PPIR semantic indexing also handles synthetic items through a tracked expanded path.

Changes

Memoized AST pipeline

Layer / File(s) Summary
Shared file AST lowering
baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs, baml_language/crates/baml_compiler2_hir/src/lib.rs
Adds equality support for EnvVarRef, introduces the cached FileAst result and file_ast query, and updates HIR semantic indexing to consume it.
PPIR expansion and indexed semantics
baml_language/crates/baml_compiler2_ppir/src/lib.rs
Reuses cached AST items for PPIR collection and expansion, adds synthetic-item semantic-index handling, and tracks function_body.
LSP validation integration
baml_language/crates/baml_lsp2_actions/src/check.rs
Uses cached AST items for associated-type binding validation.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SourceFile
  participant file_ast
  participant HIR
  participant PPIR
  participant LSP
  SourceFile->>file_ast: Lower syntax tree once
  file_ast-->>HIR: Return items, diagnostics, env_var_refs
  file_ast-->>PPIR: Return AST items
  file_ast-->>LSP: Return AST items
  PPIR->>HIR: Build expanded semantic index when synthetic items exist
Loading

Possibly related PRs

  • BoundaryML/baml#3635: Modifies the PPIR expansion pipeline around block attributes, alias bodies, and expansion items.
  • BoundaryML/baml#4016: Overlaps with the file_ast and PPIR semantic-index changes.

Poem

A rabbit found one AST,
Cached it neatly, built it fast.
HIR and PPIR shared the view,
LSP hopped along there too.
“No duplicate lowering today!”
And bounced through the code away.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: adding a tracked file_ast query to lower CST→AST once per file.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch perf/reland-file-ast

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@vercel
vercel Bot temporarily deployed to Preview – beps July 15, 2026 20:34 Inactive
@github-actions

Copy link
Copy Markdown

⏭️ Performance benchmarks were skipped

Perf benchmarks (CodSpeed) are opt-in on pull requests — they no longer run on every push. They always run automatically after merge to canary/main.

To run them on this PR, do any of the following, then push a commit (or re-run CI):

  • Add RUN_CODSPEED=1 to the PR description, or
  • Include run-perf or /perf in the PR title or any commit message.

@vercel
vercel Bot temporarily deployed to Preview – promptfiddle2 July 15, 2026 20:42 Inactive
@github-actions

Copy link
Copy Markdown

Binary size checks passed

7 passed

Artifact Platform File Gzip Gated on Baseline Delta Status
baml-cli Linux 🔒 22.8 MB 9.8 MB file 22.8 MB +72.9 KB (+0.3%) OK
packed-program Linux 🔒 16.3 MB 6.8 MB file 16.3 MB +36.8 KB (+0.2%) OK
baml-cli macOS 🔒 17.6 MB 8.5 MB file 17.5 MB +49.9 KB (+0.3%) OK
packed-program macOS 🔒 12.6 MB 6.0 MB file 12.6 MB +192 B (+0.0%) OK
baml-cli Windows 🔒 19.2 MB 8.7 MB file 18.7 MB +428.0 KB (+2.3%) OK
packed-program Windows 🔒 13.6 MB 6.1 MB file 13.5 MB +17.9 KB (+0.1%) OK
bridge_wasm WASM 15.2 MB 🔒 4.3 MB gzip 4.3 MB +5.1 KB (+0.1%) OK

🔒 = the size this artifact is GATED on (ceiling + delta). Binaries gate on file size (installed binary); WASM gates on gzip (download size). The other size is shown for information only.


Generated by cargo size-gate · workflow run

@hellovai

Copy link
Copy Markdown
Contributor Author

Folded into the combined re-landing PR #4054 (per maintainer preference for a single PR post-tool-merge). Branch kept for provenance; individual before/after measurements remain in this PR's description.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant