Lower lambdas in surrounding function - #4282
Conversation
Lambdas lack a lot of the stuff needed by `FunctionDef`
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
⏭️ 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):
|
📝 WalkthroughWalkthroughLambda bodies now use ChangesLambda arena migration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Source
participant ASTLowering
participant ExprBody
participant TIRInference
participant LSP
Source->>ASTLowering: lower lambda expression
ASTLowering->>ExprBody: allocate lambda body nodes
ASTLowering->>TIRInference: provide LambdaDef and body root
TIRInference->>ExprBody: infer reachable nodes
LSP->>ExprBody: index reachable non-lambda nodes
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 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 |
Binary size checks passed✅ 7 passed
Generated by |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
baml_language/crates/baml_project/src/db.rs (1)
775-775: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
call_sites_by_source_exprnow scans the whole shared arena instead of the current test/lambda's own subtree.Since lambda bodies (and, for tests, every sibling
test { ... }registration in a file) now share oneExprBodyarena,call_sites_by_source_expr's flatbody.exprs.iter()scan is no longer scoped to the callable being visualized. At Line 775,test_bodyis the entire$init_test_*registration body for the file (every test's lambda body lives in it), so building the CFG for one test now re-scans every other test's body too. At Line 861 the same helper additionally now walks any lambdas nested inside the function being graphed (previously excluded, since lambda bodies were separate arenas).This doesn't produce wrong output —
expand_user_function_calls_in_graph's subsequentgraph.nodeslookup still filters out call sites that aren't part of the bounded CFG — but it turns "visualize one test" into O(all tests' combined body size) instead of O(that test's body size), which is quadratic across a file with many tests in the playground.Bound the scan the same way
tokens/index.rs/usages.rsalready do for this migration, usingExprBody::reachable_excluding_lambdasfrom an explicit root:⚡ Proposed fix: bound the scan by root
-fn call_sites_by_source_expr(body: &baml_compiler2_ast::ExprBody) -> Vec<(u32, String)> { - use baml_compiler2_ast::Expr; - - let mut calls = Vec::new(); - for (expr_id, expr) in body.exprs.iter() { - let (Expr::Call { callee, .. } | Expr::OptionalCall { callee, .. }) = expr else { - continue; - }; - - let Expr::Path(segments) = &body.exprs[*callee] else { - continue; - }; - - let callee_name = segments - .iter() - .map(AsRef::<str>::as_ref) - .collect::<Vec<_>>() - .join("."); - calls.push((expr_id.into_raw().into_u32(), callee_name)); - } - calls -} +fn call_sites_by_source_expr( + body: &baml_compiler2_ast::ExprBody, + root: Option<baml_compiler2_ast::ExprId>, +) -> Vec<(u32, String)> { + use baml_compiler2_ast::{BodyNode, Expr}; + + let mut calls = Vec::new(); + let Some(root) = root else { + return calls; + }; + for node in body.reachable_excluding_lambdas(root) { + let BodyNode::Expr(expr_id) = node else { + continue; + }; + let (Expr::Call { callee, .. } | Expr::OptionalCall { callee, .. }) = &body.exprs[expr_id] + else { + continue; + }; + + let Expr::Path(segments) = &body.exprs[*callee] else { + continue; + }; + + let callee_name = segments + .iter() + .map(AsRef::<str>::as_ref) + .collect::<Vec<_>>() + .join("."); + calls.push((expr_id.into_raw().into_u32(), callee_name)); + } + calls +}Then thread a
root: Option<ExprId>throughexpand_user_function_calls_in_graphand passexpr_body.root_exprat Line 861 andtest_lambda.bodyat Line 775.Also applies to: 861-861, 879-963, 998-1019
🤖 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` at line 775, Bound call-site scanning in expand_user_function_calls_in_graph to the callable’s subtree by adding a root: Option<ExprId> parameter and using ExprBody::reachable_excluding_lambdas from that root instead of scanning the full arena. Pass test_lambda.body when processing test bodies and expr_body.root_expr when processing functions, and thread the root through all affected callers and helper logic.
🧹 Nitpick comments (2)
baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs (1)
686-688: 📐 Maintainability & Code Quality | 🔵 TrivialTracked
BUG:— doubly-nestedtestis dropped without a diagnostic.
testset "A" { test "B" { test "C" {} } }compiles clean whileCnever runs. The newa_test_inside_a_test_does_not_registertest pins the current behavior, so this is a known gap rather than a regression.Want me to open an issue to add a
NestedTestNotRegisteredlowering diagnostic on this arm?🤖 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/lower_expr_body.rs` around lines 686 - 688, Track this known gap by opening an issue for a NestedTestNotRegistered lowering diagnostic when processing a test nested inside another test in the relevant lowering arm. Preserve the current behavior and the a_test_inside_a_test_does_not_register test; do not change registration or add unrelated diagnostics.baml_language/crates/baml_compiler2_tir/src/throw_inference.rs (1)
419-443: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueNit:
body_nodesis walked twice per call here.
collect_catch_arm_bodies(body)(line 410) already builds the full node list, and this loop rebuilds it. Computing it once and passing the slice to both would halve the traversal + allocation on every function.🤖 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_tir/src/throw_inference.rs` around lines 419 - 443, Reuse the node list produced by collect_catch_arm_bodies in the throw-inference flow instead of calling body_nodes(body) again. Update collect_catch_arm_bodies and its caller as needed to pass the shared slice to both catch-arm analysis and the subsequent throw scan, preserving the existing filtering and fact insertion behavior.
🤖 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/ast.rs`:
- Around line 903-906: Update the stale lambda documentation at
baml_language/crates/baml_compiler2_ast/src/ast.rs:903-906 by removing
references to FunctionDef and FunctionBodyDef::Expr, and describe Expr::Lambda
as referencing LambdaDef with a body stored as ExprId in the shared arena. Also
revise lower_spawn_expr documentation at
baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs:854-869 to
describe lowering a synthesized LambdaDef into the current arena and replace the
Expr::Lambda(func_def) wording with the LambdaDef payload terminology.
In `@baml_language/crates/baml_lsp2_actions/src/annotations.rs`:
- Around line 246-249: Update the comment above the synthesized test/testset
registration skip in the annotation traversal to remove the reference to the
nonexistent Expr::Lambda arm. Describe that lambda arguments are still
recursively visited because their bodies reside in the same arena and are
processed by the surrounding loop.
---
Outside diff comments:
In `@baml_language/crates/baml_project/src/db.rs`:
- Line 775: Bound call-site scanning in expand_user_function_calls_in_graph to
the callable’s subtree by adding a root: Option<ExprId> parameter and using
ExprBody::reachable_excluding_lambdas from that root instead of scanning the
full arena. Pass test_lambda.body when processing test bodies and
expr_body.root_expr when processing functions, and thread the root through all
affected callers and helper logic.
---
Nitpick comments:
In `@baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs`:
- Around line 686-688: Track this known gap by opening an issue for a
NestedTestNotRegistered lowering diagnostic when processing a test nested inside
another test in the relevant lowering arm. Preserve the current behavior and the
a_test_inside_a_test_does_not_register test; do not change registration or add
unrelated diagnostics.
In `@baml_language/crates/baml_compiler2_tir/src/throw_inference.rs`:
- Around line 419-443: Reuse the node list produced by collect_catch_arm_bodies
in the throw-inference flow instead of calling body_nodes(body) again. Update
collect_catch_arm_bodies and its caller as needed to pass the shared slice to
both catch-arm analysis and the subsequent throw scan, preserving the existing
filtering and fact insertion 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 Plus
Run ID: 00b1ed2d-a2f8-43f8-8344-4ec151201ae5
⛔ Files ignored due to path filters (1)
baml_language/crates/baml_tests/snapshots/diagnostic_errors/void_return_type/baml_tests__diagnostic_errors__void_return_type__05_diagnostics.snapis excluded by!**/*.snap
📒 Files selected for processing (24)
baml_language/ARCHITECTURE.mdbaml_language/crates/baml_compiler2_ast/src/ast.rsbaml_language/crates/baml_compiler2_ast/src/disambiguate.rsbaml_language/crates/baml_compiler2_ast/src/lib.rsbaml_language/crates/baml_compiler2_ast/src/lower_cst.rsbaml_language/crates/baml_compiler2_ast/src/lower_expr_body.rsbaml_language/crates/baml_compiler2_ast/src/traverse.rsbaml_language/crates/baml_compiler2_hir/src/builder.rsbaml_language/crates/baml_compiler2_mir/src/lower.rsbaml_language/crates/baml_compiler2_tir/src/builder.rsbaml_language/crates/baml_compiler2_tir/src/inference.rsbaml_language/crates/baml_compiler2_tir/src/throw_inference.rsbaml_language/crates/baml_compiler2_tir/src/throws_analysis.rsbaml_language/crates/baml_compiler2_visualization/src/control_flow/from_ast.rsbaml_language/crates/baml_compiler2_visualization/src/control_flow/mod.rsbaml_language/crates/baml_lsp2_actions/src/annotations.rsbaml_language/crates/baml_lsp2_actions/src/completions.rsbaml_language/crates/baml_lsp2_actions/src/describe.rsbaml_language/crates/baml_lsp2_actions/src/tokens/index.rsbaml_language/crates/baml_lsp2_actions/src/usages.rsbaml_language/crates/baml_project/src/db.rsbaml_language/crates/baml_tests/src/compiler2_tir/inference.rsbaml_language/crates/baml_tests/src/compiler2_tir/mod.rsbaml_language/crates/baml_tests/src/compiler2_tir/phase8_exceptions.rs
| /// Lambda expression: anonymous function in expression position. | ||
| /// Reuses `FunctionDef` with synthetic name `"<anonymous function>"`. | ||
| /// The lambda's body gets its own `ExprBody` via `FunctionBodyDef::Expr`. | ||
| Lambda(Box<FunctionDef>), | ||
| Lambda(Box<LambdaDef>), |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Doc comments left behind by the shared-arena migration. Both sites still describe lambdas as reusing FunctionDef and owning their own ExprBody/source map — precisely the model this PR replaces. The code is correct; only the prose is stale, and these are the two places a reader looks first to learn the lambda representation.
baml_language/crates/baml_compiler2_ast/src/ast.rs#L903-L906: drop theFunctionDef/FunctionBodyDef::Exprsentences onExpr::Lambdaand point at [LambdaDef], noting the body is anExprIdin this arena.baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs#L854-L869: rewordlower_spawn_expr's doc so the synthesized body is aLambdaDeflowered into the current arena, not "a fresh 0-arg lambda (its ownExprBody+ source map)"; also update "Expr::Lambda(func_def)" to reflect theLambdaDefpayload.
📍 Affects 2 files
baml_language/crates/baml_compiler2_ast/src/ast.rs#L903-L906(this comment)baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs#L854-L869
🤖 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/ast.rs` around lines 903 - 906,
Update the stale lambda documentation at
baml_language/crates/baml_compiler2_ast/src/ast.rs:903-906 by removing
references to FunctionDef and FunctionBodyDef::Expr, and describe Expr::Lambda
as referencing LambdaDef with a body stored as ExprId in the shared arena. Also
revise lower_spawn_expr documentation at
baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs:854-869 to
describe lowering a synthesized LambdaDef into the current arena and replace the
Expr::Lambda(func_def) wording with the LambdaDef payload terminology.
| // Skip synthesized test/testset registration calls — their | ||
| // `name` / `body` / `collector` / `runner` arguments are codegen, | ||
| // not user-facing. We still recurse into their lambda arguments | ||
| // (the actual test bodies) via the `Expr::Lambda` arm below. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale comment: the Expr::Lambda arm it references no longer exists.
Lambda arguments are now reached because their bodies live in this same arena and are visited by this loop — not "via the Expr::Lambda arm below".
📝 Proposed comment fix
// Skip synthesized test/testset registration calls — their
// `name` / `body` / `collector` / `runner` arguments are codegen,
- // not user-facing. We still recurse into their lambda arguments
- // (the actual test bodies) via the `Expr::Lambda` arm below.
+ // not user-facing. Their lambda arguments (the actual test bodies)
+ // are still covered: those expressions live in this same arena and
+ // are visited by this loop.📝 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.
| // Skip synthesized test/testset registration calls — their | |
| // `name` / `body` / `collector` / `runner` arguments are codegen, | |
| // not user-facing. We still recurse into their lambda arguments | |
| // (the actual test bodies) via the `Expr::Lambda` arm below. | |
| // Skip synthesized test/testset registration calls — their | |
| // `name` / `body` / `collector` / `runner` arguments are codegen, | |
| // not user-facing. Their lambda arguments (the actual test bodies) | |
| // are still covered: those expressions live in this same arena and | |
| // are visited by this loop. |
🤖 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_lsp2_actions/src/annotations.rs` around lines 246 -
249, Update the comment above the synthesized test/testset registration skip in
the annotation traversal to remove the reference to the nonexistent Expr::Lambda
arm. Describe that lambda arguments are still recursively visited because their
bodies reside in the same arena and are processed by the surrounding loop.
…n TIR
Probed the ticket's bad_pop repro against current TIR: pop() on a
generic receiver types T | null (both engines agree), but TIR still
accepts the single-arm match { let x: T => x } over T | null with no
diagnostic - the soundness hole lives in exhaustiveness, and its fix
PR #3904 was parked for this rework. The bad_pop fixture lands at S10
(patterns) as tir: fails. B-1010 was fixed in TIR by the base commit's
lambda rework (#4282); the ticket is stale against canary.
…n TIR
Probed the ticket's bad_pop repro against current TIR: pop() on a
generic receiver types T | null (both engines agree), but TIR still
accepts the single-arm match { let x: T => x } over T | null with no
diagnostic - the soundness hole lives in exhaustiveness, and its fix
PR #3904 was parked for this rework. The bad_pop fixture lands at S10
(patterns) as tir: fails. B-1010 was fixed in TIR by the base commit's
lambda rework (#4282); the ticket is stale against canary.
1 participant
Lambdas use their parent function's
TypeExpr/TypeRefarenas, now have a separateLambdaDefAST type and generally lower in their surrounding function rather than as a separate system. This improves the internal representation and also is part of the groundwork for a rust-like inference system.Summary by CodeRabbit
Improvements
Documentation
Tests