feat(trunk-ir): add arena IR transform passes and rewrite helpers (#439)#459
Merged
feat(trunk-ir): add arena IR transform passes and rewrite helpers (#439)#459
Conversation
- Add DCE pass using use-chain analysis to eliminate dead operations - Add Global DCE pass with call graph analysis for dead function removal - Add SCF-to-CF lowering pass using split_block/inline_region_blocks/RAUW - Add arena rewrite helpers: split_block, inline_region_blocks, erase_op - Add add_block_arg and detach_op methods to IrContext - Fix register_pure_op\!/register_isolated_op\! macros to use raw_ident_str\! instead of stringify\!, enabling correct DCE detection for ops like arith.const Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughPorts three transform passes (DCE, Global DCE, SCF→CF) to the arena IR, adds IrContext APIs for block/operation manipulation, introduces rewrite helpers for block/region mutation, and exposes new transforms and rewrite modules with unit tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant GlobalDCE as Global DCE Pass
participant CallGraph as Call Graph Builder
participant Reachability as Reachability Analyzer
participant Removal as Removal Engine
participant IrContext
Caller->>GlobalDCE: eliminate_dead_functions(ctx, module)
GlobalDCE->>CallGraph: scan module for func.call/func.tail_call/func.constant
CallGraph-->>GlobalDCE: call_graph
GlobalDCE->>Reachability: compute_reachability(entry_points, call_graph)
Reachability-->>GlobalDCE: reachable_functions
GlobalDCE->>Removal: identify unreachable func.func ops
Removal->>IrContext: erase_op(unreachable)
IrContext-->>Removal: op removed
Removal-->>GlobalDCE: removed list/count
GlobalDCE-->>Caller: GlobalDceResult
sequenceDiagram
participant Caller
participant ScfToCf as SCF→CF Lowering
participant BlockOps as Block/Region Helpers
participant OpGen as CF Op Generator
participant Remapper as RAUW Remapper
participant IrContext
Caller->>ScfToCf: lower_scf_to_cf(ctx, module)
ScfToCf->>IrContext: traverse regions/blocks
ScfToCf->>BlockOps: split_block(block, scf_op)
BlockOps-->>ScfToCf: tail/merge blocks
ScfToCf->>BlockOps: inline_region_blocks(then/else/body)
BlockOps-->>ScfToCf: inlined blocks
ScfToCf->>OpGen: emit cf.br / cf.cond_br (replace yields/continues)
OpGen-->>ScfToCf: cf ops inserted
ScfToCf->>Remapper: replace_all_uses(scf_results, merge_args)
Remapper-->>IrContext: values remapped
ScfToCf-->>Caller: lowered module
Estimated code review effort🎯 5 (Critical) | ⏱️ ~100+ minutes Possibly related issues
Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #459 +/- ##
=======================================
Coverage 78.54% 78.54%
=======================================
Files 10 10
Lines 4935 4935
=======================================
Hits 3876 3876
Misses 1059 1059
🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
crates/trunk-ir/src/arena/rewrite/helpers.rs (1)
79-109: Guardinline_region_blocksagainstsrc_region == dest_region.With
insert_before: Some(_), self-inlining can currently clear the list and then panic on missing anchor. Add an explicit precondition or dedicated no-op path.Proposed guard
pub fn inline_region_blocks( ctx: &mut IrContext, src_region: RegionRef, dest_region: RegionRef, insert_before: Option<BlockRef>, ) -> Vec<BlockRef> { + assert!( + src_region != dest_region, + "inline_region_blocks: src_region and dest_region must differ" + ); // Take blocks from src let src_blocks: SmallVec<[BlockRef; 4]> = std::mem::take(&mut ctx.region_mut(src_region).blocks);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/trunk-ir/src/arena/rewrite/helpers.rs` around lines 79 - 109, inline_region_blocks must guard against src_region == dest_region to avoid clearing the region then panicking when insert_before anchor is missing; add an early no-op return right at the start of inline_region_blocks (before calling std::mem::take or ctx.region_mut) that checks if src_region == dest_region and returns an empty Vec<BlockRef>, so no mutation or insertion logic runs when inlining into the same region (ensure the check covers both insert_before Some(_) and None cases).crates/trunk-ir/src/arena/transforms/scf_to_cf.rs (2)
107-115: Consider cachingSymbolconstants for repeated checks.
Symbol::new("scf"),Symbol::new("if"), etc. are created on each call. IfSymbol::newdoesn't intern/cache, this could add overhead. Consider usingconstorstaticsymbols if the pattern is hot.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/trunk-ir/src/arena/transforms/scf_to_cf.rs` around lines 107 - 115, The is_scf_control_flow function repeatedly constructs Symbol::new("scf"), Symbol::new("if"), Symbol::new("loop"), and Symbol::new("switch") on each call; change these to cached symbols (e.g., module-level const/static values using a const fn or a Lazy/static initializer such as once_cell/lazy_static) and use those cached identifiers in the comparisons inside is_scf_control_flow to avoid reallocating Symbols on every invocation (update the symbol names you reference: SCF, IF, LOOP, SWITCH or similar and replace the Symbol::new(...) uses in that function).
263-265: Result type inference only checks first case.
find_yield_typeis called only oncases.first()anddefault_region. If the first case happens to have an empty yield but subsequent cases have values, the result type could be missed. This might be fine if well-formed IR guarantees all cases yield the same type or none, but worth noting.Consider checking all case regions
- let result_ty = find_yield_type(ctx, cases.first().map(|(_, r)| *r)) - .or_else(|| find_yield_type(ctx, default_region)); + let result_ty = cases + .iter() + .find_map(|(_, r)| find_yield_type(ctx, Some(*r))) + .or_else(|| find_yield_type(ctx, default_region));🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/trunk-ir/src/arena/transforms/scf_to_cf.rs` around lines 263 - 265, The current result_ty computation only calls find_yield_type on cases.first(), which misses yields in later cases; update it to scan all case regions (the cases collection) and pick the first non-None yield type returned by find_yield_type, falling back to find_yield_type(ctx, default_region) if none found. Locate the sites using cases and find_yield_type (and the result_ty binding) and replace the single-case lookup with an iterator-based search (e.g., iterate cases -> extract each region from the (_, r) tuple -> call find_yield_type for each until Some) so the inferred result_ty reflects any case that yields a value.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@crates/trunk-ir/src/arena/rewrite/helpers.rs`:
- Around line 79-109: inline_region_blocks must guard against src_region ==
dest_region to avoid clearing the region then panicking when insert_before
anchor is missing; add an early no-op return right at the start of
inline_region_blocks (before calling std::mem::take or ctx.region_mut) that
checks if src_region == dest_region and returns an empty Vec<BlockRef>, so no
mutation or insertion logic runs when inlining into the same region (ensure the
check covers both insert_before Some(_) and None cases).
In `@crates/trunk-ir/src/arena/transforms/scf_to_cf.rs`:
- Around line 107-115: The is_scf_control_flow function repeatedly constructs
Symbol::new("scf"), Symbol::new("if"), Symbol::new("loop"), and
Symbol::new("switch") on each call; change these to cached symbols (e.g.,
module-level const/static values using a const fn or a Lazy/static initializer
such as once_cell/lazy_static) and use those cached identifiers in the
comparisons inside is_scf_control_flow to avoid reallocating Symbols on every
invocation (update the symbol names you reference: SCF, IF, LOOP, SWITCH or
similar and replace the Symbol::new(...) uses in that function).
- Around line 263-265: The current result_ty computation only calls
find_yield_type on cases.first(), which misses yields in later cases; update it
to scan all case regions (the cases collection) and pick the first non-None
yield type returned by find_yield_type, falling back to find_yield_type(ctx,
default_region) if none found. Locate the sites using cases and find_yield_type
(and the result_ty binding) and replace the single-case lookup with an
iterator-based search (e.g., iterate cases -> extract each region from the (_,
r) tuple -> call find_yield_type for each until Some) so the inferred result_ty
reflects any case that yields a value.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
crates/trunk-ir/src/arena/context.rscrates/trunk-ir/src/arena/mod.rscrates/trunk-ir/src/arena/rewrite/helpers.rscrates/trunk-ir/src/arena/rewrite/mod.rscrates/trunk-ir/src/arena/transforms/dce.rscrates/trunk-ir/src/arena/transforms/global_dce.rscrates/trunk-ir/src/arena/transforms/mod.rscrates/trunk-ir/src/arena/transforms/scf_to_cf.rscrates/trunk-ir/src/op_interface.rs
- `inline_region_blocks`: add early-return guard when `src_region == dest_region` to prevent a panic caused by aliased mutable borrow of the same region; add test `inline_region_blocks_same_region_is_noop` - `lower_scf_switch`: scan all case regions (not just the first) when inferring `result_ty` from yield values, so a yield in any case is found even when earlier cases are void; add test `lower_scf_switch_yield_in_later_case` Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
crates/trunk-ir/src/arena/transforms/scf_to_cf.rs (1)
139-143: Drop the staleremove_op_from_blockcall inlower_scf_if.After
split_block,scf_opis no longer inblock;detach_op(scf_op)already handles removal from its actual parent. Keeping both calls obscures intent and risks future maintenance mistakes.♻️ Proposed cleanup
@@ - // Remove the scf op from the original block - ctx.remove_op_from_block(block, scf_op); @@ - // Remove the scf op (split_block moved it to merge_block) + // Remove the scf op (split_block moved it to merge_block) ctx.detach_op(scf_op);Also applies to: 159-160
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/trunk-ir/src/arena/transforms/scf_to_cf.rs` around lines 139 - 143, The call to ctx.remove_op_from_block(block, scf_op) in lower_scf_if is stale because split_block moves scf_op out of block and detach_op(scf_op) already removes the op from its real parent; remove the redundant ctx.remove_op_from_block(...) invocations (the occurrences after split_block and the similar one around lines with split_block/detach_op) so only detach_op(scf_op) handles removal and intent is clear.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/trunk-ir/src/arena/transforms/scf_to_cf.rs`:
- Around line 263-337: The code can create a merge block expecting an argument
(result_ty Some) while some branches emit zero args, producing malformed
cf.br→merge edges; before inlining/emitting branches (around
split_block/merge_block and calls to inline_region_blocks and
replace_yield_with_br) validate that when result_ty.is_some() every case region
and the default region yields exactly one value of that type (use
find_yield_type on each region or inspect inlined blocks), and if any region
yields a mismatched count or type, fail fast (return an Err / emit a diagnostic)
instead of emitting cf::brs with no args; alternatively, if intended to
tolerate, ensure you synthesize a proper SSA value and pass it to cf::br when
creating default_block or replacing yields so branch arg counts always match
merge_block's BlockArg.
---
Nitpick comments:
In `@crates/trunk-ir/src/arena/transforms/scf_to_cf.rs`:
- Around line 139-143: The call to ctx.remove_op_from_block(block, scf_op) in
lower_scf_if is stale because split_block moves scf_op out of block and
detach_op(scf_op) already removes the op from its real parent; remove the
redundant ctx.remove_op_from_block(...) invocations (the occurrences after
split_block and the similar one around lines with split_block/detach_op) so only
detach_op(scf_op) handles removal and intent is clear.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
crates/trunk-ir/src/arena/rewrite/helpers.rscrates/trunk-ir/src/arena/transforms/scf_to_cf.rs
Abi functions (marked with the `abi` attribute) are externally callable and must be treated as reachability roots so that their callees are not incorrectly eliminated. Previously, abi functions were preserved by a special guard in `filter_region`, but that guard only kept the functions themselves — any function called exclusively from an abi function was still silently removed. Fix the root cause: add abi functions to `reachability_roots` during analysis so the BFS traversal naturally keeps their entire transitive call graph alive. The now-redundant abi preservation guard in `filter_region` is removed. Rename the internal `entry_points` field to `reachability_roots` to reflect that abi functions, main/wasm exports, and custom roots are all treated uniformly as BFS starting points. The public API `GlobalDceConfig::extra_entry_points` is unchanged. Add test `abi_function_callees_are_reachable` to cover the fixed case. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Remove stale `remove_op_from_block` call in `lower_scf_if`: after `split_block`, the scf op is already moved to the merge block, so the removal was a silent no-op; only `detach_op` is needed. - Fix `lower_scf_switch` result_ty determination: infer from the op's actual results instead of scanning yield ops, consistent with `lower_scf_if`. Since `scf.switch` has no result, `result_ty` is always `None`, eliminating the merge block / branch arg count mismatch. - Make `replace_yield_with_br` defensive by truncating yield values to the target block's expected arg count, preventing malformed `cf.br` ops. - Remove unused `find_yield_type` helper. - Update test: rename to `lower_scf_switch_no_result` and correct the assertion from 1 merge block arg to 0. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
crates/trunk-ir/src/arena/transforms/global_dce.rs (1)
377-391: Consider documenting the memory trade-off.The comment explains why
remove_opisn't called, but orphaned ops remain in the arena. This is acceptable for a compilation pass (arena is discarded after use), but consider adding a brief note about expected arena lifetime to help future maintainers.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/trunk-ir/src/arena/transforms/global_dce.rs` around lines 377 - 391, Add a brief note explaining the arena lifetime and memory trade-off where ops are detached but not freed: update the comment inside global_dce.rs next to the block that detaches funcs (the if branch that checks dialect == self.syms.dialect_func, name == self.syms.func, uses extract_func_name, pushes to removed, and calls ctx.remove_op_from_block) to state that ctx.remove_op is intentionally not used to avoid region/result ownership complexity, that orphaned ops remain in the arena but this is acceptable because the arena is short-lived and discarded after this compilation pass, and mention that if the arena lifetime changes future maintainers must revisit using ctx.remove_op for proper reclamation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@crates/trunk-ir/src/arena/transforms/global_dce.rs`:
- Around line 377-391: Add a brief note explaining the arena lifetime and memory
trade-off where ops are detached but not freed: update the comment inside
global_dce.rs next to the block that detaches funcs (the if branch that checks
dialect == self.syms.dialect_func, name == self.syms.func, uses
extract_func_name, pushes to removed, and calls ctx.remove_op_from_block) to
state that ctx.remove_op is intentionally not used to avoid region/result
ownership complexity, that orphaned ops remain in the arena but this is
acceptable because the arena is short-lived and discarded after this compilation
pass, and mention that if the arena lifetime changes future maintainers must
revisit using ctx.remove_op for proper reclamation.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
crates/trunk-ir/src/arena/transforms/global_dce.rscrates/trunk-ir/src/arena/transforms/scf_to_cf.rs
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
1 participant
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.
Summary
Implements the arena IR transform passes for GitHub Issue #439, building on the arena-based IR core (Milestone 1, #452) and rewrite infrastructure (#437/#458).
nested::maincorrectly matches"main")scf.if,scf.while) to unstructured control flow usingsplit_block,inline_region_blocks, and RAUWsplit_block,inline_region_blocks, anderase_oputilities incrates/trunk-ir/src/arena/rewrite/helpers.rsadd_block_arganddetach_opmethods for block/operation manipulationregister_pure_op!/register_isolated_op!macros now useraw_ident_str!instead ofstringify!, enabling correct pure-op detection for raw identifier ops likearith.constBug Fixes Included
detach_opinstead of incorrect block removal aftersplit_blocknested::mainArenaModule: private field access made safe withOption-returningbody()TypeConverter::is_empty()now correctly considers materializer stateTest Plan
trunk-irtests pass (cargo nextest run -p trunk-ir)cargo clippy -p trunk-ir)scf.ifandscf.whilelowered correctly to basic blocksCloses #439
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Tests