feat(trunk-ir): add arena IR rewrite infrastructure (#437)#458
Merged
Conversation
Implements walk utilities and a full pattern-based rewrite framework for the arena IR, analogous to the existing Salsa-based infrastructure: - arena/walk.rs: walk_region/walk_block/walk_op/walk_typed with WalkAction (Advance/Skip) control via ControlFlow - arena/rewrite/pattern.rs: ArenaRewritePattern trait - arena/rewrite/rewriter.rs: PatternRewriter with RAUW-based mutations (replace_op, insert_op, erase_op, add_module_op) - arena/rewrite/type_converter.rs: ArenaTypeConverter for dialect lowering - arena/rewrite/conversion_target.rs: ArenaConversionTarget with legality checks (legal/illegal/dynamically-legal operations) - arena/rewrite/applicator.rs: PatternApplicator with fixpoint iteration - op_interface.rs: is_pure_arena, is_removable_arena, is_isolated_arena helpers delegating to the existing registry-based checks Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro 📒 Files selected for processing (9)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds an arena-based rewrite and traversal subsystem: pattern applicator and rewriter with fixpoint iteration, conversion target legality checks, a pluggable type converter/materializer, recursive IR walkers, ArenaModule helpers, and TypeData builder plus related arena query APIs. Changes
Sequence Diagram(s)sequenceDiagram
participant App as PatternApplicator
participant Ctx as IrContext
participant Pat as ArenaRewritePattern
participant Rw as PatternRewriter
participant Conv as ArenaConversionTarget
App->>App: Snapshot top-level ops
loop iterate until fixpoint or max_iterations
App->>Ctx: traverse module body (walk)
loop per op
App->>Pat: match_and_rewrite(&mut ctx, op, rewriter)
alt pattern matches
Pat->>Rw: insert/replace/erase/add_module_op
Rw-->>App: mutations accumulated
App->>Ctx: apply_mutations(ctx, op, mutations)
else no match
Pat-->>App: no-op
end
end
end
App->>Conv: verify(&ctx, module_body)
Conv-->>App: list of IllegalOp (if any)
sequenceDiagram
participant User as Caller
participant TC as ArenaTypeConverter
participant ConvFn as ConversionFn
participant Mat as Materializer
participant Ctx as IrContext
User->>TC: convert_type(ctx, ty)
TC->>ConvFn: try conversion f(ctx, ty)
alt conversion returns Some
ConvFn-->>TC: Some(new_ty)
TC-->>User: return Some(new_ty)
else none match
ConvFn-->>TC: None
TC-->>User: None
end
User->>TC: materialize(ctx, loc, value, from, to)
TC->>Mat: call materializer(...)
Mat->>Ctx: create ops, new value
Ctx-->>Mat: MaterializeResult{value, ops}
Mat-->>TC: Some(MaterializeResult)
TC-->>User: MaterializeResult
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 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 #458 +/- ##
=======================================
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: 2
🧹 Nitpick comments (1)
crates/trunk-ir/src/arena/walk.rs (1)
27-31: Consider removing per-nodeto_vec()snapshots in read-only walkers.These clones allocate at every recursion step. Iterating directly over slices can reduce traversal overhead for large modules.
Allocation-light traversal sketch
- let blocks: Vec<BlockRef> = ctx.region(region).blocks.to_vec(); - for block in blocks { + for &block in ctx.region(region).blocks.iter() { walk_block(ctx, block, f)?; } - let ops: Vec<OpRef> = ctx.block(block).ops.to_vec(); - for op in ops { + for &op in ctx.block(block).ops.iter() { walk_op(ctx, op, f)?; } - let regions: Vec<RegionRef> = ctx.op(op).regions.to_vec(); - for region in regions { + for ®ion in ctx.op(op).regions.iter() { walk_region(ctx, region, f)?; }Also applies to: 40-44, 58-60
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/trunk-ir/src/arena/walk.rs` around lines 27 - 31, The code currently takes per-node snapshots via ctx.region(region).blocks.to_vec() which causes needless allocations; change the loops to iterate directly over the slice returned by ctx.region(region).blocks (e.g., for block in ctx.region(region).blocks.iter() or for &block in ...) and call walk_block(ctx, block, f) without cloning the Vec. Replace the three occurrences (the to_vec() usage at the shown block and the similar sites around lines 40-44 and 58-60) so the walker reads the slice directly to avoid per-recursion allocations while preserving the same read-only semantics.
🤖 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/rewrite/mod.rs`:
- Around line 25-48: Make the ArenaModule wrapper safe by hiding its inner OpRef
and avoiding unchecked region indexing: change the struct from pub struct
ArenaModule(pub OpRef) to pub struct ArenaModule(OpRef) (make the field
private), update new to remain the only public constructor (keep returning
Option<Self>), change the methods signatures to take &self (pub fn op(&self) ->
OpRef) and pub fn body(&self, ctx: &IrContext) ->
Option<super::refs::RegionRef>, and implement body to check
ctx.op(self.0).regions.get(0).cloned() (or return Result with an explanatory
error) instead of unconditionally indexing; this prevents external construction
and panics when the module has no regions.
In `@crates/trunk-ir/src/arena/rewrite/type_converter.rs`:
- Around line 87-90: is_empty() currently returns true when conversions is empty
even if a materializer is present; update the method (pub fn is_empty(&self) ->
bool) to consider the materializer as well by returning true only when
conversions.is_empty() AND self.materializer.is_none(), so any set materializer
makes is_empty() return false; locate the method using the symbol is_empty and
adjust the boolean expression to include checking self.materializer.
---
Nitpick comments:
In `@crates/trunk-ir/src/arena/walk.rs`:
- Around line 27-31: The code currently takes per-node snapshots via
ctx.region(region).blocks.to_vec() which causes needless allocations; change the
loops to iterate directly over the slice returned by ctx.region(region).blocks
(e.g., for block in ctx.region(region).blocks.iter() or for &block in ...) and
call walk_block(ctx, block, f) without cloning the Vec. Replace the three
occurrences (the to_vec() usage at the shown block and the similar sites around
lines 40-44 and 58-60) so the walker reads the slice directly to avoid
per-recursion allocations while preserving the same read-only semantics.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
crates/trunk-ir/src/arena/mod.rscrates/trunk-ir/src/arena/rewrite/applicator.rscrates/trunk-ir/src/arena/rewrite/conversion_target.rscrates/trunk-ir/src/arena/rewrite/mod.rscrates/trunk-ir/src/arena/rewrite/pattern.rscrates/trunk-ir/src/arena/rewrite/rewriter.rscrates/trunk-ir/src/arena/rewrite/type_converter.rscrates/trunk-ir/src/arena/walk.rscrates/trunk-ir/src/op_interface.rs
- Make ArenaModule's inner OpRef field private; replace direct construction with ArenaModule::new().expect() in tests - Change body() to return Option<RegionRef> instead of panicking when the module op has no regions; update all callers in applicator.rs - Fix ArenaTypeConverter::is_empty() to also consider materializer presence - Remove unnecessary to_vec() allocations in walk_region, walk_block, and walk_op (ctx is &IrContext, so no mutation possible during iteration) - Add 9 unit tests across mod.rs and type_converter.rs Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…ction sites
Add `TypeDataBuilder` with a fluent API mirroring `OperationDataBuilder`.
Defaults to empty params and attrs (the common case), with `.param()`,
`.params()`, and `.attr()` methods for incremental construction.
Migrate all 16 `TypeData { ... }` struct literal sites across 8 files to
use the builder. Clean up now-unused `BTreeMap` and `smallvec` imports in
test modules where the builder removes the need for them.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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 Milestone 2 of the arena IR porting effort: walk utilities and a full pattern-based rewrite framework for the arena IR, analogous to the existing Salsa-based infrastructure but without lifetimes and without functional rebuilds.
Closes #437
What was added
arena/walk.rs:walk_region,walk_block,walk_op,walk_typedfor recursive traversal. UsesControlFlow<B, WalkAction>withAdvance/Skipvariants to give callers control over descent.arena/rewrite/pattern.rs:ArenaRewritePatterntrait — the arena counterpart ofRewritePattern.arena/rewrite/rewriter.rs:PatternRewriterwith RAUW-based mutation methods:replace_op,insert_op,erase_op,add_module_op. Operand access viarewriter.operand(i)mirrors the Salsa convention.arena/rewrite/type_converter.rs:ArenaTypeConverterfor registering per-dialect type-lowering callbacks used during dialect conversion.arena/rewrite/conversion_target.rs:ArenaConversionTargetwith legality checks (legal / illegal / dynamically-legal operations) and averifypass that collectsIllegalOperrors.arena/rewrite/applicator.rs:PatternApplicatorwith visitor-based fixpoint iteration. Snapshots block op lists before each pass; skips deleted ops by checkingparent_blockvalidity. ExposesApplyResultwith iteration count, total change count, and fixpoint status.arena/rewrite/mod.rs:ArenaModulewrapper (thinOpRefnewtype forcore.module), verifies dialect/name on construction, exposesbody,ops,first_block, andnamehelpers.op_interface.rs:is_pure_arena,is_removable_arena,is_isolated_arena— arena-native helpers that delegate to the existing inventory-based registry checks, enabling DCE and isolation queries without Salsa.Why
The Salsa-based rewrite infrastructure recreates entire IR nodes on each mutation. The arena IR uses direct in-place mutation + RAUW (Replace All Uses With), which removes the need for value maps and produces ~5x less boilerplate per pass. This milestone makes it possible to begin porting actual lowering passes (boxing, closures, evidence) to the arena backend.
Key design decisions
OpRef/TypeRef/BlockRef/RegionRefentity indices rather than borrowed Salsa-tracked references. TheIrContextis always passed explicitly.replace_oprewires all uses of the old results to the new op's results in-place, without building a side-table value map.PatternApplicatoriterates until no pattern fires ormax_iterationsis reached, consistent with MLIR'sapplyPatternsAndFoldGreedilyapproach.WalkAction::Skip— callers can suppress descent into nested regions (e.g., to avoid crossing isolation boundaries).Test plan
cargo nextest run -p trunk-ir)Skipsemantics,PatternRewriterRAUW correctness,ArenaConversionTargetlegality checks,PatternApplicatorfixpoint iteration and max-iterations guard,ArenaModulewrapper construction and accessors🤖 Generated with Claude Code
Summary by CodeRabbit