perf(compiler2): tracked file_ast query — lower CST→AST once per file - #4043
perf(compiler2): tracked file_ast query — lower CST→AST once per file#4043hellovai wants to merge 1 commit into
Conversation
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>
|
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. |
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📝 WalkthroughWalkthroughThe change centralizes per-file CST-to-AST lowering in a Salsa-tracked ChangesMemoized AST pipeline
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
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
⏭️ Performance benchmarks were skippedPerf benchmarks (CodSpeed) are opt-in on pull requests — they no longer run on every push. They always run automatically after merge to To run them on this PR, do any of the following, then push a commit (or re-run CI):
|
Binary size checks passed✅ 7 passed
Generated by |
|
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. |
What
On
canarythe 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_indexbaml_compiler2_ppir::file_semantic_indexbaml_compiler2_ppir::ppir_expansion_itemscollect_block_attrs/collect_alias_bodies)baml_lsp2_actionscheck passRepeated 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
baml_compiler2_hir::file_ast(db, file)— a new#[salsa::tracked]query that lowers the CST once per file and sharesitems+ loweringdiagnostics+env_var_refs. All six consumers now read it instead of re-lowering. (EnvVarRefgainsPartialEq, EqsoFileAstcan usePartialEqfor Salsa early-cutoff.)file_semantic_indexdelegates 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 asfile_semantic_index_expandedfor files that actually have*$streamcompanions. The delegation also unifies scope identity for expansion-free files, which cutsinfer_scope_typesexecutions (15,590 → 13,331).function_bodyis now#[salsa::tracked]returningArc<FunctionBody>(mirroring HIR'sfunction_body). MIR lowering fetches the callee body at every direct-call site; the untracked version cloned the entireExprBodyarena each time.Each changed site carries a short comment explaining what was being recomputed.
Before / after (cold compile)
Measured with the
tools_compile_profileharness from #4038 (not committed here), disk cache disabled viaBAML_NO_BYTECODE_CACHE=1, oncrates/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):Load-independent evidence, same direction:
/usr/bin/time -l): 93.7 G → 76.4 G (−18.5 %), stable across 3 interleaved roundsfile_astexecutes 131× (exactly once per file) instead of lowering inline ~6× per filefile_semantic_index_expandedruns only 78× (files that really have*$streamexpansions); the other 53 files reuse HIR's indexinfer_scope_types: 15,590 → 13,331 executionsTests
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 withUV_PYTHON=3.12), and the Node fixtures needingsdk_tests/crates/typescript_node/setup.shrun once; after thatsdk_test_python_pydantic2(14/14) andsdk_test_typescript_node(15/15) pass..snap.newfiles) — outputs are byte-identical, as required for this track.cargo clippy --workspace --all-targets --all-features -- -D warnings: clean.cargo fmtcheck: clean.Provenance
Re-derived against current
canary; reference implementation: #4016 (perf/compiler2-cold-compile), audit items 3, 9, 13.