Repository navigation
perf(emit): skip redundant Pass-1 MIR lowering [run-perf] - #3846
Closed
codeshaunted wants to merge 7 commits into
Closed
codeshaunted wants to merge 7 commits into
codeshaunted wants to merge 7 commits into
Conversation
Bytecode emit lowered every function twice: once in Pass 1 only to read its global-slot name and intrinsic/await-any skip flag, then again in Pass 4 to actually emit code. lower_function is not memoized, so this doubled all MIR lowering work -- including the ~676-function stdlib that is recompiled from embedded source on every invocation. Derive both Pass-1 needs without lowering: the slot name from def_to_item_ref (identical to mir.item_ref in every lowering branch), and the skip set by matching the function body directly (only FunctionBody::Builtin lowers to MirFunctionKind::Builtin; Expr/Missing bodies always lower to Bytecode). Pass 4 still lowers to emit code, so the emitted Program is byte-identical. Clean-compile wall time (baml check, median of 9 runs): 1 fn: 0.35s -> 0.28s (-20%) 200 fns: 0.37s -> 0.30s (-19%) 5000 fns: 1.22s -> 0.93s (-24%)
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds a precompiled stdlib emission prefix, a resume-capable bytecode emission path, builtins-first class type-tag ordering, AST and HIR Borsh support, and project/CLI entry points that use the prefix-based compilation flow. ChangesPrecompiled emission prefix and serialization
AST and HIR Borsh serialization
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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 are running for this PRCodSpeed perf benchmarks were triggered because this PR opted in. Results will appear in the CodSpeed check / dashboard once they finish. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
No description provided. |
Merging this PR will improve performance by 15%
Performance Changes
Tip Curious why this is faster? Comment Comparing 
|
…[run-perf] The stdlib (~676 functions / 8287 lines) was recompiled from embedded source on every invocation. Its bytecode is deterministic, so compile it once at build time and embed the result. - Restructure emit so stdlib (`<builtin>/`) files form a stable prefix of the global-slot space, object pool, and class type tags (Passes 1-4 run per-block: stdlib, then user). Align build_class_type_tags in MIR lowering to the same builtins-first order -- type-switch arms must match runtime class.type_tag or type switches miss and throw Unreachable. - Add EmitState and split emit into emit_core (Passes 1-4, resumable from a prefix) and emit_trailing (Passes 4.5-8, over all files). generate_project_bytecode_with_prefix resumes from a precompiled stdlib EmitState, emitting only user files on top. - New crate baml_builtins2_prebuilt: build.rs runs precompile_stdlib once and borsh-serializes the EmitState (861KB) into the binary; CLI check/test/run inject it via ProjectDatabase::get_bytecode_with_prefix. Compiler access routes through baml_db/baml_project per the crate-dependency policy. baml check floor ~289ms -> ~143ms (~2x). Full-compile path unchanged. Verified: baml_tests 2619/2619, corpus via prefix 1948/1948.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
baml_language/crates/baml_project/src/db.rs (1)
445-462: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMinor duplication with
get_bytecode.
get_bytecodeandget_bytecode_with_prefixrepeat the same error-gate andCompileOptionsconstruction. Optional: extract the shared gate (e.g. a small helper returningErrearly or the preparedopts) to keep the two paths in sync. TheOptLevel::Two/prefix wiring itself is consistent with the build-time prefix and CLI callers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/baml_project/src/db.rs` around lines 445 - 462, The error gate and CompileOptions setup in get_bytecode_with_prefix duplicate the same logic from get_bytecode, so extract the shared “project has errors” check and/or options construction into a small helper and reuse it from both methods. Keep the existing GenerateProjectBytecodeWithPrefix call, OptLevel::Two, and prefix handling unchanged, and make sure both paths continue to return the same LoweringError::ProjectHasErrors behavior.baml_language/crates/baml_builtins2_prebuilt/build.rs (1)
33-36: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
cargo:warningemits a build warning on every compile.Using
cargo:warning=for an informational byte-count prints a yellow warning on each build of this crate, adding noise. Consider gating it or dropping it once the artifact size is verified.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/baml_builtins2_prebuilt/build.rs` around lines 33 - 36, The `build.rs` logging in `baml_builtins2_prebuilt` is using `cargo:warning` for an informational byte-count, which makes every compile emit a warning. Update the `println!` in the build script to avoid warning-level output from the `stdlib EmitState prefix` message, either by removing it after verification or gating it behind a debug/opt-in condition so `build.rs` stays quiet during normal builds.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@baml_language/crates/baml_builtins2_prebuilt/build.rs`:
- Around line 1-6: The module doc comment is stale and describes the wrong build
artifact; it currently says the build script compiles a Program into
stdlib_program.bin, but the build flow in build.rs and the matching
include_bytes! in lib.rs actually produce an EmitState prefix in
stdlib_prefix.bin. Update the top-level doc comment in build.rs to reflect the
real artifact type and filename so it matches the behavior of the build script
and its consumer.
In `@baml_language/crates/baml_compiler2_emit/src/lib.rs`:
- Around line 1258-1263: Write the function object into its reserved global slot
instead of appending it, because `program.globals.len()` may not match the slot
already reserved in `globals[&fq_name]` when resuming from a prefix. Update the
function emission logic in `baml_language/crates/baml_compiler2_emit/src/lib.rs`
around `program.add_object`, `function_indices`, and `function_global_indices`
so it uses the Pass 1 global index for `fq_name`, pads any missing entries with
`ConstValue::Null`, and stores the function at that exact slot to keep bytecode
references and later let initialization aligned.
- Around line 739-748: precompile_stdlib is no longer stdlib-only because it
forwards user_files into emit_core, which can cause user objects/globals to be
baked into the stdlib prefix and duplicated on resume. Update precompile_stdlib
in lib.rs to keep the prefix build limited to stdlib_files by passing an empty
user slice to emit_core, or add an assertion that user_files.is_empty() before
the call. Use the existing precompile_stdlib, compiler2_all_files, and emit_core
symbols to locate the change.
---
Nitpick comments:
In `@baml_language/crates/baml_builtins2_prebuilt/build.rs`:
- Around line 33-36: The `build.rs` logging in `baml_builtins2_prebuilt` is
using `cargo:warning` for an informational byte-count, which makes every compile
emit a warning. Update the `println!` in the build script to avoid warning-level
output from the `stdlib EmitState prefix` message, either by removing it after
verification or gating it behind a debug/opt-in condition so `build.rs` stays
quiet during normal builds.
In `@baml_language/crates/baml_project/src/db.rs`:
- Around line 445-462: The error gate and CompileOptions setup in
get_bytecode_with_prefix duplicate the same logic from get_bytecode, so extract
the shared “project has errors” check and/or options construction into a small
helper and reuse it from both methods. Keep the existing
GenerateProjectBytecodeWithPrefix call, OptLevel::Two, and prefix handling
unchanged, and make sure both paths continue to return the same
LoweringError::ProjectHasErrors behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: a122c67e-ccfd-49c6-9136-3b26dfc04d82
⛔ Files ignored due to path filters (3)
baml_language/Cargo.lockis excluded by!**/*.lockbaml_language/crates/baml_tests/tests/bytecode_format/snapshots/bytecode_format__bytecode_display_expanded.snapis excluded by!**/*.snapbaml_language/crates/baml_tests/tests/bytecode_format/snapshots/bytecode_format__bytecode_display_expanded_unoptimized.snapis excluded by!**/*.snap
📒 Files selected for processing (12)
baml_language/Cargo.tomlbaml_language/crates/baml_builtins2_prebuilt/Cargo.tomlbaml_language/crates/baml_builtins2_prebuilt/build.rsbaml_language/crates/baml_builtins2_prebuilt/src/lib.rsbaml_language/crates/baml_cli/Cargo.tomlbaml_language/crates/baml_cli/src/check_command.rsbaml_language/crates/baml_cli/src/run_command.rsbaml_language/crates/baml_cli/src/test_command.rsbaml_language/crates/baml_compiler2_emit/Cargo.tomlbaml_language/crates/baml_compiler2_emit/src/lib.rsbaml_language/crates/baml_compiler2_mir/src/lower.rsbaml_language/crates/baml_project/src/db.rs
✅ Files skipped from review due to trivial changes (1)
- baml_language/crates/baml_compiler2_emit/Cargo.toml
Comment on lines
+739
to
+748
| pub fn precompile_stdlib(db: &dyn baml_compiler2_mir::Db, opt: OptLevel) -> EmitState { | ||
| let all_files = compiler2_all_files(db); | ||
| let (stdlib_files, user_files): (Vec<baml_base::SourceFile>, Vec<baml_base::SourceFile>) = | ||
| all_files | ||
| .iter() | ||
| .copied() | ||
| .partition(|f| f.path(db).to_string_lossy().starts_with("<builtin>/")); | ||
| let alias_caches = build_alias_caches(db, &all_files); | ||
| emit_core(db, opt, &stdlib_files, &user_files, &alias_caches, None) | ||
| } |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Keep precompile_stdlib stdlib-only.
Line 747 passes user_files into the prefix build, so a caller using a normal project DB would bake user objects/globals into the “stdlib” prefix and then emit those user files again on resume. Pass an empty user slice, or assert user_files.is_empty().
Proposed fix
- emit_core(db, opt, &stdlib_files, &user_files, &alias_caches, None)
+ debug_assert!(
+ user_files.is_empty(),
+ "precompile_stdlib must be called with a stdlib-only database"
+ );
+ emit_core(db, opt, &stdlib_files, &[], &alias_caches, None)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| pub fn precompile_stdlib(db: &dyn baml_compiler2_mir::Db, opt: OptLevel) -> EmitState { | |
| let all_files = compiler2_all_files(db); | |
| let (stdlib_files, user_files): (Vec<baml_base::SourceFile>, Vec<baml_base::SourceFile>) = | |
| all_files | |
| .iter() | |
| .copied() | |
| .partition(|f| f.path(db).to_string_lossy().starts_with("<builtin>/")); | |
| let alias_caches = build_alias_caches(db, &all_files); | |
| emit_core(db, opt, &stdlib_files, &user_files, &alias_caches, None) | |
| } | |
| pub fn precompile_stdlib(db: &dyn baml_compiler2_mir::Db, opt: OptLevel) -> EmitState { | |
| let all_files = compiler2_all_files(db); | |
| let (stdlib_files, user_files): (Vec<baml_base::SourceFile>, Vec<baml_base::SourceFile>) = | |
| all_files | |
| .iter() | |
| .copied() | |
| .partition(|f| f.path(db).to_string_lossy().starts_with("<builtin>/")); | |
| let alias_caches = build_alias_caches(db, &all_files); | |
| debug_assert!( | |
| user_files.is_empty(), | |
| "precompile_stdlib must be called with a stdlib-only database" | |
| ); | |
| emit_core(db, opt, &stdlib_files, &[], &alias_caches, None) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_compiler2_emit/src/lib.rs` around lines 739 - 748,
precompile_stdlib is no longer stdlib-only because it forwards user_files into
emit_core, which can cause user objects/globals to be baked into the stdlib
prefix and duplicated on resume. Update precompile_stdlib in lib.rs to keep the
prefix build limited to stdlib_files by passing an empty user slice to
emit_core, or add an assertion that user_files.is_empty() before the call. Use
the existing precompile_stdlib, compiler2_all_files, and emit_core symbols to
locate the change.
Comment on lines
+1258
to
+1263
| let fn_obj_idx = program.add_object(Object::Function(Box::new(compiled_fn))); | ||
| program.function_indices.insert(fq_name.clone(), fn_obj_idx); | ||
| program | ||
| .function_global_indices | ||
| .insert(fq_name, program.globals.len()); | ||
| program.add_global(ConstValue::Object(ObjectIndex::from_raw(fn_obj_idx))); |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Write functions into their reserved global slot.
When resuming from a prefix that already reserved let slots, program.globals.len() can be before globals[&fq_name]. Appending here makes bytecode call one slot while function_global_indices points at another, and later let initialization can overwrite function globals. Use the Pass 1 slot and pad holes with Null.
Proposed fix
let fn_obj_idx = program.add_object(Object::Function(Box::new(compiled_fn)));
program.function_indices.insert(fq_name.clone(), fn_obj_idx);
+ let global_slot = *globals
+ .get(&fq_name)
+ .expect("compiled function must have a reserved global slot");
+ while program.globals.len() < global_slot {
+ program.add_global(ConstValue::Null);
+ }
+ debug_assert_eq!(program.globals.len(), global_slot);
program
.function_global_indices
- .insert(fq_name, program.globals.len());
+ .insert(fq_name, global_slot);
program.add_global(ConstValue::Object(ObjectIndex::from_raw(fn_obj_idx)));📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| let fn_obj_idx = program.add_object(Object::Function(Box::new(compiled_fn))); | |
| program.function_indices.insert(fq_name.clone(), fn_obj_idx); | |
| program | |
| .function_global_indices | |
| .insert(fq_name, program.globals.len()); | |
| program.add_global(ConstValue::Object(ObjectIndex::from_raw(fn_obj_idx))); | |
| let fn_obj_idx = program.add_object(Object::Function(Box::new(compiled_fn))); | |
| program.function_indices.insert(fq_name.clone(), fn_obj_idx); | |
| let global_slot = *globals | |
| .get(&fq_name) | |
| .expect("compiled function must have a reserved global slot"); | |
| while program.globals.len() < global_slot { | |
| program.add_global(ConstValue::Null); | |
| } | |
| debug_assert_eq!(program.globals.len(), global_slot); | |
| program | |
| .function_global_indices | |
| .insert(fq_name, global_slot); | |
| program.add_global(ConstValue::Object(ObjectIndex::from_raw(fn_obj_idx))); |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_compiler2_emit/src/lib.rs` around lines 1258 -
1263, Write the function object into its reserved global slot instead of
appending it, because `program.globals.len()` may not match the slot already
reserved in `globals[&fq_name]` when resuming from a prefix. Update the function
emission logic in `baml_language/crates/baml_compiler2_emit/src/lib.rs` around
`program.add_object`, `function_indices`, and `function_global_indices` so it
uses the Pass 1 global index for `fq_name`, pads any missing entries with
`ConstValue::Null`, and stores the function at that exact slot to keep bytecode
references and later let initialization aligned.
…-perf] Foundation for serializing the HIR/AST so the stdlib's typecheck artifacts can be precompiled and loaded (skipping stdlib parse/resolve at runtime). The AST embeds three foreign types that lack borsh impls -- text_size::TextRange (spans), la_arena::Idx (arena indices), la_arena::Arena (expression arenas) -- and the orphan rule forbids implementing Borsh on them directly. borsh 1.6's `#[borsh(serialize_with/deserialize_with)]` field attributes are the escape hatch: this module provides the helper fns (TextRange, Idx, Option<Idx>, Arena<T>, Arena<TextRange>). Arena is dense/alloc-only, so re-allocing elements in iteration order on load reproduces the original Idx values. borsh keeps one serialization format across the precompiled-stdlib artifacts. Validated by a round-trip unit test. Next: derive Borsh across the AST + HIR ItemTree using these helpers, then a build-time stdlib HIR artifact + query interception for `<builtin>/` files.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs (1)
148-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the
Nonebranch in the round-trip test.The new test only exercises
opt_idx: Some(i0), so the zero-tag path inserialize_opt_idx/deserialize_opt_idxis still untested.As per coding guidelines, "Prefer writing Rust unit tests over integration tests where possible."Suggested tweak
- let sample = Sample { - span: TextRange::new(5.into(), 11.into()), - idx: i0, - opt_idx: Some(i0), - arena, - spans, - }; - - let bytes = borsh::to_vec(&sample).unwrap(); - let decoded: Sample = borsh::from_slice(&bytes).unwrap(); - assert_eq!(sample, decoded); + for opt_idx in [Some(i0), None] { + let sample = Sample { + span: TextRange::new(5.into(), 11.into()), + idx: i0, + opt_idx, + arena: arena.clone(), + spans: spans.clone(), + }; + + let bytes = borsh::to_vec(&sample).unwrap(); + let decoded: Sample = borsh::from_slice(&bytes).unwrap(); + assert_eq!(sample, decoded); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs` around lines 148 - 167, The round-trip test in foreign_types_round_trip_through_borsh only covers opt_idx = Some(i0), leaving the None serialization/deserialization path in serialize_opt_idx and deserialize_opt_idx untested. Update foreign_types_round_trip_through_borsh or add a second unit test to also construct a Sample with opt_idx: None, then borsh::to_vec/borsh::from_slice it and assert equality so both branches are exercised.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs`:
- Around line 64-69: The deserialize_opt_idx helper currently accepts any
non-zero discriminant as Some, so invalid Option<Idx<T>> tags are silently
treated as valid. Update deserialize_opt_idx in borsh_helpers.rs to explicitly
match only the expected tags for None and Some, and return an Error for any
other byte value before calling deserialize_idx. Keep the change localized to
deserialize_opt_idx so malformed input is rejected consistently.
---
Nitpick comments:
In `@baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs`:
- Around line 148-167: The round-trip test in
foreign_types_round_trip_through_borsh only covers opt_idx = Some(i0), leaving
the None serialization/deserialization path in serialize_opt_idx and
deserialize_opt_idx untested. Update foreign_types_round_trip_through_borsh or
add a second unit test to also construct a Sample with opt_idx: None, then
borsh::to_vec/borsh::from_slice it and assert equality so both branches are
exercised.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d53cef6d-9315-4a2d-b006-772b69837939
⛔ Files ignored due to path filters (1)
baml_language/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (3)
baml_language/crates/baml_compiler2_ast/Cargo.tomlbaml_language/crates/baml_compiler2_ast/src/borsh_helpers.rsbaml_language/crates/baml_compiler2_ast/src/lib.rs
✅ Files skipped from review due to trivial changes (1)
- baml_language/crates/baml_compiler2_ast/src/lib.rs
Comment on lines
+64
to
+69
| pub fn deserialize_opt_idx<T, R: Read>(r: &mut R) -> Result<Option<Idx<T>>, Error> { | ||
| if u8::deserialize_reader(r)? == 0 { | ||
| Ok(None) | ||
| } else { | ||
| Ok(Some(deserialize_idx(r)?)) | ||
| } |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Reject invalid Option<Idx<T>> discriminants.
Line 65 currently treats any non-zero tag as Some, so malformed bytes like 2 deserialize successfully instead of failing. That widens the binary format and can hide corrupted precompiled artifacts.
Suggested fix
pub fn deserialize_opt_idx<T, R: Read>(r: &mut R) -> Result<Option<Idx<T>>, Error> {
- if u8::deserialize_reader(r)? == 0 {
- Ok(None)
- } else {
- Ok(Some(deserialize_idx(r)?))
- }
+ match u8::deserialize_reader(r)? {
+ 0 => Ok(None),
+ 1 => Ok(Some(deserialize_idx(r)?)),
+ tag => Err(Error::new(
+ std::io::ErrorKind::InvalidData,
+ format!("invalid Option<Idx<T>> discriminant: {tag}"),
+ )),
+ }
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| pub fn deserialize_opt_idx<T, R: Read>(r: &mut R) -> Result<Option<Idx<T>>, Error> { | |
| if u8::deserialize_reader(r)? == 0 { | |
| Ok(None) | |
| } else { | |
| Ok(Some(deserialize_idx(r)?)) | |
| } | |
| pub fn deserialize_opt_idx<T, R: Read>(r: &mut R) -> Result<Option<Idx<T>>, Error> { | |
| match u8::deserialize_reader(r)? { | |
| 0 => Ok(None), | |
| 1 => Ok(Some(deserialize_idx(r)?)), | |
| tag => Err(Error::new( | |
| std::io::ErrorKind::InvalidData, | |
| format!("invalid Option<Idx<T>> discriminant: {tag}"), | |
| )), | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs` around lines 64
- 69, The deserialize_opt_idx helper currently accepts any non-zero discriminant
as Some, so invalid Option<Idx<T>> tags are silently treated as valid. Update
deserialize_opt_idx in borsh_helpers.rs to explicitly match only the expected
tags for None and Some, and return an Error for any other byte value before
calling deserialize_idx. Keep the change localized to deserialize_opt_idx so
malformed input is rejected consistently.
Make the full compiler2 AST borsh-serializable (foundation for the precomputed stdlib HIR artifact). All 57 ast.rs types now derive Borsh; foreign-typed fields (Idx/Arena/TextRange, and tuple/Option/Vec/HashSet of them) use the borsh_helpers via `#[borsh(serialize_with/deserialize_with)]`. - baml_base: derive Borsh on TypePath (Vec<Name>, trivially serializable). - AstSourceMap secondary diagnostic sub-span maps (member/path-segment/label, HashMap<ExprId, ..>) are `#[borsh(skip)]`: they only refine diagnostic locations, never consulted for the frozen/error-free stdlib the artifact is built from. Primary spans (the per-kind Arena<TextRange>) are preserved. Additive (new trait impls only); existing AST tests pass (72), behavior unchanged. Next: derive Borsh across the HIR ItemTree.
Make the compiler2 HIR item tree borsh-serializable (next step toward the precomputed stdlib artifact; builds on the AST borsh sweep). - LocalItemId<T>: derive Borsh with `_phantom` skipped and an empty bound override on `packed` so no `T: Borsh` bound is added; add manual Ord/PartialOrd over `packed` (no `T: Ord` bound) -- borsh sorts HashMap keys, so LocalItemId map keys must be Ord once ItemTree derives Borsh. - ItemKind: derive Borsh + Ord (it keys the collision-tracking map). - All 23 item_tree.rs types derive Borsh; bare TextRange spans use the AST crate's borsh_helpers. ItemTreeSourceMap name-span side-tables are `#[borsh(skip)]` (diagnostic-only, never consulted for the frozen stdlib). Additive; hir compiles + clippy clean. Next: build-time stdlib HIR artifact + query interception for `<builtin>/` files.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs`:
- Around line 140-159: The HashSet serialization in serialize_idx_set is
non-deterministic because iteration order is arbitrary, so sort the serialized
Idx values by their raw integer representation before writing to guarantee
stable bytes. Update deserialize_idx_set to validate input by rejecting
duplicate indices instead of silently inserting them into the HashSet. Keep the
fix localized to serialize_idx_set and deserialize_idx_set in borsh_helpers so
prebuilt artifacts remain deterministic and malformed inputs fail fast.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 6e8e565a-20ef-48db-8d3c-718ce00784de
📒 Files selected for processing (3)
baml_language/crates/baml_base/src/core_types.rsbaml_language/crates/baml_compiler2_ast/src/ast.rsbaml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs
Comment on lines
+140
to
+159
| pub fn serialize_idx_set<T, W: Write>( | ||
| set: &std::collections::HashSet<Idx<T>>, | ||
| w: &mut W, | ||
| ) -> Result<(), Error> { | ||
| len_as_u32(set.len()).serialize(w)?; | ||
| for idx in set { | ||
| serialize_idx(idx, w)?; | ||
| } | ||
| Ok(()) | ||
| } | ||
|
|
||
| pub fn deserialize_idx_set<T, R: Read>( | ||
| r: &mut R, | ||
| ) -> Result<std::collections::HashSet<Idx<T>>, Error> { | ||
| let len = u32::deserialize_reader(r)?; | ||
| let mut set = std::collections::HashSet::with_capacity(len as usize); | ||
| for _ in 0..len { | ||
| set.insert(deserialize_idx(r)?); | ||
| } | ||
| Ok(set) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
ast-grep run --pattern 'struct Idx<$T> { $BODY }' --lang rust baml_language/crates/baml_compiler2_ast/src
ast-grep run --pattern 'pub fn into_raw(&self) -> $RET' --lang rust baml_language/crates/baml_compiler2_ast/src
ast-grep outline baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rsRepository: ast-grep
Length of output: 2008
🏁 Script executed:
#!/bin/bash
# Find the Idx struct definition in the local repo
# First, locate files containing "struct Idx"
rg "struct Idx<T>" --type rust -A 10
# Inspect the borsh_helpers.rs file to see imports and Idx usage context
cat -n baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs | head -50
# Check for Idx definition in the same file or nearby files
rg "impl.*Idx" --type rust -A 5
# Check Error definition
rg "pub.*Error" --type rust baml_language/crates/baml_compiler2_ast/src/ | head -20Repository: BoundaryML/baml
Length of output: 153
🏁 Script executed:
#!/bin/bash
# Read the full content of borsh_helpers.rs to see imports and Idx definition
cat -n baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rsRepository: BoundaryML/baml
Length of output: 9906
🏁 Script executed:
#!/bin/bash
# Check if Idx is Copy by searching for traits or usage
rg "impl.*Copy.*Idx" --type rust
# Check la_arena documentation or source if available locally, or assume Copy based on common pattern
# Check Error construction in borsh
rg "Error::new" --type rust baml_language/crates/baml_compiler2_ast/src/ | head -5
# Check if HashSet tests exist
rg "idx_set" --type rust baml_language/crates/baml_compiler2_ast/src/
# Check if Idx is Copy by looking at la_arena usage or docs
rg "use la_arena" --type rust baml_language/crates/baml_compiler2_ast/src/ -A 5Repository: BoundaryML/baml
Length of output: 153
Canonicalize HashSet serialization to ensure deterministic byte output.
HashSet iteration order in Rust is arbitrary, but prebuilt artifacts require deterministic serialization. Sorting indices by their raw integer value before writing ensures equivalent sets produce identical bytes. Additionally, rejecting duplicate indices during deserialization prevents silent data corruption from malformed inputs.
[large_fixes_required]
Proposed fix
pub fn serialize_idx_set<T, W: Write>(
set: &std::collections::HashSet<Idx<T>>,
w: &mut W,
) -> Result<(), Error> {
len_as_u32(set.len()).serialize(w)?;
- for idx in set {
- serialize_idx(idx, w)?;
+ let mut indices: Vec<_> = set.iter().copied().collect();
+ indices.sort_by_key(|idx| idx.into_raw().into_u32());
+ for idx in &indices {
+ serialize_idx(idx, w)?;
}
Ok(())
}
pub fn deserialize_idx_set<T, R: Read>(
r: &mut R,
) -> Result<std::collections::HashSet<Idx<T>>, Error> {
let len = u32::deserialize_reader(r)?;
let mut set = std::collections::HashSet::with_capacity(len as usize);
for _ in 0..len {
- set.insert(deserialize_idx(r)?);
+ let idx = deserialize_idx(r)?;
+ if !set.insert(idx) {
+ return Err(Error::new(
+ std::io::ErrorKind::InvalidData,
+ "duplicate Idx<T> in serialized HashSet",
+ ));
+ }
}
Ok(set)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| pub fn serialize_idx_set<T, W: Write>( | |
| set: &std::collections::HashSet<Idx<T>>, | |
| w: &mut W, | |
| ) -> Result<(), Error> { | |
| len_as_u32(set.len()).serialize(w)?; | |
| for idx in set { | |
| serialize_idx(idx, w)?; | |
| } | |
| Ok(()) | |
| } | |
| pub fn deserialize_idx_set<T, R: Read>( | |
| r: &mut R, | |
| ) -> Result<std::collections::HashSet<Idx<T>>, Error> { | |
| let len = u32::deserialize_reader(r)?; | |
| let mut set = std::collections::HashSet::with_capacity(len as usize); | |
| for _ in 0..len { | |
| set.insert(deserialize_idx(r)?); | |
| } | |
| Ok(set) | |
| pub fn serialize_idx_set<T, W: Write>( | |
| set: &std::collections::HashSet<Idx<T>>, | |
| w: &mut W, | |
| ) -> Result<(), Error> { | |
| len_as_u32(set.len()).serialize(w)?; | |
| let mut indices: Vec<_> = set.iter().copied().collect(); | |
| indices.sort_by_key(|idx| idx.into_raw().into_u32()); | |
| for idx in &indices { | |
| serialize_idx(idx, w)?; | |
| } | |
| Ok(()) | |
| } | |
| pub fn deserialize_idx_set<T, R: Read>( | |
| r: &mut R, | |
| ) -> Result<std::collections::HashSet<Idx<T>>, Error> { | |
| let len = u32::deserialize_reader(r)?; | |
| let mut set = std::collections::HashSet::with_capacity(len as usize); | |
| for _ in 0..len { | |
| let idx = deserialize_idx(r)?; | |
| if !set.insert(idx) { | |
| return Err(Error::new( | |
| std::io::ErrorKind::InvalidData, | |
| "duplicate Idx<T> in serialized HashSet", | |
| )); | |
| } | |
| } | |
| Ok(set) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@baml_language/crates/baml_compiler2_ast/src/borsh_helpers.rs` around lines
140 - 159, The HashSet serialization in serialize_idx_set is non-deterministic
because iteration order is arbitrary, so sort the serialized Idx values by their
raw integer representation before writing to guarantee stable bytes. Update
deserialize_idx_set to validate input by rejecting duplicate indices instead of
silently inserting them into the HashSet. Keep the fix localized to
serialize_idx_set and deserialize_idx_set in borsh_helpers so prebuilt artifacts
remain deterministic and malformed inputs fail fast.
…[run-perf]
Two changes that cut the stdlib's cost in `baml check`/`run`/`test`:
1. Precompiled stdlib HIR (load path). Build + embed the frozen stdlib's
pre-lowered AST so file_semantic_index skips lex/parse/lower for `<builtin>/`
files and re-runs only the semantic-index builder (which re-interns all `'db`
handles in the runtime DB -> identical result to from-source).
- hir::precompiled: PrecompiledFile (borsh: items + env_var_refs + range) +
an OnceLock cache; unset -> parse as before (tests/LSP unaffected).
- baml_compiler2_ast: derive Borsh on EnvVarRef.
- baml_db: re-export baml_compiler2_ast (prebuilt crate routes via baml_db).
- baml_builtins2_prebuilt build.rs emits stdlib_hir.bin (561KB / 53 files,
asserts zero lowering diagnostics); CLI installs it at startup.
2. Skip re-type-checking the frozen stdlib's bodies in collect_compiler2_-
diagnostics (`<builtin>/` files). Validated at artifact-build time; user
typechecking needs only their declared signatures, so re-running TIR body
inference + lints over them every compile is wasted -- the dominant stdlib
cost in `check` (parse/build is comparatively cheap).
baml check floor: ~143ms -> ~70ms (~2x here; ~5x vs the original ~360ms).
Verified: corpus 1948/1948 via the fast-path, full baml_tests 2619/2619, user
errors still report. Remaining toward ~15ms: cache FileSemanticIndex (skip the
builder) + TIR signatures -- the HIR serialization here is the foundation.
…un-perf] The first `baml check` lazily resolves the whole stdlib type graph through `package_interface(<stdlib pkg>)` (signatures) and `package_implements_registry` (interface-impl rules); a trivial user file pulls in ~62ms of this. Both return pure data (Ty + names; no `'db` lifetime, no salsa handles), so they serialize with borsh and the cached value is returned directly -- no reconstruction. - TIR cache (the win): build + embed each frozen stdlib package's resolved `PackageInterface` and `ImplementsRegistry`; `package_interface` / `package_implements_registry` short-circuit for stdlib packages. Borsh on PackageInterface/ExportedType/ExportedFunction/FunctionThrowSets + ImplementsRegistry (+5 helpers); `precompiled_tir` cache; installed at startup. - HIR cache (upgraded): `file_semantic_index` for builtin files now rehydrates a full precompiled `FileSemanticIndex` (skips parse AND the builder), re-interning scope_ids + symbol_contributions loc-handles in the runtime DB. Used by PPIR `package_items`. Borsh on the scope-tree (Scope/ScopeKind/ScopeBindings/ LocalBinding/BindingId/DefinitionSite/PathResolution/FileScopeId) + DefinitionKind. - baml_compiler2_tir gains a `borsh` dep; baml_builtins2_prebuilt build.rs emits stdlib_tir.bin + stdlib_implements.bin; install_precompiled_stdlib installs all. baml check floor: ~70ms -> ~35.7ms (stable, 82 runs); ~10x vs the original ~360ms. Verified: full baml_tests 2619/2619, corpus 1948/1948, run=232, objtest=250, user errors still report. Remaining toward ~15ms: cache PPIR package_items (the next ~10-19ms; 'db-bound, re-intern) -- diminishing returns.
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Bytecode emit (
generate_project_bytecode_with_opt) lowered every function twice:lower_functiononly to read each function's global-slot name (mir.item_ref) and its intrinsic/await-any skip flag (mir.kind).lower_functionis not memoized, so this doubled all MIR lowering work — including the ~676-function / 8,287-line stdlib, which is embedded as source (include_str!) and recompiled from scratch on every invocation. For a typical project this lowering is ~the entire fixed compile cost.Change
Derive both of Pass 1's needs without lowering:
def_to_item_ref(Definition::Function(func_loc))— identical tomir.item_refin everylower_functionbranch (Expr / Builtin / Missing all setitem_refto exactly this).FunctionBody::Builtinever lowers toMirFunctionKind::Builtin(Expr/Missing always lower toBytecode), somatches!(function_body(..), FunctionBody::Builtin(Intrinsic | AwaitAny))is equivalent to the oldmatches!(mir.kind, ...).Pass 4 is untouched and still lowers to emit code, so the emitted
Programis byte-identical.function_bodyis the same salsa querylower_functionalready calls first, so Pass 1 now does no redundant work.Results
Clean-compile wall time (
baml check, median of 9 runs), same fat-LTO release profile before/after:The bytecode stage drops from ~0.25s to ~0.17s. Behavioral output verified identical on a recursive + loop program (result
232before and after).Follow-up
The remaining fixed cost is the stdlib being recompiled from embedded source on every run; precompiling it to bytecode once is a separate, larger win tracked next.
Summary by CodeRabbit
New Features
Bug Fixes
Other