Skip to content

fix: thread inferred generic type args to callee frames at runtime#3698

Closed
hellovai wants to merge 2 commits into
canaryfrom
hellovai/iterators2
Closed

fix: thread inferred generic type args to callee frames at runtime#3698
hellovai wants to merge 2 commits into
canaryfrom
hellovai/iterators2

Conversation

@hellovai

@hellovai hellovai commented Jun 4, 2026

Copy link
Copy Markdown
Contributor

What was wrong

A generic call whose type args are inferred rather than written ran its callee with T = unknown. TIR inferred the bindings, used them to compute the result type, then threw them away; MIR seeded frame.type_args from receiver class args ++ explicit <...> args only. On current canary (d968182, verified pre-merge):

Program Wrong (canary) Right (this PR)
function tn<T>(a: T[]) -> string { reflect.type_of<T>().to_string() }tn([1,2,3]) "unknown" "int"
Same, class static: Holder.tn([1,2,3]) "unknown" "int"
let it: IterA<int, never> = ArrA.of([1,2,3]); it.collect().length() (2 implementors) 0 — wrong implementor runs 3
match (ArrBag.make([5,6,7])) { let a: ArrBag<int> => ... } falls to _ arm — instance carries class_type_args=[unknown] matches ArrBag<int>
fwd<int>(true || false) where fwd<T>(x: bool) -> bool { x } null (pre-existing emitter miscompile, surfaced by this work) true

The dispatch fallout is the nasty part: instances built by inferred static ctors are born with class_type_args = [unknown], every pinned is_type C<int> dispatch guard fails, and the interface-dispatch switch's last candidate arm is an unguarded else — so whichever implementor is last in hash order silently runs. That made the bug look "fragile" (renames, extra declarations in the file, or method names flip the arm order and mask it), and it is why #3667's regression suite stayed green: those tests pass by luck of arm order — their instances carry the same broken [unknown] args (now pinned by a direct IsType regression test that fails without this fix).

What the fix does

1. TIR (baml_compiler2_tir): check_call_inner records the post-inference bindings per call expr in call_inferred_type_args (sorted Vec<(Name, Ty)>), plumbed through ScopeInference / DefaultParameterInference with the same save/restore discipline as call_plans (parameter defaults + inline lambda inference).

2. MIR (baml_compiler2_mir): at the call-lowering merge point, emit_inferred_call_type_arg_ops emits those bindings as leading LoadType frame args in the callee's De Bruijn order (callee_frame_generic_params = owner params ++ fn params, mirroring enclosing_generic_params). Conservative gates preserve the status quo when seeding can't be done soundly: explicit <...> args, sys-ops (positional Rust glue), function values (no resolution), misaligned receiver prefixes, all-unknown seeds. Bindings nesting a TypeVar outside the caller's generic params — a rigid Self from a Self-pinned call — are demoted to unknown: naive templating lowers them to void, which would seed Self[] as void[] (caught by adversarial review; guarded by nested_self_inferred_binding_not_voided).

Also: the path-callee branch no longer prepends a stray Constant::Null receiver to self-less static methods (ArrA.of(...)). The VM's arity-based pop silently tolerated the orphan, but it corrupts the leading type-arg window once calls carry seeds.

3. EMIT (baml_compiler2_emit): pre-existing ShortCircuit miscompile, reproducible on canary without this PR. The emitter assumed the destination local was always stack-carried; when the stack-carry pass rejects the carry (e.g. a LoadType operand pushes before the use), the taken path never stored the slot while eval_rhs did — a half-applied phi, so a || b fed into such a call reads an unwritten local. The taken path now stores whenever the destination is not stack-carried (&& via a small trampoline). Bonus: the previously-failing if_else_assigned_then_passed_to_call shape now passes.

Verification

  • Clean A/B on the same canary base: all five wrong behaviors above reproduce on origin/canary and flip with this PR.
  • In-BAML suite 1646/1646; MIR/codegen/bytecode/format snapshot targets green; workspace Rust tests green; clippy/fmt clean.
  • Snapshot churn is exactly the expected categories: +load_type seeds at inferred generic call sites, -const null stray receivers, store/load pairs where short-circuit carries are now correctly rejected.
  • 4 new regression tests in ns_inferred_generic_type_args (free-fn reflect, instance IsType after inferred ctor, the lazy-iterator collect shape with 2 implementors, nested-Self demotion).
  • Adversarially reviewed (multi-agent: TIR soundness / MIR De Bruijn alignment / emit fix / blast radius): 2 confirmed findings, both addressed in this PR.

Known remaining (pre-existing, intentionally out of scope)

  • The interface-dispatch switch does not thread a dispatched method's own inferred fn-level generics (interface method with its own <U> reading U at runtime).
  • Function values (let f = ArrA.of; f([1])) stay unseeded — no resolution at the call site.
  • The dispatch switch's last candidate arm is still an unguarded else; a guarded-miss panic would be a better failure mode than silent fallthrough if type args ever regress.
  • Test-block let bindings consumed by a later statement in the block lower to Constant(Null) — separate lowering bug discovered during this work.

Perf note

Every inferred generic call now carries LoadType seeds. The existing bench corpus is monomorphic and cannot observe this path, so a green bench run is not evidence — this PR adds a compute/generic-apply-inferred-1m speedtest workload that exercises it; a release-mode A/B (total cycles + allocs/iter) against canary is the meaningful pre-merge measurement.

🤖 Generated with Claude Code

Summary by CodeRabbit

Release Notes

  • New Features

    • Generic function and method calls now automatically infer missing type arguments at call sites and properly propagate them to callees at runtime, enabling more flexible generic code without explicit type specifications.
  • Tests

    • Added comprehensive test coverage for inferred generic type argument dispatch, including interface methods, default methods, and nested generic scenarios.

hellovai and others added 2 commits June 4, 2026 08:58
A generic call whose type args are inferred rather than written —
`ArrA.of([1, 2, 3])`, `tn([1, 2, 3])` — ran its callee with
`T = unknown`: TIR inferred the bindings but discarded them after
computing the result type, and MIR seeded `frame.type_args` from
receiver class args ++ explicit `<...>` args only. Instances built by
inferred static ctors were born with `class_type_args = [unknown]`,
every runtime `IsType C<int>` test against them failed, and the
interface-dispatch switch — whose last candidate arm is an unguarded
else — silently ran whichever implementor was last in hash order.
(#3667's regression suite passed by luck of that arm order; this is the
residual root cause behind the iterator-blocker bug class.)

Three-part fix:

* TIR: `check_call_inner` now records the post-inference bindings per
  call expr (`call_inferred_type_args`, sorted by param name), plumbed
  through `ScopeInference`/`DefaultParameterInference` with the same
  save/restore discipline as `call_plans`.
* MIR: `emit_inferred_call_type_arg_ops` seeds the callee frame with
  those bindings as leading `LoadType` args in the callee's De Bruijn
  order (`callee_frame_generic_params` = owner ++ fn params, mirroring
  `enclosing_generic_params`). Conservative gates: explicit args,
  sys-ops (positional Rust glue), function values, misaligned receiver
  prefixes, all-unknown seeds. Bindings nesting a TypeVar outside the
  caller's generic params (a rigid `Self`) are demoted to `unknown` —
  naive templating would lower them to `void`. Also stop prepending a
  stray `null` receiver to self-less static path calls (`ArrA.of(...)`);
  the VM's arity-based pop tolerated the orphan, but it corrupts the
  leading type-arg window.
* EMIT: fix a pre-existing ShortCircuit miscompile this surfaced
  (`fwd<int>(true || false)` returned null on canary): when the
  stack-carry pass rejects the destination carry, the taken path never
  stored the slot while `eval_rhs` did — a half-applied phi. The taken
  path now stores whenever the destination is not stack-carried.

Regression tests: free-fn reflect, instance `IsType` after inferred
ctor, the full lazy-iterator collect shape, and the nested-`Self`
demotion, plus a `generic-apply-inferred-1m` speedtest workload (the
existing bench corpus is monomorphic and cannot observe this path).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@vercel

vercel Bot commented Jun 4, 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 Jun 4, 2026 11:04pm
promptfiddle Ready Ready Preview, Comment Jun 4, 2026 11:04pm
promptfiddle2 Ready Ready Preview, Comment Jun 4, 2026 11:04pm

Request Review

@coderabbitai

coderabbitai Bot commented Jun 4, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

The compiler now records inferred generic type arguments at call sites during type inference and propagates them through MIR lowering to seed callee frames at runtime, enabling type dispatch and reflection in called functions. A short-circuit bytecode fix improves destination slot storage logic.

Changes

Inferred Generic Type Arguments Dispatch

Layer / File(s) Summary
TIR: Collect and record inferred call-site type arguments
baml_language/crates/baml_compiler2_tir/src/builder.rs, baml_language/crates/baml_compiler2_tir/src/inference.rs
TypeInferenceBuilder now collects inferred generic type-argument bindings (for calls that omit explicit <...> args) and persists them into a call_inferred_type_args map. During finish(), this map is returned alongside scope inference; ScopeInference threads it through Salsa construction and provides iterator helpers for MIR to query the bindings. Inferred maps are correctly saved/restored during nested parameter-default and lambda-body inference.
MIR: Thread inferred type-args through lowering context
baml_language/crates/baml_compiler2_mir/src/lower.rs (context setup)
LoweringContext adds a call_inferred_type_args field initialized from aggregated TIR scope results. The merge_scope closure in both LoweringContext::new and LoweringContext::new_for_let gathers iter_call_inferred_type_args() from body scopes, descendant scopes, and parameter-default scopes, making inferred bindings available during the entire function lowering pass.
MIR: Compute frame layout and emit inferred type-arg operands
baml_language/crates/baml_compiler2_mir/src/lower.rs (helpers)
callee_frame_generic_params() computes the De Bruijn layout of a callee's generic parameters by inspecting the item tree (distinguishing out-of-body impl methods, interface default methods, and class methods). emit_inferred_call_type_arg_ops() reads call_inferred_type_args, verifies receiver-prefix alignment, demotes non-resolvable bindings to unknown, and emits LoadType operands for remaining slots in the callee frame.
MIR: Refactor call lowering and integrate type-arg emission
baml_language/crates/baml_compiler2_mir/src/lower.rs (call lowering)
lower_call is refactored to correctly extract receiver and argument operands for method calls via MemberAccess and multi-segment Path callees, including a proper "does this method take self?" check. sys-op detection is hoisted earlier to skip inferred-type-arg seeding for system operations. Inferred LoadType operands are appended after receiver-provided class type args and explicit type args, enabling callee frame.type_args runtime resolution.
Bytecode: Short-circuit destination storage
baml_language/crates/baml_compiler2_emit/src/emit.rs
Terminator::ShortCircuit emission now inspects destination's LocalStoreBehavior to determine whether the taken-path result must be stored into a slot. For &&, a trampoline patches the JumpIfFalse landing to a post-true-path store block when needed; for `
Integration and tests
baml_language/crates/baml_lsp2_actions/src/check.rs, baml_language/crates/baml_tests/baml_src/ns_inferred_generic_type_args/inferred_generic_type_args.baml, baml_language/tools/speedtest/workloads/compute/generic-apply-inferred-1m.md
LSP type-checker adapts to the new finish() tuple. New end-to-end tests validate that inferred type-argument bindings are preserved into callee frames via reflect.type_of<T>() runtime reflection, match guards discriminating inferred static-constructor instances, interface-dispatched default methods, and nested self-type bindings. A benchmark exercises million-iteration inferred-call frame seeding.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

  • BoundaryML/baml#3667: Both PRs modify baml_compiler2_mir/src/lower.rs to ensure inferred generic type arguments are seeded into the dispatched callee frame via emitting LoadType operands, making interface dispatch part of the same inferred type-arg propagation mechanism.
  • BoundaryML/baml#3460: The MIR lowering changes seeding callee type_args via inferred generic call-site bindings directly complement the retrieved PR's reflect.type_of<T>() runtime reflection machinery that consumes those frame type_args.

Poem

🐰 Generics once lived in the dark, unknown and shy,
But now they're threaded through the lowering sky—
TIR whispers "here's the type!" to frames below,
While short-circuits store their values true and slow,
And reflect can see at last what was hid from eye! 🎯

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title directly and accurately summarizes the primary fix: threading inferred generic type arguments to callee frames at runtime, which is the core objective of this multi-file compiler bug fix.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ 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 hellovai/iterators2
⚔️ Resolve merge conflicts
  • Resolve merge conflict in branch hellovai/iterators2

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 and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
baml_language/crates/baml_compiler2_mir/src/lower.rs (1)

6057-6156: ⚡ Quick win

Add unit coverage for the inferred-call type-arg emitter.

The new ordering and alignment rules in emit_inferred_call_type_arg_ops are the core behavior change here, but the added unit tests only exercise interface guard helpers. A small in-module test for a static generic-class call and a receiver-seeded method call would lock in the owner/fn slot ordering and the misalignment gates cheaply. As per coding guidelines, "Prefer writing Rust unit tests over integration tests where possible."

🤖 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_mir/src/lower.rs` around lines 6057 -
6156, Add in-module Rust unit tests that directly exercise
emit_inferred_call_type_arg_ops to cover (1) a static generic-class call where
owner params precede fn params to assert the emitted frame type arg operands
follow owner-then-fn ordering, and (2) a receiver-seeded method call where
already_seeded equals owner params to assert proper emission and when
already_seeded is misaligned (non-zero but not equal to owner params) to assert
it returns empty. Use the same module-level helpers/state setup the file uses:
populate call_inferred_type_args for the expr_id, set
resolutions/path_member_resolutions so func_loc resolves to a function (so
callee_frame_generic_params returns owner_params and fn_params), and vary
already_seeded; assert on the returned Vec<Operand> contents and that
all-unknown or misaligned cases return Vec::new(). Ensure tests live in the same
file as the function (unit tests) and reference emit_inferred_call_type_arg_ops,
call_inferred_type_args, resolutions/path_member_resolutions, and
callee_frame_generic_params so they validate ordering and misalignment gates.
🤖 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.

Nitpick comments:
In `@baml_language/crates/baml_compiler2_mir/src/lower.rs`:
- Around line 6057-6156: Add in-module Rust unit tests that directly exercise
emit_inferred_call_type_arg_ops to cover (1) a static generic-class call where
owner params precede fn params to assert the emitted frame type arg operands
follow owner-then-fn ordering, and (2) a receiver-seeded method call where
already_seeded equals owner params to assert proper emission and when
already_seeded is misaligned (non-zero but not equal to owner params) to assert
it returns empty. Use the same module-level helpers/state setup the file uses:
populate call_inferred_type_args for the expr_id, set
resolutions/path_member_resolutions so func_loc resolves to a function (so
callee_frame_generic_params returns owner_params and fn_params), and vary
already_seeded; assert on the returned Vec<Operand> contents and that
all-unknown or misaligned cases return Vec::new(). Ensure tests live in the same
file as the function (unit tests) and reference emit_inferred_call_type_arg_ops,
call_inferred_type_args, resolutions/path_member_resolutions, and
callee_frame_generic_params so they validate ordering and misalignment gates.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 509f0e45-5708-494f-8305-7895a7348bf3

📥 Commits

Reviewing files that changed from the base of the PR and between 654e9c1 and 30c563f.

⛔ Files ignored due to path filters (39)
  • baml_language/crates/baml_tests/snapshots/baml_src/builtins.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/closures.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/deep_copy.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/fs.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/inferred_generic_type_args.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/instantiation_expr.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/json_auto_derive.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/lambdas.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/type_error_repro.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/baml_src/type_reflection.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/__baml_std__/baml_tests__compiles____baml_std____04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/__baml_std__/baml_tests__compiles____baml_std____06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/__testing_std__/baml_tests__compiles____testing_std____04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/__testing_std__/baml_tests__compiles____testing_std____06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/closures/baml_tests__compiles__closures__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/closures/baml_tests__compiles__closures__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/comment_in_type/baml_tests__compiles__comment_in_type__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/comment_in_type/baml_tests__compiles__comment_in_type__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/json_llm_return_type/baml_tests__compiles__json_llm_return_type__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/json_llm_return_type/baml_tests__compiles__json_llm_return_type__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/lambda_advanced/baml_tests__compiles__lambda_advanced__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/lambda_advanced/baml_tests__compiles__lambda_advanced__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/lambda_basic/baml_tests__compiles__lambda_basic__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/lambda_basic/baml_tests__compiles__lambda_basic__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/llm_image_outputs/baml_tests__compiles__llm_image_outputs__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/llm_image_outputs/baml_tests__compiles__llm_image_outputs__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/optional_function_parameters/baml_tests__compiles__optional_function_parameters__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/optional_function_parameters/baml_tests__compiles__optional_function_parameters__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/stream_llm_inferred_typeargs/baml_tests__compiles__stream_llm_inferred_typeargs__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/stream_llm_inferred_typeargs/baml_tests__compiles__stream_llm_inferred_typeargs__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/testset_vibes_nested/baml_tests__compiles__testset_vibes_nested__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/testset_vibes_nested/baml_tests__compiles__testset_vibes_nested__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/type_builder_errors/baml_tests__compiles__type_builder_errors__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/type_builder_errors/baml_tests__compiles__type_builder_errors__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/type_builder_test/baml_tests__compiles__type_builder_test__04_5_mir.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/snapshots/compiles/type_builder_test/baml_tests__compiles__type_builder_test__06_codegen.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/tests/bytecode_format/snapshots/bytecode_format__bytecode_display_expanded.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/tests/bytecode_format/snapshots/bytecode_format__bytecode_display_expanded_unoptimized.snap is excluded by !**/*.snap
  • baml_language/crates/baml_tests/tests/bytecode_format/snapshots/bytecode_format__bytecode_display_textual.snap is excluded by !**/*.snap
📒 Files selected for processing (7)
  • baml_language/crates/baml_compiler2_emit/src/emit.rs
  • baml_language/crates/baml_compiler2_mir/src/lower.rs
  • baml_language/crates/baml_compiler2_tir/src/builder.rs
  • baml_language/crates/baml_compiler2_tir/src/inference.rs
  • baml_language/crates/baml_lsp2_actions/src/check.rs
  • baml_language/crates/baml_tests/baml_src/ns_inferred_generic_type_args/inferred_generic_type_args.baml
  • baml_language/tools/speedtest/workloads/compute/generic-apply-inferred-1m.md

@aaronvg

aaronvg commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

superseded by this (which adds the type args in the CallPlan instead)
#3702

@aaronvg aaronvg closed this Jun 5, 2026
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.

2 participants