From 2b5f0c9430a9be1dc16bef10819f656ad2e03fab Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Thu, 5 Nov 2026 09:29:00 -0500 Subject: [PATCH 01/14] feat(ir): replace only selected uses in Rewriter Add replace_uses_if, replace_all_uses_except_stmt and replace_all_uses_except_stmts, which redirect only the use sites a predicate accepts. ReplacedUses is recorded only once the old value has no uses left. Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/rewriter.rs | 122 +++++++- crates/kirin-ir/tests/rewriter.rs | 334 +++++++++++++++++++++- crates/kirin-ir/tests/rewriter_ungraph.rs | 36 +++ docs/design/rewrite-engine.md | 2 + 4 files changed, 480 insertions(+), 14 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index b50a57c22..db2e34bc8 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -23,7 +23,9 @@ //! applied are available for diagnosis only. //! //! Scope is operand/yield rewriting (including bulk result redirection through -//! [`Rewriter::replace_results`]), block-body statement surgery +//! [`Rewriter::replace_results`], and selective redirection through +//! [`Rewriter::replace_uses_if`] and +//! [`Rewriter::replace_all_uses_except_stmt`]), block-body statement surgery //! ([`Rewriter::erase_statement`], [`Rewriter::insert_before`], //! [`Rewriter::insert_after`], [`Rewriter::replace_statement`], //! [`Rewriter::insert_terminator`]), `CFG` @@ -601,6 +603,75 @@ impl<'a, L: Dialect> Rewriter<'a, L> { Ok(self.redirect_uses(&[(old, new)])) } + /// Replace the uses of `old` that `select` accepts with `new`, returning + /// the number of use slots rewritten. Replacing a value with itself is a + /// no-op returning `Ok(0)`. + /// + /// `select` is called with the stage and each use site of `old` (a + /// statement operand or a `DiGraph` yield slot), all before anything is + /// written, so it sees the stage unchanged and can look up the using + /// statement, e.g. its parent block. The sites it rejects stay uses of + /// `old`. Otherwise this is [`Rewriter::replace_all_uses`] restricted to + /// the selected sites. + /// + /// Rejects, leaving the stage untouched, if `old` or `new` is unknown. + /// Records the events of [`Rewriter::replace_all_uses`] for the rewritten + /// slots, except that [`MutationEvent::ReplacedUses`] is recorded only if + /// `select` accepted every use, since it promises that every use of `old` + /// moved. + pub fn replace_uses_if( + &mut self, + old: impl Into, + new: impl Into, + select: impl FnMut(&StageInfo, Use) -> bool, + ) -> Result { + let old: SSAValue = old.into(); + let new: SSAValue = new.into(); + if old.get_info(self.stage).is_none() { + return Err(RewriteError::UnknownValue(old)); + } + if new.get_info(self.stage).is_none() { + return Err(RewriteError::UnknownValue(new)); + } + if old == new { + return Ok(0); + } + Ok(self.redirect_uses_where(&[(old, new)], select)) + } + + /// Replace every use of `old` with `new` except `except`'s own operands, + /// returning the number of use slots rewritten. + /// + /// Only `except`'s operands are kept: uses by statements nested in its + /// bodies are replaced (for those, use [`Rewriter::replace_uses_if`]). + /// Rejects, leaving the stage untouched, if `old`, `new` or `except` is + /// unknown. Records the events of [`Rewriter::replace_uses_if`]. + pub fn replace_all_uses_except_stmt( + &mut self, + old: impl Into, + new: impl Into, + except: Statement, + ) -> Result { + self.replace_all_uses_except_stmts(old, new, &[except]) + } + + /// [`Rewriter::replace_all_uses_except_stmt`] keeping the operands of every + /// statement in `except`. + pub fn replace_all_uses_except_stmts( + &mut self, + old: impl Into, + new: impl Into, + except: &[Statement], + ) -> Result { + for &stmt in except { + self.live_statement(stmt)?; + } + self.replace_uses_if(old, new, |_, site| match site { + Use::StatementOperand { stmt, .. } => !except.contains(&stmt), + Use::DiGraphYield { .. } => true, + }) + } + /// Redirect **every** result of `stmt` at once, one replacement per result /// in result order, exact arity required. /// @@ -667,14 +738,28 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// *to* it. A sequential loop would instead let a later pair carry away the /// sites an earlier one just moved. pub(super) fn redirect_uses(&mut self, pairs: &[(SSAValue, SSAValue)]) -> usize { - // Snapshot the use sites so we can mutate the arenas while walking them - // (the borrows of the `old` infos end here). + self.redirect_uses_where(pairs, |_, _| true) + } + + /// [`Rewriter::redirect_uses`], rewriting only the use sites `select` + /// accepts. `select` sees every site before anything is written. The + /// sites it rejects stay uses of `old`. + fn redirect_uses_where( + &mut self, + pairs: &[(SSAValue, SSAValue)], + mut select: impl FnMut(&StageInfo, Use) -> bool, + ) -> usize { + // Snapshot the selected use sites so we can mutate the arenas while + // walking them (the borrows of the `old` infos end here). let sites: Vec> = pairs .iter() .map(|(old, _)| { old.get_info(self.stage) .map(|info| info.uses().to_vec()) .unwrap_or_default() + .into_iter() + .filter(|&site| select(self.stage, site)) + .collect() }) .collect(); @@ -754,13 +839,24 @@ impl<'a, L: Dialect> Rewriter<'a, L> { // Transfer the use records in two phases: drop every moved list before // installing any of them, so a value appearing on both sides of the - // substitution keeps only the sites redirected to it. - for (&(old, _), moved) in pairs.iter().zip(&moved) { - if !moved.is_empty() - && let Some(info) = old.get_info_mut(self.stage) - { - info.uses_mut().clear(); - } + // substitution keeps only the sites redirected to it. A selected site + // that no longer applied was stale, and is dropped too. Whether `old` + // is left with no uses is read here, before any replacement gains its + // sites: a value can be both. + let mut emptied: Vec = Vec::with_capacity(pairs.len()); + for (&(old, _), sites) in pairs.iter().zip(&sites) { + let Some(info) = old.get_info_mut(self.stage) else { + emptied.push(false); + continue; + }; + let uses = info.uses_mut(); + if sites.len() == uses.len() { + uses.clear(); + } else if !sites.is_empty() { + let sites: HashSet = sites.iter().copied().collect(); + uses.retain(|site| !sites.contains(site)); + } + emptied.push(uses.is_empty()); } for (&(_, new), moved) in pairs.iter().zip(&moved) { if !moved.is_empty() @@ -791,8 +887,10 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.events .push(MutationEvent::ChangedDiGraphYields { graph }); } - for (&(old, new), moved) in pairs.iter().zip(&moved) { - if !moved.is_empty() { + // `ReplacedUses` promises every use of `old` moved; a partial + // redirection is covered by the events above only. + for ((&(old, new), moved), &emptied) in pairs.iter().zip(&moved).zip(&emptied) { + if !moved.is_empty() && emptied { self.events.push(MutationEvent::ReplacedUses { old, new }); } } diff --git a/crates/kirin-ir/tests/rewriter.rs b/crates/kirin-ir/tests/rewriter.rs index 1b351c0cf..9e27e8b4d 100644 --- a/crates/kirin-ir/tests/rewriter.rs +++ b/crates/kirin-ir/tests/rewriter.rs @@ -1003,10 +1003,22 @@ fn replace_results_substitutes_simultaneously() { let mut stage = stage.finalize().unwrap(); // a -> b and b -> x at once. - { + let events = { let mut rw = Rewriter::new(&mut stage); rw.replace_results(split, &[b, x]).unwrap(); - } + rw.drain_events() + }; + // Every use of `a` and of `b` moved, so both get a `ReplacedUses`, even + // though `b` gained the site that read `a`. + assert_eq!( + events, + vec![ + MutationEvent::ChangedOperands { stmt: reads_a }, + MutationEvent::ChangedOperands { stmt: reads_b }, + MutationEvent::ReplacedUses { old: a, new: b }, + MutationEvent::ReplacedUses { old: b, new: x }, + ] + ); // The site that read `a` reads `b`; the site that read `b` reads `x`. match reads_a.definition(&stage) { @@ -1086,3 +1098,321 @@ fn replace_results_rejects_bad_arity_and_dead_ids() { assert_eq!(events, vec![MutationEvent::ErasedStatement { stmt: nop }]); assert_eq!((uses_set(&stage, a), uses_set(&stage, b)), before); } + +/// `%a` read by `u0` in the outer block and by `u1` / `u2` in the two blocks +/// `owner` owns; `outer`'s argument is the replacement. +struct NestedUses { + stage: StageInfo, + a: SSAValue, + new: SSAValue, + u0: Statement, + u1: Statement, + u2: Statement, + b1: Block, + owner: Statement, +} + +fn nested_uses() -> NestedUses { + let mut stage = new_stage(); + let s_src = stage.statement().definition(BuilderDialect::Nop).new(); + let a = stage + .ssa() + .ty(TestType::I32) + .kind(BuilderSSAKind::Result(s_src, 0)) + .new(); + let u0 = stage.statement().definition(BuilderDialect::Use(a)).new(); + let u1 = stage.statement().definition(BuilderDialect::Use(a)).new(); + let u2 = stage.statement().definition(BuilderDialect::Use(a)).new(); + let b1 = stage.block().stmt(u1).new(); + let b2 = stage.block().stmt(u2).new(); + let owner = stage + .statement() + .definition(BuilderDialect::OwnBlocks(b1, b2)) + .new(); + let outer = stage + .block() + .argument(TestType::I32) + .stmt(s_src) + .stmt(u0) + .stmt(owner) + .new(); + let stage = stage.finalize().unwrap(); + let new = SSAValue::from(outer.expect_info(&stage).arguments()[0]); + NestedUses { + stage, + a, + new, + u0, + u1, + u2, + b1, + owner, + } +} + +fn used(stage: &StageInfo, stmt: Statement) -> SSAValue { + match stmt.definition(stage) { + BuilderDialect::Use(value) => *value, + other => panic!("expected Use, got {other:?}"), + } +} + +#[test] +fn replace_uses_if_rewrites_only_the_selected_sites() { + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + .. + } = nested_uses(); + + let events = { + let mut rewriter = Rewriter::new(&mut stage); + let replaced = rewriter + .replace_uses_if(a, new, |_, site| site == so(u1, 0)) + .unwrap(); + assert_eq!(replaced, 1); + rewriter.drain_events() + }; + + assert_eq!( + (used(&stage, u0), used(&stage, u1), used(&stage, u2)), + (a, new, a) + ); + // The rejected sites stay uses of `a`. + assert_eq!(uses_set(&stage, a), HashSet::from([so(u0, 0), so(u2, 0)])); + assert_eq!(uses_set(&stage, new), HashSet::from([so(u1, 0)])); + // Not every use moved, so no `ReplacedUses`. + assert_eq!(events, vec![MutationEvent::ChangedOperands { stmt: u1 }]); + verify_derived(&stage).unwrap(); +} + +#[test] +fn replace_uses_if_can_look_up_the_using_statement() { + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + b1, + .. + } = nested_uses(); + + let mut rewriter = Rewriter::new(&mut stage); + let replaced = rewriter + .replace_uses_if(a, new, |stage, site| match site { + Use::StatementOperand { stmt, .. } => { + *stmt.parent(stage) == Some(StatementParent::Block(b1)) + } + Use::DiGraphYield { .. } => false, + }) + .unwrap(); + assert_eq!(replaced, 1); + drop(rewriter); + + assert_eq!( + (used(&stage, u0), used(&stage, u1), used(&stage, u2)), + (a, new, a) + ); + verify_derived(&stage).unwrap(); +} + +#[test] +fn replace_uses_if_selecting_every_use_records_replaced_uses() { + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + .. + } = nested_uses(); + + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!(rewriter.replace_uses_if(a, new, |_, _| true), Ok(3)); + let events = rewriter.drain_events(); + drop(rewriter); + + assert_eq!( + events, + vec![ + MutationEvent::ChangedOperands { stmt: u0 }, + MutationEvent::ChangedOperands { stmt: u1 }, + MutationEvent::ChangedOperands { stmt: u2 }, + MutationEvent::ReplacedUses { old: a, new }, + ] + ); + assert!(uses_set(&stage, a).is_empty()); + verify_derived(&stage).unwrap(); +} + +#[test] +fn replace_all_uses_except_stmt_keeps_only_its_own_operands() { + // Excepting `u0` keeps its operand. + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + .. + } = nested_uses(); + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!(rewriter.replace_all_uses_except_stmt(a, new, u0), Ok(2)); + drop(rewriter); + assert_eq!( + (used(&stage, u0), used(&stage, u1), used(&stage, u2)), + (a, new, new) + ); + verify_derived(&stage).unwrap(); + + // Excepting `owner` keeps nothing: the uses nested in its blocks are not + // its own operands. + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + owner, + .. + } = nested_uses(); + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!(rewriter.replace_all_uses_except_stmt(a, new, owner), Ok(3)); + drop(rewriter); + assert_eq!( + (used(&stage, u0), used(&stage, u1), used(&stage, u2)), + (new, new, new) + ); + verify_derived(&stage).unwrap(); +} + +#[test] +fn replace_all_uses_except_stmts_keeps_the_operands_of_each() { + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + .. + } = nested_uses(); + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!( + rewriter.replace_all_uses_except_stmts(a, new, &[u0, u1]), + Ok(1) + ); + drop(rewriter); + assert_eq!( + (used(&stage, u0), used(&stage, u1), used(&stage, u2)), + (a, a, new) + ); + assert_eq!(uses_set(&stage, a), HashSet::from([so(u0, 0), so(u1, 0)])); + verify_derived(&stage).unwrap(); +} + +#[test] +fn replace_uses_if_rejects_unknown_ids_leaving_the_stage_untouched() { + let NestedUses { + mut stage, + a, + new, + u0, + u1, + u2, + .. + } = nested_uses(); + let dead_value = stage.ssa_arena().next_id(); + let dead_stmt = stage.statement_arena().next_id(); + + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!( + rewriter.replace_uses_if(dead_value, new, |_, _| true), + Err(RewriteError::UnknownValue(dead_value)) + ); + assert_eq!( + rewriter.replace_uses_if(a, dead_value, |_, _| true), + Err(RewriteError::UnknownValue(dead_value)) + ); + assert_eq!( + rewriter.replace_all_uses_except_stmt(a, new, dead_stmt), + Err(RewriteError::UnknownStatement(dead_stmt)) + ); + assert_eq!( + rewriter.replace_all_uses_except_stmts(a, new, &[u0, dead_stmt]), + Err(RewriteError::UnknownStatement(dead_stmt)) + ); + assert!(rewriter.events().is_empty()); + drop(rewriter); + + assert_eq!( + (used(&stage, u0), used(&stage, u1), used(&stage, u2)), + (a, a, a) + ); + verify_derived(&stage).unwrap(); +} + +#[test] +fn replace_uses_if_separates_a_digraph_operand_from_its_yield() { + let mut stage = new_stage(); + let s_src = stage.statement().definition(BuilderDialect::Nop).new(); + let a = stage + .ssa() + .ty(TestType::I32) + .kind(BuilderSSAKind::Result(s_src, 0)) + .new(); + let consumer = stage.statement().definition(BuilderDialect::Use(a)).new(); + let dg = stage + .digraph() + .node(s_src) + .node(consumer) + .yield_value(a) + .new(); + let block = stage.block().argument(TestType::I32).new(); + let mut stage = stage.finalize().unwrap(); + let b = SSAValue::from(block.expect_info(&stage).arguments()[0]); + + // Only the yield: the operand keeps `a`, and with it the edge. + let events = { + let mut rewriter = Rewriter::new(&mut stage); + let is_yield = + |_: &StageInfo, site| matches!(site, Use::DiGraphYield { .. }); + assert_eq!(rewriter.replace_uses_if(a, b, is_yield), Ok(1)); + rewriter.drain_events() + }; + assert_eq!( + events, + vec![MutationEvent::ChangedDiGraphYields { graph: dg }] + ); + assert_eq!(dg.expect_info(&stage).yields()[0], b); + assert_eq!(used(&stage, consumer), a); + assert_eq!(dg.expect_info(&stage).graph().edge_count(), 1); + assert_eq!(uses_set(&stage, a), HashSet::from([so(consumer, 0)])); + verify_derived(&stage).unwrap(); + + // Then the operand: the edge goes, since `b` is not a member result. + let events = { + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!(rewriter.replace_uses_if(a, b, |_, _| true), Ok(1)); + rewriter.drain_events() + }; + assert_eq!( + events, + vec![ + MutationEvent::ChangedOperands { stmt: consumer }, + MutationEvent::ChangedDiGraphTopology { graph: dg }, + MutationEvent::ReplacedUses { old: a, new: b }, + ] + ); + assert_eq!(dg.expect_info(&stage).graph().edge_count(), 0); + verify_derived(&stage).unwrap(); +} diff --git a/crates/kirin-ir/tests/rewriter_ungraph.rs b/crates/kirin-ir/tests/rewriter_ungraph.rs index c98e9beb5..1cb2c1eba 100644 --- a/crates/kirin-ir/tests/rewriter_ungraph.rs +++ b/crates/kirin-ir/tests/rewriter_ungraph.rs @@ -814,3 +814,39 @@ fn replace_statement_on_ungraph_member_rejects_role_changes() { assert_topology(&stage, ug); verify_derived(&stage).unwrap(); } + +#[test] +fn replace_uses_if_moves_only_the_selected_reader_to_the_new_wire() { + let Fixture { + mut stage, + ug, + w, + v, + .. + } = fixture(); + let mut rw = Rewriter::new(&mut stage); + let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); + let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); + let c = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); + assert_eq!(edges(rw.stage(), ug), HashMap::from([(w, pair(a, b))])); + rw.drain_events(); + + // Only `a` moves: `w` keeps one reader and loses its edge, and `v` gains + // one, joining `c`. + let only_a = + |_: &StageInfo, site| site == Use::StatementOperand { stmt: a, index: 0 }; + assert_eq!(rw.replace_uses_if(w, v, only_a), Ok(1)); + assert_eq!(edges(rw.stage(), ug), HashMap::from([(v, pair(a, c))])); + assert_eq!( + rw.drain_events(), + vec![ + MutationEvent::ChangedOperands { stmt: a }, + MutationEvent::ChangedUnGraphTopology { graph: ug }, + ] + ); + drop(rw); + + assert_eq!(*b.definition(&stage), BuilderDialect::Use(w)); + assert_topology(&stage, ug); + verify_derived(&stage).unwrap(); +} diff --git a/docs/design/rewrite-engine.md b/docs/design/rewrite-engine.md index 892445a34..f7d36487c 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -1792,6 +1792,8 @@ stage remains valid. |---|---|---|---| | `replace_operand` | Rewrite one operand slot; move its use record `old → new`; for a `DiGraph` member, move the slot's edge; for an `UnGraph` node, resync the wires it leaves and joins. | `ChangedOperands` (none if slot unchanged), then `ChangedDiGraphTopology` or `ChangedUnGraphTopology` if an edge changed | operand-index and live-value checked; an `UnGraph` wire may exceed two readers (reported by `verify_derived`) | | `replace_all_uses` | Rewrite every operand and `DiGraph` yield slot reading `old` to `new`; transfer use records and move the edges of rewritten member slots (resyncing affected `UnGraph` wires). | one `ChangedOperands` per affected statement, one `ChangedDiGraphTopology` / `ChangedUnGraphTopology` per graph whose edges changed, one `ChangedDiGraphYields` per affected graph, then one `ReplacedUses` | returns count of slots rewritten; self-replace is `Ok(0)` | +| `replace_uses_if` | `replace_all_uses` restricted to the use sites a predicate accepts; the predicate gets the stage and each site, all before any write, and rejected sites stay uses of `old`. | as `replace_all_uses` for the rewritten slots; `ReplacedUses` only if every use was selected | returns count of slots rewritten; self-replace is `Ok(0)`; `old` and `new` must be live | +| `replace_all_uses_except_stmt` / `_stmts` | `replace_uses_if` keeping the operands of the given statement(s); uses nested in their bodies are still replaced. | as `replace_uses_if` | the excepted statements must be live | | `erase_statement` | Tombstone the statement, unlink it from its block list (or remove its graph node and the edges touching it, or drop an `UnGraph` edge statement from the member map), drop its operand uses, tombstone its (unused) result values. Statement-owned blocks, `Cfg`s, `DiGraph`s, and `UnGraph`s nested in it are tombstoned with it. Erasing a terminator clears the block's terminator and removes its outgoing edges, leaving the block unterminated. | `ErasedStatement` (covers the nested bodies; no CFG or graph events), then `ChangedTerminator` for a terminator, `ChangedDiGraphTopology` for a `DiGraph` member, or `ChangedUnGraphTopology` for an `UnGraph` member | parented by a block or graph; no live result uses (a `DiGraph` yield is a use, and so is a node reading an edge statement's wire) | | `new_result` | Allocate a typed result value for a statement not inserted yet; its arena slot stays empty (not live) until an insert binds it. | none | a value its `Rewriter` never binds is tombstoned when the `Rewriter` is dropped (`run_pass` drops it before verifying) | | `insert_before` / `insert_after` | Allocate a statement, splice it adjacent to `anchor`, bind its results, add operand uses; return the new id. | `InsertedStatement` | anchor is block-parented; only `insert_before` may anchor on the terminator, which appends to the end of the chain; `definition` is a non-terminator with no successors; operands must be live; results must be distinct, unbound `new_result` values | From f949f89cb7432443d1deaa63678e5abbb0dd382c Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Thu, 5 Nov 2026 11:07:55 -0500 Subject: [PATCH 02/14] refactor(ir): replace Rewriter::drain_events with into_events The event log now only grows while a Rewriter lives, so positions in events() stay valid. Tests track positions with an EventCursor helper. Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/pass.rs | 8 +- crates/kirin-ir/src/rewrite/rewriter.rs | 6 +- crates/kirin-ir/tests/common.rs | 21 ++++++ crates/kirin-ir/tests/quarantine.rs | 2 +- .../tests/replace_statement_bodies.rs | 6 +- crates/kirin-ir/tests/rewriter.rs | 41 ++++++----- crates/kirin-ir/tests/rewriter_cfg.rs | 27 ++++--- crates/kirin-ir/tests/rewriter_digraph.rs | 26 ++++--- crates/kirin-ir/tests/rewriter_ungraph.rs | 73 +++++++++++-------- docs/design/rewrite-engine.md | 2 +- 10 files changed, 126 insertions(+), 86 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/pass.rs b/crates/kirin-ir/src/rewrite/pass.rs index 2d258a961..bd76cb0c2 100644 --- a/crates/kirin-ir/src/rewrite/pass.rs +++ b/crates/kirin-ir/src/rewrite/pass.rs @@ -38,7 +38,9 @@ where let mut rewriter = Rewriter::new(stage); let panic_or_result = catch_unwind(AssertUnwindSafe(|| pass(&mut rewriter))); - let events = rewriter.drain_events(); + // Taking the events drops the rewriter, tombstoning the results the pass + // never bound and ending the borrow of `stage` before it is verified. + let events = rewriter.into_events(); let pass_result = match panic_or_result { Ok(pass_result) => pass_result, @@ -51,10 +53,6 @@ where Err(error) => return (events, Err(QuarantineCause::Pass(Box::new(error)))), }; - // Tombstone the results the pass never bound, ending the borrow of - // `stage` before it is verified. - drop(rewriter); - // Verify the generated IR if let Err(error) = verify_derived(stage) { return (events, Err(QuarantineCause::Verify(error))); diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index db2e34bc8..b2e2a79cc 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -483,8 +483,10 @@ impl<'a, L: Dialect> Rewriter<'a, L> { &self.events } - /// Take all recorded events, leaving the log empty. - pub fn drain_events(&mut self) -> Vec { + /// End the rewrite and return every event it recorded, in application + /// order. The log only grows while the `Rewriter` lives, so a position in + /// [`events`](Self::events) stays valid until then. + pub fn into_events(mut self) -> Vec { std::mem::take(&mut self.events) } diff --git a/crates/kirin-ir/tests/common.rs b/crates/kirin-ir/tests/common.rs index 0bc0521bc..db2ed238f 100644 --- a/crates/kirin-ir/tests/common.rs +++ b/crates/kirin-ir/tests/common.rs @@ -396,3 +396,24 @@ pub fn make_compute( pub fn new_stage() -> BuilderStageInfo { BuilderStageInfo::default() } + +/// A position in a [`Rewriter`]'s event log, for checking each edit's events +/// in turn. +#[allow(dead_code)] +#[derive(Default)] +pub struct EventCursor(usize); + +#[allow(dead_code)] +impl EventCursor { + /// Move past every event recorded so far. + pub fn advance(&mut self, rw: &Rewriter<'_, L>) { + self.0 = rw.events().len(); + } + + /// The events recorded since the cursor's position, moving past them. + pub fn read_and_advance(&mut self, rw: &Rewriter<'_, L>) -> Vec { + let events = rw.events()[self.0..].to_vec(); + self.advance(rw); + events + } +} diff --git a/crates/kirin-ir/tests/quarantine.rs b/crates/kirin-ir/tests/quarantine.rs index ed09c8151..3086bafb9 100644 --- a/crates/kirin-ir/tests/quarantine.rs +++ b/crates/kirin-ir/tests/quarantine.rs @@ -38,7 +38,7 @@ fn fixture(cause: QuarantineCause) -> Fixture { let mut rewriter = Rewriter::new(&mut stage); rewriter.erase_statement(nop).expect("nop is erasable"); - let events = rewriter.drain_events(); + let events = rewriter.events().to_vec(); drop(rewriter); Fixture { diff --git a/crates/kirin-ir/tests/replace_statement_bodies.rs b/crates/kirin-ir/tests/replace_statement_bodies.rs index 32524aa85..550253d6f 100644 --- a/crates/kirin-ir/tests/replace_statement_bodies.rs +++ b/crates/kirin-ir/tests/replace_statement_bodies.rs @@ -77,7 +77,7 @@ fn swapping_the_bodies_keeps_them() { BuilderDialect::OwnBlocks(f.else_block, f.then_block), ) .unwrap(); - rewriter.drain_events() + rewriter.events().to_vec() }; assert_eq!( events, @@ -103,7 +103,7 @@ fn a_dropped_body_is_erased_and_a_gained_one_adopted() { rewriter .replace_statement(f.owner, BuilderDialect::OwnBlocks(f.then_block, f.spare)) .unwrap(); - rewriter.drain_events() + rewriter.events().to_vec() }; assert_eq!( events, @@ -132,7 +132,7 @@ fn assert_rejected(f: &mut Fixture, definition: BuilderDialect, error: RewriteEr let (result, events) = { let mut rewriter = Rewriter::new(&mut f.stage); let result = rewriter.replace_statement(f.owner, definition); - (result, rewriter.drain_events()) + (result, rewriter.events().to_vec()) }; assert_eq!(result, Err(error)); assert!(events.is_empty()); diff --git a/crates/kirin-ir/tests/rewriter.rs b/crates/kirin-ir/tests/rewriter.rs index 9e27e8b4d..868fb2833 100644 --- a/crates/kirin-ir/tests/rewriter.rs +++ b/crates/kirin-ir/tests/rewriter.rs @@ -5,7 +5,7 @@ mod common; use std::collections::HashSet; -use common::{BuilderDialect, TestType, make_split, make_wire, new_stage}; +use common::{BuilderDialect, EventCursor, TestType, make_split, make_wire, new_stage}; use kirin_ir::*; /// The def-use set of `value` (order-agnostic; `SSAInfo::uses` is a bag). @@ -57,10 +57,11 @@ fn replace_all_uses_rewrites_operands_and_records_events() { }; let mut rewriter = Rewriter::new(&mut stage); + let mut log = EventCursor::default(); let replaced = rewriter.replace_all_uses(real_x, real_y).unwrap(); assert_eq!(replaced, 3, "Add uses x twice, Use once"); assert_eq!( - rewriter.drain_events(), + log.read_and_advance(&rewriter), vec![ MutationEvent::ChangedOperands { stmt: add }, MutationEvent::ChangedOperands { stmt: use_stmt }, @@ -74,7 +75,7 @@ fn replace_all_uses_rewrites_operands_and_records_events() { // A second replacement finds nothing left to rewrite and records nothing. assert_eq!(rewriter.replace_all_uses(real_x, real_y).unwrap(), 0); - assert!(rewriter.events().is_empty()); + assert!(log.read_and_advance(&rewriter).is_empty()); drop(rewriter); match add.definition(&stage) { @@ -141,7 +142,7 @@ fn replace_operand_updates_single_slot_and_rejects_illegal_edits() { assert_eq!(rewriter.replace_operand(add, 1, real_y).unwrap(), real_y); assert_eq!( - rewriter.drain_events(), + rewriter.events().to_vec(), vec![MutationEvent::ChangedOperands { stmt: add }], "only the one real edit is recorded" ); @@ -244,7 +245,7 @@ fn replace_all_uses_maintains_def_use_index_including_yields() { let mut rewriter = Rewriter::new(&mut stage); let replaced = rewriter.replace_all_uses(a, real_b).unwrap(); assert_eq!(replaced, 2, "one operand slot + one yield slot"); - rewriter.drain_events() + rewriter.events().to_vec() }; // After: a has no uses; b absorbed both (operand + yield). @@ -311,7 +312,7 @@ fn erase_statement_unlinks_and_maintains_index() { let mut rw = Rewriter::new(&mut stage); rw.erase_statement(s1).unwrap(); assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![MutationEvent::ErasedStatement { stmt: s1 }] ); } @@ -350,7 +351,7 @@ fn erase_statement_tombstones_nested_bodies() { let mut rw = Rewriter::new(&mut stage); rw.erase_statement(owner).unwrap(); assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![MutationEvent::ErasedStatement { stmt: owner }] ); } @@ -398,7 +399,7 @@ fn erase_statement_erases_a_terminator_and_its_edges() { let events = { let mut rw = Rewriter::new(&mut stage); rw.erase_statement(jump).unwrap(); - rw.drain_events() + rw.events().to_vec() }; assert_eq!( events, @@ -478,7 +479,7 @@ fn insert_before_and_after_splice_and_maintain_index() { let mut rw = Rewriter::new(&mut stage); let n0 = rw.insert_before(s0, BuilderDialect::Use(real_y)).unwrap(); let n1 = rw.insert_after(s0, BuilderDialect::Use(real_x)).unwrap(); - (n0, n1, rw.drain_events()) + (n0, n1, rw.events().to_vec()) }; let order: Vec = block.statements(&stage).collect(); @@ -615,7 +616,7 @@ fn insert_binds_fresh_results() { // One event per successful insert; the rejections recorded none. assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![ MutationEvent::InsertedStatement { stmt: split }, MutationEvent::InsertedStatement { stmt: user }, @@ -644,7 +645,7 @@ fn insert_before_terminator_fills_a_terminator_only_block() { let (n0, events) = { let mut rw = Rewriter::new(&mut stage); let n0 = rw.insert_before(ret, BuilderDialect::Nop).unwrap(); - (n0, rw.drain_events()) + (n0, rw.events().to_vec()) }; assert_eq!(events, vec![MutationEvent::InsertedStatement { stmt: n0 }]); @@ -691,7 +692,7 @@ fn replace_statement_swaps_def_and_updates_uses() { rw.replace_statement(s0, BuilderDialect::Add(real_y, real_x)) .unwrap(); assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![MutationEvent::ReplacedStatement { stmt: s0 }] ); } @@ -880,7 +881,7 @@ fn replace_statement_maintains_predecessors_when_retargeting() { BuilderDialect::Branch(Successor::from_block(new_target)), ) .unwrap(); - rw.drain_events() + rw.events().to_vec() }; // The reverse control-flow index followed the edge. @@ -933,7 +934,7 @@ fn replace_results_redirects_every_result_and_unblocks_erase() { let events = { let mut rw = Rewriter::new(&mut stage); rw.replace_results(dup, &[orig_a, orig_b]).unwrap(); - rw.drain_events() + rw.events().to_vec() }; // Both results moved to the originals; the duplicate is now unread. @@ -1006,7 +1007,7 @@ fn replace_results_substitutes_simultaneously() { let events = { let mut rw = Rewriter::new(&mut stage); rw.replace_results(split, &[b, x]).unwrap(); - rw.drain_events() + rw.events().to_vec() }; // Every use of `a` and of `b` moved, so both get a `ReplacedUses`, even // though `b` gained the site that read `a`. @@ -1091,7 +1092,7 @@ fn replace_results_rejects_bad_arity_and_dead_ids() { // Replacing every result by itself is a no-op. assert_eq!(rw.replace_results(split, &[a, b]), Ok(())); - rw.drain_events() + rw.events().to_vec() }; // Only the erase recorded anything; no rejected call moved a use list. @@ -1175,7 +1176,7 @@ fn replace_uses_if_rewrites_only_the_selected_sites() { .replace_uses_if(a, new, |_, site| site == so(u1, 0)) .unwrap(); assert_eq!(replaced, 1); - rewriter.drain_events() + rewriter.events().to_vec() }; assert_eq!( @@ -1236,7 +1237,7 @@ fn replace_uses_if_selecting_every_use_records_replaced_uses() { let mut rewriter = Rewriter::new(&mut stage); assert_eq!(rewriter.replace_uses_if(a, new, |_, _| true), Ok(3)); - let events = rewriter.drain_events(); + let events = rewriter.events().to_vec(); drop(rewriter); assert_eq!( @@ -1387,7 +1388,7 @@ fn replace_uses_if_separates_a_digraph_operand_from_its_yield() { let is_yield = |_: &StageInfo, site| matches!(site, Use::DiGraphYield { .. }); assert_eq!(rewriter.replace_uses_if(a, b, is_yield), Ok(1)); - rewriter.drain_events() + rewriter.events().to_vec() }; assert_eq!( events, @@ -1403,7 +1404,7 @@ fn replace_uses_if_separates_a_digraph_operand_from_its_yield() { let events = { let mut rewriter = Rewriter::new(&mut stage); assert_eq!(rewriter.replace_uses_if(a, b, |_, _| true), Ok(1)); - rewriter.drain_events() + rewriter.events().to_vec() }; assert_eq!( events, diff --git a/crates/kirin-ir/tests/rewriter_cfg.rs b/crates/kirin-ir/tests/rewriter_cfg.rs index b4d9b9ac4..bcfc213f9 100644 --- a/crates/kirin-ir/tests/rewriter_cfg.rs +++ b/crates/kirin-ir/tests/rewriter_cfg.rs @@ -6,7 +6,7 @@ mod common; use std::collections::HashSet; -use common::{BuilderDialect, TestType, make_split, new_stage}; +use common::{BuilderDialect, EventCursor, TestType, make_split, new_stage}; use kirin_ir::*; /// The def-use set of `value` (order-agnostic; `SSAInfo::uses` is a bag). @@ -61,7 +61,7 @@ fn append_and_insert_blocks_link_into_cfg() { let after = rw .insert_block_after(entry, [TestType::I64, TestType::I32]) .unwrap(); - (tail, before, after, rw.drain_events()) + (tail, before, after, rw.events().to_vec()) }; assert_eq!( @@ -161,7 +161,7 @@ fn erase_block_tombstones_contents_and_maintains_mirrors() { let events = { let mut rw = Rewriter::new(&mut stage); rw.erase_block(dead).unwrap(); - rw.drain_events() + rw.events().to_vec() }; assert_eq!( events, @@ -279,6 +279,7 @@ fn erase_block_at_cfg_ends_updates_head_and_tail() { #[test] fn erase_block_rejects_values_used_outside_without_mutating() { + let mut log = EventCursor::default(); let mut stage = new_stage(); // middle(y) { split -> r0 } -> exit { use r0; use_arg } @@ -307,7 +308,7 @@ fn erase_block_rejects_values_used_outside_without_mutating() { // An argument read outside the block is rejected the same way. rw.replace_operand(use_r0, 0, y).unwrap(); - rw.drain_events(); + log.advance(&rw); assert_eq!( rw.erase_block(middle), Err(RewriteError::BlockArgumentInUse { @@ -317,7 +318,7 @@ fn erase_block_rejects_values_used_outside_without_mutating() { ); // Nothing was touched. - assert!(rw.events().is_empty()); + assert!(log.read_and_advance(&rw).is_empty()); assert!(middle.get_info(rw.stage()).is_some()); assert!(split.get_info(rw.stage()).is_some()); assert_eq!(preds(rw.stage(), exit), HashSet::from([middle])); @@ -415,7 +416,7 @@ fn split_block_moves_suffix_and_leaves_original_unterminated() { let (tail, events) = { let mut rw = Rewriter::new(&mut stage); let tail = rw.split_block(s1).unwrap(); - (tail, rw.drain_events()) + (tail, rw.events().to_vec()) }; assert_eq!( events, @@ -558,7 +559,7 @@ fn insert_terminator_wires_operands_and_edges() { BuilderDialect::Jump(Successor::from_block(target), vec![a]), ) .unwrap(); - (jump, rw.drain_events()) + (jump, rw.events().to_vec()) }; assert_eq!( events, @@ -694,7 +695,7 @@ fn merge_blocks_splices_successor_and_redirects_arguments() { let events = { let mut rw = Rewriter::new(&mut stage); rw.merge_blocks(entry, succ, &[real_a]).unwrap(); - rw.drain_events() + rw.events().to_vec() }; assert_eq!( events, @@ -975,12 +976,16 @@ fn retarget_successor_rewrites_every_matching_slot() { let mut stage = stage.finalize().unwrap(); let mut rw = Rewriter::new(&mut stage); + let mut log = EventCursor::default(); assert_eq!(rw.retarget_successor(cond, t, u).unwrap(), 2); - assert_eq!(rw.drain_events(), vec![MutationEvent::ChangedCfg { cfg }]); + assert_eq!( + log.read_and_advance(&rw), + vec![MutationEvent::ChangedCfg { cfg }] + ); // Nothing left to retarget: no-op, no event. assert_eq!(rw.retarget_successor(cond, t, u).unwrap(), 0); - assert!(rw.events().is_empty()); + assert!(log.read_and_advance(&rw).is_empty()); assert_eq!( cond.definition(rw.stage()), @@ -1093,7 +1098,7 @@ fn insert_block_arg_shifts_arguments_and_rewrites_edges() { }) .unwrap(); assert_eq!(asked, vec![(jump, 0)], "one value per incoming edge"); - (SSAValue::from(fresh), rw.drain_events()) + (SSAValue::from(fresh), rw.events().to_vec()) }; assert_eq!( events, diff --git a/crates/kirin-ir/tests/rewriter_digraph.rs b/crates/kirin-ir/tests/rewriter_digraph.rs index 7a89b75e3..e5bc6129d 100644 --- a/crates/kirin-ir/tests/rewriter_digraph.rs +++ b/crates/kirin-ir/tests/rewriter_digraph.rs @@ -6,7 +6,7 @@ mod common; use std::collections::{HashMap, HashSet}; -use common::{BuilderDialect, TestType, make_compute, make_split, new_stage}; +use common::{BuilderDialect, EventCursor, TestType, make_compute, make_split, new_stage}; use kirin_ir::*; use petgraph::visit::{EdgeRef, IntoEdgeReferences}; @@ -84,7 +84,7 @@ fn insert_node_adds_operand_edges_and_binds_results() { .unwrap(); // Two slots reading the same member result get two parallel edges. let n3 = rw.insert_node(dg, BuilderDialect::Add(rp, rp)).unwrap(); - let events = rw.drain_events(); + let events = rw.events().to_vec(); drop(rw); let topology = MutationEvent::ChangedDiGraphTopology { graph: dg }; @@ -206,7 +206,7 @@ fn erase_statement_removes_digraph_member_and_its_edges() { rw.erase_statement(u).unwrap(); // With `u` gone, `q`'s result is unused and `q` can go too. rw.erase_statement(q).unwrap(); - let events = rw.drain_events(); + let events = rw.events().to_vec(); drop(rw); let topology = MutationEvent::ChangedDiGraphTopology { graph: dg }; @@ -264,7 +264,7 @@ fn replace_yield_moves_the_yield_use() { // Writing the value already there records nothing. assert_eq!(rw.replace_yield(dg, 0, x), Ok(x)); assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![MutationEvent::ChangedDiGraphYields { graph: dg }] ); // No longer yielded, `p` can be erased. @@ -305,7 +305,7 @@ fn erase_statement_tombstones_owned_digraph() { let mut rw = Rewriter::new(&mut stage); rw.erase_statement(owner).unwrap(); assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![MutationEvent::ErasedStatement { stmt: owner }] ); } @@ -323,6 +323,7 @@ fn erase_statement_tombstones_owned_digraph() { #[test] fn replace_operand_on_member_retargets_its_edge() { + let mut log = EventCursor::default(); let mut stage = new_stage(); let (split, x, y) = make_split(&mut stage); @@ -340,22 +341,22 @@ fn replace_operand_on_member_retargets_its_edge() { // Member result → member result: the edge moves to the new producer. assert_eq!(rw.replace_operand(c, 0, rq), Ok(rp)); assert_eq!(edges(rw.stage(), dg), HashMap::from([((q, c, rq), 1)])); - assert_eq!(rw.drain_events(), vec![operands, topology]); + assert_eq!(log.read_and_advance(&rw), vec![operands, topology]); // Member result → outside value: the edge goes. rw.replace_operand(c, 0, x).unwrap(); assert!(edges(rw.stage(), dg).is_empty()); - assert_eq!(rw.drain_events(), vec![operands, topology]); + assert_eq!(log.read_and_advance(&rw), vec![operands, topology]); // Outside → outside: no edge either way, so no topology event. rw.replace_operand(c, 0, y).unwrap(); assert!(edges(rw.stage(), dg).is_empty()); - assert_eq!(rw.drain_events(), vec![operands]); + assert_eq!(log.read_and_advance(&rw), vec![operands]); // Outside → member result: an edge appears. rw.replace_operand(c, 0, rp).unwrap(); assert_eq!(edges(rw.stage(), dg), HashMap::from([((p, c, rp), 1)])); - assert_eq!(rw.drain_events(), vec![operands, topology]); + assert_eq!(log.read_and_advance(&rw), vec![operands, topology]); drop(rw); assert_eq!(uses_set(&stage, rp), HashSet::from([so(c, 0)])); @@ -389,7 +390,7 @@ fn replace_all_uses_retargets_every_member_edge() { let mut stage = stage.finalize().unwrap(); let mut rw = Rewriter::new(&mut stage); assert_eq!(rw.replace_all_uses(rp, rq), Ok(4)); - let events = rw.drain_events(); + let events = rw.events().to_vec(); // One topology and one yields event for the graph, however many slots. let expected = [ @@ -423,6 +424,7 @@ fn replace_all_uses_retargets_every_member_edge() { #[test] fn replace_statement_on_member_updates_operand_edges() { + let mut log = EventCursor::default(); let mut stage = new_stage(); // p → n → d, with q a second producer. @@ -451,7 +453,7 @@ fn replace_statement_on_member_updates_operand_edges() { rw.replace_statement(n, BuilderDialect::Compute(rq, placeholder)) .unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ReplacedStatement { stmt: n }, MutationEvent::ChangedDiGraphTopology { graph: dg }, @@ -466,7 +468,7 @@ fn replace_statement_on_member_updates_operand_edges() { rw.replace_statement(n, BuilderDialect::Compute(rq, placeholder)) .unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![MutationEvent::ReplacedStatement { stmt: n }] ); drop(rw); diff --git a/crates/kirin-ir/tests/rewriter_ungraph.rs b/crates/kirin-ir/tests/rewriter_ungraph.rs index 1cb2c1eba..2d99f750e 100644 --- a/crates/kirin-ir/tests/rewriter_ungraph.rs +++ b/crates/kirin-ir/tests/rewriter_ungraph.rs @@ -6,7 +6,7 @@ mod common; use std::collections::HashMap; -use common::{BuilderDialect, TestType, make_split, make_wire, new_stage}; +use common::{BuilderDialect, EventCursor, TestType, make_split, make_wire, new_stage}; use kirin_ir::*; use petgraph::visit::{EdgeRef, IntoEdgeReferences}; @@ -162,7 +162,7 @@ fn insert_ungraph_node_connects_the_readers_of_a_wire() { let d = rw .insert_ungraph_node(ug, BuilderDialect::Gate(boundary, capture)) .unwrap(); - let events = rw.drain_events(); + let events = rw.events().to_vec(); drop(rw); let topology = MutationEvent::ChangedUnGraphTopology { graph: ug }; @@ -236,7 +236,7 @@ fn insert_ungraph_edge_adds_a_wire_nodes_can_read() { let n2 = rw .insert_ungraph_node(ug, BuilderDialect::Use(r.into())) .unwrap(); - let events = rw.drain_events(); + let events = rw.events().to_vec(); drop(rw); let topology = MutationEvent::ChangedUnGraphTopology { graph: ug }; @@ -327,6 +327,7 @@ fn ungraph_inserts_reject_without_mutating() { #[test] fn erase_ungraph_node_drops_its_edges_and_frees_the_wire() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -338,12 +339,12 @@ fn erase_ungraph_node_drops_its_edges_and_frees_the_wire() { let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let c = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); - rw.drain_events(); + log.advance(&rw); // `a` is not the last node, so the last one moves into its slot. rw.erase_statement(a).unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ErasedStatement { stmt: a }, MutationEvent::ChangedUnGraphTopology { graph: ug }, @@ -367,6 +368,7 @@ fn erase_ungraph_node_drops_its_edges_and_frees_the_wire() { #[test] fn erase_ungraph_node_with_used_results_is_rejected() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -381,13 +383,13 @@ fn erase_ungraph_node_with_used_results_is_rejected() { let _reader = rw .insert_ungraph_node(ug, BuilderDialect::Use(r.into())) .unwrap(); - rw.drain_events(); + log.advance(&rw); assert_eq!( rw.erase_statement(n), Err(RewriteError::StatementResultsInUse(n)) ); - assert!(rw.events().is_empty()); + assert!(log.read_and_advance(&rw).is_empty()); drop(rw); assert!(n.get_info(&stage).is_some()); @@ -396,6 +398,7 @@ fn erase_ungraph_node_with_used_results_is_rejected() { #[test] fn erase_edge_statement_once_no_node_reads_its_wire() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -406,7 +409,7 @@ fn erase_edge_statement_once_no_node_reads_its_wire() { } = fixture(); let mut rw = Rewriter::new(&mut stage); let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); - rw.drain_events(); + log.advance(&rw); // A node still reads the wire. assert_eq!( @@ -414,10 +417,10 @@ fn erase_edge_statement_once_no_node_reads_its_wire() { Err(RewriteError::StatementResultsInUse(w_stmt)) ); rw.erase_statement(a).unwrap(); - rw.drain_events(); + log.advance(&rw); rw.erase_statement(w_stmt).unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ErasedStatement { stmt: w_stmt }, MutationEvent::ChangedUnGraphTopology { graph: ug }, @@ -472,7 +475,7 @@ fn erase_statement_tombstones_owned_ungraph() { let mut rw = Rewriter::new(&mut stage); rw.erase_statement(owner).unwrap(); assert_eq!( - rw.drain_events(), + rw.events().to_vec(), vec![MutationEvent::ErasedStatement { stmt: owner }] ); } @@ -492,6 +495,7 @@ fn erase_statement_tombstones_owned_ungraph() { #[test] fn replace_operand_swaps_endpoints_through_an_over_full_wire() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -505,7 +509,7 @@ fn replace_operand_swaps_endpoints_through_an_over_full_wire() { let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let c = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); let d = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); - rw.drain_events(); + log.advance(&rw); let topology = MutationEvent::ChangedUnGraphTopology { graph: ug }; // Moving b onto v first leaves v with three readers: no edge, and invalid @@ -522,14 +526,14 @@ fn replace_operand_swaps_endpoints_through_an_over_full_wire() { }] ); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![MutationEvent::ChangedOperands { stmt: b }, topology] ); // Moving d onto w completes it: w joins a — d, v joins c — b. rw.replace_operand(d, 0, w).unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![MutationEvent::ChangedOperands { stmt: d }, topology] ); drop(rw); @@ -544,6 +548,7 @@ fn replace_operand_swaps_endpoints_through_an_over_full_wire() { #[test] fn replace_operand_off_a_wire_drops_its_edge() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -554,13 +559,13 @@ fn replace_operand_off_a_wire_drops_its_edge() { let mut rw = Rewriter::new(&mut stage); let _a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); - rw.drain_events(); + log.advance(&rw); // A capture is not a wire: b leaves w, so w's edge goes. rw.replace_operand(b, 0, capture).unwrap(); assert!(edges(rw.stage(), ug).is_empty()); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ChangedOperands { stmt: b }, MutationEvent::ChangedUnGraphTopology { graph: ug }, @@ -572,10 +577,10 @@ fn replace_operand_off_a_wire_drops_its_edge() { let r = rw.new_result(TestType::I32); rw.insert_ungraph_node(ug, BuilderDialect::Compute(capture, r)) .unwrap(); - rw.drain_events(); + log.advance(&rw); rw.replace_operand(b, 0, r).unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![MutationEvent::ChangedOperands { stmt: b }] ); drop(rw); @@ -586,6 +591,7 @@ fn replace_operand_off_a_wire_drops_its_edge() { #[test] fn replace_all_uses_moves_every_reader_to_the_new_wire() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -596,13 +602,13 @@ fn replace_all_uses_moves_every_reader_to_the_new_wire() { let mut rw = Rewriter::new(&mut stage); let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let c = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); - rw.drain_events(); + log.advance(&rw); // One reader each: merging w into v gives v exactly two. assert_eq!(rw.replace_all_uses(w, v), Ok(1)); assert_eq!(edges(rw.stage(), ug), HashMap::from([(v, pair(a, c))])); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ChangedOperands { stmt: a }, MutationEvent::ChangedUnGraphTopology { graph: ug }, @@ -630,6 +636,7 @@ fn replace_all_uses_moves_every_reader_to_the_new_wire() { #[test] fn replace_results_redirects_an_edge_statements_readers() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -641,7 +648,7 @@ fn replace_results_redirects_an_edge_statements_readers() { let mut rw = Rewriter::new(&mut stage); let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); - rw.drain_events(); + log.advance(&rw); // w's readers move to v, so the edge now carries v, and w_stmt is free // to go. @@ -672,7 +679,7 @@ fn replace_all_uses_across_two_ungraphs() { let mut stage = stage.finalize().unwrap(); let mut rw = Rewriter::new(&mut stage); assert_eq!(rw.replace_all_uses(x, w1), Ok(2)); - let events = rw.drain_events(); + let events = rw.events().to_vec(); drop(rw); // In ug1, `b` joins `w1`; in ug2, `w1` is not a wire, so nothing changes. @@ -692,6 +699,7 @@ fn replace_all_uses_across_two_ungraphs() { #[test] fn replace_statement_on_ungraph_node_moves_its_wires() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -703,13 +711,13 @@ fn replace_statement_on_ungraph_node_moves_its_wires() { let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let c = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); - rw.drain_events(); + log.advance(&rw); // b now reads v instead of w: w loses its edge, v gains one. rw.replace_statement(b, BuilderDialect::Use(v)).unwrap(); assert_eq!(edges(rw.stage(), ug), HashMap::from([(v, pair(b, c))])); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ReplacedStatement { stmt: b }, MutationEvent::ChangedUnGraphTopology { graph: ug }, @@ -719,7 +727,7 @@ fn replace_statement_on_ungraph_node_moves_its_wires() { // The same operands again change no topology. rw.replace_statement(b, BuilderDialect::Use(v)).unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![MutationEvent::ReplacedStatement { stmt: b }] ); drop(rw); @@ -731,6 +739,7 @@ fn replace_statement_on_ungraph_node_moves_its_wires() { #[test] fn replace_statement_on_edge_statement_keeps_its_wire() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -742,7 +751,7 @@ fn replace_statement_on_edge_statement_keeps_its_wire() { let mut rw = Rewriter::new(&mut stage); let a = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); - rw.drain_events(); + log.advance(&rw); // The result slot is overwritten with the statement's own result, so `w` // stays the wire and keeps its readers. @@ -750,7 +759,7 @@ fn replace_statement_on_edge_statement_keeps_its_wire() { rw.replace_statement(w_stmt, BuilderDialect::Wire(placeholder)) .unwrap(); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![MutationEvent::ReplacedStatement { stmt: w_stmt }] ); drop(rw); @@ -770,6 +779,7 @@ fn replace_statement_on_edge_statement_keeps_its_wire() { #[test] fn replace_statement_on_ungraph_member_rejects_role_changes() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -782,7 +792,7 @@ fn replace_statement_on_ungraph_member_rejects_role_changes() { let node = rw .insert_ungraph_node(ug, BuilderDialect::Compute(capture, r)) .unwrap(); - rw.drain_events(); + log.advance(&rw); let placeholder = ResultValue::from(TestSSAValue(0)); // A node cannot become an edge statement, nor the reverse. @@ -804,7 +814,7 @@ fn replace_statement_on_ungraph_member_rejects_role_changes() { rw.replace_statement(node, BuilderDialect::Invoke(placeholder)), Err(RewriteError::TerminatorKindMismatch(node)) ); - assert!(rw.events().is_empty()); + assert!(log.read_and_advance(&rw).is_empty()); drop(rw); assert_eq!( @@ -817,6 +827,7 @@ fn replace_statement_on_ungraph_member_rejects_role_changes() { #[test] fn replace_uses_if_moves_only_the_selected_reader_to_the_new_wire() { + let mut log = EventCursor::default(); let Fixture { mut stage, ug, @@ -829,7 +840,7 @@ fn replace_uses_if_moves_only_the_selected_reader_to_the_new_wire() { let b = rw.insert_ungraph_node(ug, BuilderDialect::Use(w)).unwrap(); let c = rw.insert_ungraph_node(ug, BuilderDialect::Use(v)).unwrap(); assert_eq!(edges(rw.stage(), ug), HashMap::from([(w, pair(a, b))])); - rw.drain_events(); + log.advance(&rw); // Only `a` moves: `w` keeps one reader and loses its edge, and `v` gains // one, joining `c`. @@ -838,7 +849,7 @@ fn replace_uses_if_moves_only_the_selected_reader_to_the_new_wire() { assert_eq!(rw.replace_uses_if(w, v, only_a), Ok(1)); assert_eq!(edges(rw.stage(), ug), HashMap::from([(v, pair(a, c))])); assert_eq!( - rw.drain_events(), + log.read_and_advance(&rw), vec![ MutationEvent::ChangedOperands { stmt: a }, MutationEvent::ChangedUnGraphTopology { graph: ug }, diff --git a/docs/design/rewrite-engine.md b/docs/design/rewrite-engine.md index f7d36487c..258b2d5df 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -1777,7 +1777,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { // --- Immediate change/event log (current prototype API) --- pub fn events(&self) -> &[MutationEvent]; - pub fn drain_events(&mut self) -> Vec; + pub fn into_events(self) -> Vec; } ``` From 2a96592d330be1c865e3b4cd5b7a3dd07f5f37e9 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Thu, 5 Nov 2026 11:30:40 -0500 Subject: [PATCH 03/14] feat(ir): move statements with Rewriter::move_before/move_after Moves record MovedStatement; walk_mut rejects moving the visited statement in link-following orders. Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/mod.rs | 1 + crates/kirin-ir/src/rewrite/moves.rs | 79 ++++++ crates/kirin-ir/src/rewrite/rewriter.rs | 229 +++++++++------- .../src/rewrite/walk/breadth_first.rs | 6 +- .../kirin-ir/src/rewrite/walk/depth_first.rs | 4 +- crates/kirin-ir/src/rewrite/walk/mod.rs | 7 +- crates/kirin-ir/src/rewrite/walk/order.rs | 57 +++- crates/kirin-ir/src/rewrite/walk/tree.rs | 2 +- crates/kirin-ir/tests/rewriter.rs | 223 ++++++++++++++++ crates/kirin-ir/tests/walk.rs | 101 ++++++- .../src/driver/rewrite_driver.rs | 2 +- docs/design/rewrite-engine.md | 21 +- .../2026-11-04-remaining-ir-operations.md | 252 ++++++++++++++++++ 13 files changed, 863 insertions(+), 121 deletions(-) create mode 100644 crates/kirin-ir/src/rewrite/moves.rs create mode 100644 docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md diff --git a/crates/kirin-ir/src/rewrite/mod.rs b/crates/kirin-ir/src/rewrite/mod.rs index df443ca50..9973e5405 100644 --- a/crates/kirin-ir/src/rewrite/mod.rs +++ b/crates/kirin-ir/src/rewrite/mod.rs @@ -23,6 +23,7 @@ mod alloc; mod bodies; pub(crate) mod cfg; pub(crate) mod digraph; +mod moves; pub(crate) mod pass; #[cfg(test)] mod pass_tests; diff --git a/crates/kirin-ir/src/rewrite/moves.rs b/crates/kirin-ir/src/rewrite/moves.rs new file mode 100644 index 000000000..9eeda751b --- /dev/null +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -0,0 +1,79 @@ +use crate::arena::GetInfo; +use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter}; +use crate::rewrite::walk::tree; +use crate::{Dialect, Statement, StatementParent}; + +impl<'a, L: Dialect> Rewriter<'a, L> { + /// Move `stmt` to immediately before `anchor`, which may be in any block + /// body. + /// + /// The statement keeps its id, results, uses and nested bodies; only its + /// place changes. As with [`Rewriter::insert_before`], `anchor` may be a + /// block's terminator: `stmt` then goes to the end of that block's + /// statement chain. Moving a statement before itself does nothing. + /// + /// Not checked, as with [`Rewriter::insert_before`]: that every use of + /// `stmt`'s results still comes after it. + /// + /// Rejects, leaving the stage untouched, if `stmt` or `anchor` is unknown + /// or not in a block body ([`RewriteError::NotInBlockBody`]), if `stmt` is + /// a terminator ([`RewriteError::UnexpectedTerminator`]), or if `anchor` is + /// nested in `stmt`'s own bodies ([`RewriteError::MovedIntoItself`]). + /// + /// Records [`MutationEvent::MovedStatement`]. A walk visiting `stmt` + /// reports the move (see [`crate::walk`]). + pub fn move_before(&mut self, stmt: Statement, anchor: Statement) -> Result<(), RewriteError> { + self.move_relative(stmt, anchor, false) + } + + /// Move `stmt` to immediately after `anchor`, which may be in any block + /// body. + /// + /// As [`Rewriter::move_before`], except that `anchor` must not be a + /// terminator ([`RewriteError::AnchorIsTerminator`]): nothing follows one. + pub fn move_after(&mut self, stmt: Statement, anchor: Statement) -> Result<(), RewriteError> { + self.move_relative(stmt, anchor, true) + } + + fn move_relative( + &mut self, + stmt: Statement, + anchor: Statement, + after: bool, + ) -> Result<(), RewriteError> { + let block = { + let item = self.live_statement(stmt)?; + let Some(StatementParent::Block(block)) = item.parent else { + return Err(RewriteError::NotInBlockBody(stmt)); + }; + if item.definition.is_terminator() { + return Err(RewriteError::UnexpectedTerminator); + } + block + }; + // Checks the anchor; the place itself is read again once `stmt` is + // unlinked, since `stmt` may be one of its neighbours. + self.slot_at(anchor, after)?; + if stmt == anchor { + return Ok(()); + } + // Walk up the statements whose bodies enclose `anchor`. If `stmt` is + // one of them, the move would put `stmt` inside its own body. + let mut owner = tree::parent(self.stage, anchor); + while let Some(current) = owner { + if current == stmt { + return Err(RewriteError::MovedIntoItself { stmt, anchor }); + } + owner = tree::parent(self.stage, current); + } + + self.unlink_from_block(stmt, block); + let (block, prev, next) = self + .slot_at(anchor, after) + .expect("the anchor was checked above and moving `stmt` does not touch it"); + debug_assert!(stmt.get_info(self.stage).is_some()); + self.link_into_block(stmt, block, prev, next); + self.events.push(MutationEvent::MovedStatement { stmt }); + Ok(()) + } +} diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index b2e2a79cc..f51b517ec 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -28,7 +28,8 @@ //! [`Rewriter::replace_all_uses_except_stmt`]), block-body statement surgery //! ([`Rewriter::erase_statement`], [`Rewriter::insert_before`], //! [`Rewriter::insert_after`], [`Rewriter::replace_statement`], -//! [`Rewriter::insert_terminator`]), `CFG` +//! [`Rewriter::insert_terminator`], [`Rewriter::move_before`], +//! [`Rewriter::move_after`]), `CFG` //! block-list surgery in the sibling `cfg` module (append/insert/erase/split/ //! merge blocks, [`Rewriter::retarget_successor`], and block-argument insertion //! and erasure), `DiGraph` node surgery in the sibling `digraph` module @@ -139,6 +140,9 @@ pub enum MutationEvent { /// the bodies it dropped were tombstoned with everything nested in them, /// and the bodies it gained now have `owner` as their parent. ChangedBodyNesting { owner: Statement }, + /// `stmt` was moved to another place in a block body, keeping its id, + /// results, uses and nested bodies. + MovedStatement { stmt: Statement }, } /// An illegal mutation rejected by a [`Rewriter`]. @@ -162,11 +166,12 @@ pub enum RewriteError { /// Nothing can follow a block's terminator, so `insert_after` cannot anchor /// on one. (`insert_before` can: it appends to the end of the chain.) AnchorIsTerminator(Statement), - /// The inserted definition is a terminator where a non-terminator is - /// required: `insert_before`/`insert_after` do not splice terminators into - /// a block's statement list (use [`Rewriter::insert_terminator`] on a block - /// that has none), and a `DiGraph` has no terminators for `insert_node` to - /// add. The counterpart of [`RewriteError::ExpectedTerminator`]. + /// The inserted definition, or the moved statement, is a terminator where + /// a non-terminator is required: `insert_before`/`insert_after` and + /// `move_before`/`move_after` do not splice terminators into a block's + /// statement list (use [`Rewriter::insert_terminator`] on a block that has + /// none), and a `DiGraph` has no terminators for `insert_node` to add. The + /// counterpart of [`RewriteError::ExpectedTerminator`]. UnexpectedTerminator, /// The inserted or replacement definition is a non-terminator that carries /// successors. Only a block's terminator may transfer control, so such a @@ -226,6 +231,10 @@ pub enum RewriteError { /// A `replace_statement` definition names a body the statement itself is /// nested in, which would make the statement its own ancestor. BodyEnclosesOwner(Body), + /// `move_before`/`move_after` was given an anchor nested in the moved + /// statement's own bodies, which would make the statement its own + /// ancestor. + MovedIntoItself { stmt: Statement, anchor: Statement }, /// The block is not part of a CFG's block list (it is a statement-owned /// single-block body, or detached), so block-list surgery cannot apply. NotInCfgBody(Block), @@ -301,7 +310,7 @@ impl fmt::Display for RewriteError { RewriteError::UnexpectedTerminator => { write!( f, - "expected a non-terminator: terminators are only inserted with insert_terminator" + "expected a non-terminator: terminators are only inserted with insert_terminator, and are not moved" ) } RewriteError::BlockHasTerminator(block) => { @@ -377,6 +386,12 @@ impl fmt::Display for RewriteError { "the replacement names body {body:?}, which encloses the statement" ) } + RewriteError::MovedIntoItself { stmt, anchor } => { + write!( + f, + "cannot move {stmt:?} next to {anchor:?}, which is nested in it" + ) + } RewriteError::YieldIndexOutOfRange { graph, index } => { write!(f, "yield index {index} is out of range for DiGraph {graph}") } @@ -932,7 +947,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { _ => {} } } - let (block, prev, next, operands, results, successors, is_terminator) = { + let (block, operands, results, successors, is_terminator) = { let item = self.live_statement(stmt)?; let block = match item.parent { Some(StatementParent::Block(block)) => block, @@ -940,8 +955,6 @@ impl<'a, L: Dialect> Rewriter<'a, L> { }; ( block, - item.node.prev, - item.node.next, item.definition .arguments() .copied() @@ -971,37 +984,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { remove_predecessor(self.stage, target, block); } } else { - // Unlink from the block's statement linked list (fix neighbours, or - // the block head/tail when the statement is at an end). - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = next; - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.head = next; - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = prev; - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.tail = prev; - } - } - } - if let Some(block_info) = block.get_info_mut(self.stage) - && block_info.statements.len > 0 - { - block_info.statements.len -= 1; - } + self.unlink_from_block(stmt, block); } // Drop the def-use edges this statement contributed, erase the bodies @@ -1148,82 +1131,150 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.insert_relative(anchor, definition, true) } - fn insert_relative( - &mut self, + /// Where a statement goes when placed before (`after == false`) or after + /// `anchor`: its block, and the chain statements that will sit directly + /// before and after it (`None` at either end of the chain). + /// + /// The terminator sits outside the chain, so a place before it is the end + /// of the chain, and nothing can follow it ([`RewriteError::AnchorIsTerminator`]). + /// Rejects an unknown `anchor` or one not in a block body. + pub(super) fn slot_at( + &self, anchor: Statement, - definition: L, after: bool, - ) -> Result { - let (block, prev, next) = { - let item = self.live_statement(anchor)?; - let block = match item.parent { - Some(StatementParent::Block(block)) => block, - _ => return Err(RewriteError::NotInBlockBody(anchor)), - }; - match (item.definition.is_terminator(), after) { - // Nothing can follow a terminator. - (true, true) => return Err(RewriteError::AnchorIsTerminator(anchor)), - // The terminator sits outside the chain, so inserting before it - // appends to the end of the chain - (true, false) => { - let tail = block - .get_info(self.stage) - .and_then(|b| b.statements.tail().copied()); - - // The tail ends the chain, so the new statement becomes the tail. - debug_assert!(tail.is_none_or(|t| t.next(self.stage).is_none())); - - (block, tail, None) - } - (false, true) => (block, Some(anchor), item.node.next), - (false, false) => (block, item.node.prev, Some(anchor)), - } + ) -> Result<(Block, Option, Option), RewriteError> { + let item = self.live_statement(anchor)?; + let block = match item.parent { + Some(StatementParent::Block(block)) => block, + _ => return Err(RewriteError::NotInBlockBody(anchor)), }; - - if definition.is_terminator() { - return Err(RewriteError::UnexpectedTerminator); - } - if definition.successors().next().is_some() { - return Err(RewriteError::NonTerminatorWithSuccessors); + match (item.definition.is_terminator(), after) { + (true, true) => Err(RewriteError::AnchorIsTerminator(anchor)), + (true, false) => { + let tail = block + .get_info(self.stage) + .and_then(|b| b.statements.tail().copied()); + // The tail ends the chain, so the statement becomes the tail. + debug_assert!(tail.is_none_or(|t| t.next(self.stage).is_none())); + Ok((block, tail, None)) + } + (false, true) => Ok((block, Some(anchor), item.node.next)), + (false, false) => Ok((block, item.node.prev, Some(anchor))), } - let (results, operands) = self.prepare_insert(&definition)?; - let new_stmt = self.alloc_statement( - definition, - StatementParent::Block(block), - prev, - next, - results, - &operands, - ); + } - // Wire neighbours (or the block head/tail when inserting at an end). + /// Link the live, unlinked `stmt` into `block`'s chain between `prev` and + /// `next` (the block head or tail when one is `None`), making `block` its + /// parent. + pub(super) fn link_into_block( + &mut self, + stmt: Statement, + block: Block, + prev: Option, + next: Option, + ) { + if let Some(info) = stmt.get_info_mut(self.stage) { + info.parent = Some(StatementParent::Block(block)); + info.node.prev = prev; + info.node.next = next; + } match prev { Some(p) => { if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = Some(new_stmt); + info.node.next = Some(stmt); } } None => { if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.head = Some(new_stmt); + block_info.statements.head = Some(stmt); } } } match next { Some(n) => { if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = Some(new_stmt); + info.node.prev = Some(stmt); } } None => { if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.tail = Some(new_stmt); + block_info.statements.tail = Some(stmt); } } } if let Some(block_info) = block.get_info_mut(self.stage) { block_info.statements.len += 1; } + } + + /// Unlink the non-terminator `stmt` from `block`'s chain, fixing its + /// neighbours (or the block head/tail at an end) and clearing its own + /// links. Its parent is left for the caller. + pub(super) fn unlink_from_block(&mut self, stmt: Statement, block: Block) { + let Some(node) = stmt.get_info(self.stage).map(|item| item.node) else { + return; + }; + let (prev, next) = (node.prev, node.next); + match prev { + Some(p) => { + if let Some(info) = p.get_info_mut(self.stage) { + info.node.next = next; + } + } + None => { + if let Some(block_info) = block.get_info_mut(self.stage) { + block_info.statements.head = next; + } + } + } + match next { + Some(n) => { + if let Some(info) = n.get_info_mut(self.stage) { + info.node.prev = prev; + } + } + None => { + if let Some(block_info) = block.get_info_mut(self.stage) { + block_info.statements.tail = prev; + } + } + } + if let Some(block_info) = block.get_info_mut(self.stage) + && block_info.statements.len > 0 + { + block_info.statements.len -= 1; + } + if let Some(info) = stmt.get_info_mut(self.stage) { + info.node.prev = None; + info.node.next = None; + } + } + + fn insert_relative( + &mut self, + anchor: Statement, + definition: L, + after: bool, + ) -> Result { + let (block, prev, next) = self.slot_at(anchor, after)?; + + if definition.is_terminator() { + return Err(RewriteError::UnexpectedTerminator); + } + if definition.successors().next().is_some() { + return Err(RewriteError::NonTerminatorWithSuccessors); + } + let (results, operands) = self.prepare_insert(&definition)?; + let new_stmt = self.alloc_statement( + definition, + StatementParent::Block(block), + prev, + next, + results, + &operands, + ); + + self.link_into_block(new_stmt, block, prev, next); // Predecessor mirror does not need to be updated since we rejected terminators earlier. self.events diff --git a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs index 7be45879a..acfc6d4b4 100644 --- a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs @@ -81,7 +81,7 @@ pub struct BuiltOrderCursor { impl WalkOrder for BFS { type Cursor = BFSCursor; - const ALLOWS_ERASING_VISITED: bool = false; + const ALLOWS_DISPLACING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { BFSCursor { @@ -183,7 +183,7 @@ impl BFSCursor { impl WalkOrder for ReverseBFS { type Cursor = BuiltOrderCursor; - const ALLOWS_ERASING_VISITED: bool = true; + const ALLOWS_DISPLACING_VISITED: bool = true; fn start(stage: &StageInfo, root: Statement) -> Self::Cursor { let mut ordered_nodes = breadth_first(stage, root, |_| {}); @@ -194,7 +194,7 @@ impl WalkOrder for ReverseBFS { impl WalkOrder for KeepOrderReverseBFS { type Cursor = BuiltOrderCursor; - const ALLOWS_ERASING_VISITED: bool = true; + const ALLOWS_DISPLACING_VISITED: bool = true; fn start(stage: &StageInfo, root: Statement) -> Self::Cursor { let mut level_starts = Vec::new(); diff --git a/crates/kirin-ir/src/rewrite/walk/depth_first.rs b/crates/kirin-ir/src/rewrite/walk/depth_first.rs index 8b10559ad..c294bdb5e 100644 --- a/crates/kirin-ir/src/rewrite/walk/depth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/depth_first.rs @@ -33,7 +33,7 @@ pub struct PostOrderCursor { impl WalkOrder for PreOrder { type Cursor = PreOrderCursor; - const ALLOWS_ERASING_VISITED: bool = false; + const ALLOWS_DISPLACING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { PreOrderCursor { @@ -46,7 +46,7 @@ impl WalkOrder for PreOrder { impl WalkOrder for PostOrder { type Cursor = PostOrderCursor; - const ALLOWS_ERASING_VISITED: bool = false; + const ALLOWS_DISPLACING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { PostOrderCursor { diff --git a/crates/kirin-ir/src/rewrite/walk/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index 2c3baca0d..14f84d756 100644 --- a/crates/kirin-ir/src/rewrite/walk/mod.rs +++ b/crates/kirin-ir/src/rewrite/walk/mod.rs @@ -74,6 +74,11 @@ //! - merging its block into another block //! ([`Rewriter::merge_blocks`](crate::Rewriter::merge_blocks) with its //! block as `succ`), reported as [`WalkError::VisitedMoved`]; +//! - moving it with [`Rewriter::move_before`](crate::Rewriter::move_before) +//! or [`Rewriter::move_after`](crate::Rewriter::move_after), reported as +//! [`WalkError::VisitedMoved`] by [`WalkOrder::walk_mut`] (a loop driving a +//! cursor by hand checks with [`WalkOrder::check_not_moved_since`]); +//! - moving one of its ancestors, which is not checked; //! - reordering the bodies of one of its ancestors //! ([`Rewriter::replace_statement`](crate::Rewriter::replace_statement)), //! which is not checked. @@ -93,7 +98,7 @@ mod breadth_first; mod depth_first; mod order; -mod tree; +pub(crate) mod tree; pub use breadth_first::{BFS, BFSCursor, BuiltOrderCursor, KeepOrderReverseBFS, ReverseBFS}; pub use depth_first::{PostOrder, PostOrderCursor, PreOrder, PreOrderCursor}; diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index f80fc3afc..e12bb3dd5 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -1,6 +1,6 @@ use std::fmt; -use crate::{Dialect, Rewriter, StageInfo, Statement}; +use crate::{Dialect, MutationEvent, Rewriter, StageInfo, Statement}; /// A traversal order over the statements nested under a root statement. /// @@ -13,9 +13,11 @@ pub trait WalkOrder { /// The cursor walking this order. type Cursor: WalkCursor; - /// Whether an edit may erase the statement being visited, as the - /// [contract](super#contract) allows for this order. - const ALLOWS_ERASING_VISITED: bool; + /// Whether an edit may erase or move the statement being visited (see + /// the [contract](super#contract)). Only the reverse BFS orders allow it: + /// they build their order at the start, while the others find the next + /// statement through the visited one. + const ALLOWS_DISPLACING_VISITED: bool; /// A cursor at the start of the walk of the tree under `root`, `root` /// included. @@ -52,7 +54,8 @@ pub trait WalkOrder { /// [`walk`](Self::walk) with a `visit` that may edit through `rewriter`, /// within the walk's contract (see the [module docs](super)). Fails if /// `visit` erases the statement it is visiting where the order does not - /// allow it, or moves it by merging its block into another. + /// allow it, or moves it where the order does not allow it (see + /// [`check_not_moved_since`](Self::check_not_moved_since)). fn walk_mut<'r, L: Dialect>( rewriter: &mut Rewriter<'r, L>, root: Statement, @@ -60,12 +63,43 @@ pub trait WalkOrder { ) -> Result { let mut cursor = Self::start(rewriter.stage(), root); while let Some(stmt) = cursor.next(rewriter.stage())? { - if let Some(status) = visit(rewriter, stmt).steer(&mut cursor) { + let seen = rewriter.events().len(); + let result = visit(rewriter, stmt); + Self::check_not_moved_since(rewriter, stmt, seen)?; + if let Some(status) = result.steer(&mut cursor) { return Ok(status); } } Ok(WalkStatus::Completed) } + + /// Fail with [`WalkError::VisitedMoved`] if an event recorded since the + /// first `seen` events of `rewriter` moved `visited`, in an order that + /// does not allow it + /// ([`ALLOWS_DISPLACING_VISITED`](Self::ALLOWS_DISPLACING_VISITED)). + /// Panics if `rewriter` has recorded fewer than `seen` events. + /// + /// [`walk_mut`](Self::walk_mut) calls this after each visit. A loop driving + /// a cursor by hand calls it too: note `rewriter.events().len()` before + /// editing at a statement, and check before the next + /// [`WalkCursor::next`]. + fn check_not_moved_since( + rewriter: &Rewriter<'_, L>, + visited: Statement, + seen: usize, + ) -> Result<(), WalkError> { + if Self::ALLOWS_DISPLACING_VISITED { + return Ok(()); + } + let moved = rewriter.events()[seen..].iter().any( + |event| matches!(event, MutationEvent::MovedStatement { stmt } if *stmt == visited), + ); + if moved { + Err(WalkError::VisitedMoved(visited)) + } else { + Ok(()) + } + } } /// What a walk does after visiting a statement. @@ -130,10 +164,13 @@ pub enum WalkError { /// The statement being visited was erased before the walk proceeded, in a /// way its order does not allow (see the [module docs](super)). VisitedErased(Statement), - /// The statement being visited was moved before the walk proceeded, by - /// merging its block into another block + /// The statement being visited was moved before the walk proceeded, in an + /// order a move breaks: by merging its block into another block /// ([`Rewriter::merge_blocks`](crate::Rewriter::merge_blocks) with its - /// block as `succ`). + /// block as `succ`), or with + /// [`Rewriter::move_before`](crate::Rewriter::move_before) / + /// [`Rewriter::move_after`](crate::Rewriter::move_after) (see + /// [`WalkOrder::check_not_moved_since`]). VisitedMoved(Statement), } @@ -149,7 +186,7 @@ impl fmt::Display for WalkError { WalkError::VisitedMoved(stmt) => { write!( f, - "statement {stmt:?} was moved into another block while the walk was visiting it" + "statement {stmt:?} was moved while the walk was visiting it" ) } } diff --git a/crates/kirin-ir/src/rewrite/walk/tree.rs b/crates/kirin-ir/src/rewrite/walk/tree.rs index 4c6e10e46..ca573b604 100644 --- a/crates/kirin-ir/src/rewrite/walk/tree.rs +++ b/crates/kirin-ir/src/rewrite/walk/tree.rs @@ -69,7 +69,7 @@ pub(super) fn next_sibling(stage: &StageInfo, stmt: Statement) -> /// The statement owning the body `stmt` is in, or `None` for a top-level or /// dead statement. -pub(super) fn parent(stage: &StageInfo, stmt: Statement) -> Option { +pub(crate) fn parent(stage: &StageInfo, stmt: Statement) -> Option { match stmt.get_info(stage)?.parent? { StatementParent::Block(block) => match block.get_info(stage)?.parent()? { BlockParent::CFG(cfg) => cfg.get_info(stage)?.parent, diff --git a/crates/kirin-ir/tests/rewriter.rs b/crates/kirin-ir/tests/rewriter.rs index 868fb2833..8f1da4691 100644 --- a/crates/kirin-ir/tests/rewriter.rs +++ b/crates/kirin-ir/tests/rewriter.rs @@ -1417,3 +1417,226 @@ fn replace_uses_if_separates_a_digraph_operand_from_its_yield() { assert_eq!(dg.expect_info(&stage).graph().edge_count(), 0); verify_derived(&stage).unwrap(); } + +/// ```text +/// entry: p = use %v; q = nop; r = nop; owner = own_blocks { s = nop } { }; ret +/// ``` +/// `%v` is the entry block's argument. +struct MoveFixture { + stage: StageInfo, + entry: Block, + arm: Block, + v: SSAValue, + p: Statement, + q: Statement, + r: Statement, + s: Statement, + owner: Statement, + ret: Statement, +} + +fn move_fixture() -> MoveFixture { + let mut stage = new_stage(); + let v = stage.block_argument().index(0); + let p = stage.statement().definition(BuilderDialect::Use(v)).new(); + let q = stage.statement().definition(BuilderDialect::Nop).new(); + let r = stage.statement().definition(BuilderDialect::Nop).new(); + let s = stage.statement().definition(BuilderDialect::Nop).new(); + let arm = stage.block().stmt(s).new(); + let empty = stage.block().new(); + let owner = stage + .statement() + .definition(BuilderDialect::OwnBlocks(arm, empty)) + .new(); + let ret = stage.statement().definition(BuilderDialect::Return).new(); + let entry = stage + .block() + .argument(TestType::I32) + .stmt(p) + .stmt(q) + .stmt(r) + .stmt(owner) + .terminator(ret) + .new(); + let stage = stage.finalize().unwrap(); + let v = SSAValue::from(entry.expect_info(&stage).arguments()[0]); + MoveFixture { + stage, + entry, + arm, + v, + p, + q, + r, + s, + owner, + ret, + } +} + +fn order(stage: &StageInfo, block: Block) -> Vec { + block.statements(stage).collect() +} + +#[test] +fn move_before_and_after_reorder_a_block() { + let MoveFixture { + mut stage, + entry, + p, + q, + r, + owner, + ret, + .. + } = move_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + + rewriter.move_after(p, r).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [q, r, p, owner]); + rewriter.move_before(owner, q).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [owner, q, r, p]); + // Next to its own neighbour. + rewriter.move_after(q, r).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [owner, r, q, p]); + // Before the terminator: to the end of the chain. + rewriter.move_before(owner, ret).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [r, q, p, owner]); + assert_eq!( + rewriter.events().to_vec(), + vec![ + MutationEvent::MovedStatement { stmt: p }, + MutationEvent::MovedStatement { stmt: owner }, + MutationEvent::MovedStatement { stmt: q }, + MutationEvent::MovedStatement { stmt: owner }, + ] + ); + drop(rewriter); + assert_eq!(entry.terminator(&stage), Some(ret)); + verify_derived(&stage).unwrap(); +} + +#[test] +fn moving_into_and_out_of_a_nested_block_keeps_uses() { + let MoveFixture { + mut stage, + entry, + arm, + v, + p, + q, + r, + s, + owner, + .. + } = move_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + + // `p` reads `%v`; moving it changes neither its operand nor `%v`'s uses. + rewriter.move_before(p, s).unwrap(); + assert_eq!(order(rewriter.stage(), arm), [p, s]); + assert_eq!(order(rewriter.stage(), entry), [q, r, owner]); + assert_eq!( + *p.parent(rewriter.stage()), + Some(StatementParent::Block(arm)) + ); + assert_eq!(uses_set(rewriter.stage(), v), HashSet::from([so(p, 0)])); + + rewriter.move_after(s, r).unwrap(); + assert_eq!(order(rewriter.stage(), arm), [p]); + assert_eq!(order(rewriter.stage(), entry), [q, r, s, owner]); + drop(rewriter); + verify_derived(&stage).unwrap(); +} + +#[test] +fn moving_a_statement_next_to_itself_does_nothing() { + let MoveFixture { + mut stage, + entry, + p, + q, + r, + owner, + .. + } = move_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + rewriter.move_before(q, q).unwrap(); + rewriter.move_after(q, q).unwrap(); + assert!(rewriter.events().is_empty()); + assert_eq!(order(rewriter.stage(), entry), [p, q, r, owner]); +} + +#[test] +fn move_rejects_illegal_moves_leaving_the_stage_untouched() { + let MoveFixture { + mut stage, + entry, + arm, + p, + q, + r, + s, + owner, + ret, + .. + } = move_fixture(); + let dead = stage.statement_arena().next_id(); + let mut rewriter = Rewriter::new(&mut stage); + + assert_eq!( + rewriter.move_before(dead, q), + Err(RewriteError::UnknownStatement(dead)) + ); + assert_eq!( + rewriter.move_before(q, dead), + Err(RewriteError::UnknownStatement(dead)) + ); + assert_eq!( + rewriter.move_before(ret, q), + Err(RewriteError::UnexpectedTerminator) + ); + assert_eq!( + rewriter.move_after(q, ret), + Err(RewriteError::AnchorIsTerminator(ret)) + ); + // Into its own block. + assert_eq!( + rewriter.move_before(owner, s), + Err(RewriteError::MovedIntoItself { + stmt: owner, + anchor: s + }) + ); + assert!(rewriter.events().is_empty()); + drop(rewriter); + + assert_eq!(order(&stage, entry), [p, q, r, owner]); + assert_eq!(order(&stage, arm), [s]); + verify_derived(&stage).unwrap(); +} + +#[test] +fn move_rejects_graph_members() { + let mut stage = new_stage(); + let member = stage.statement().definition(BuilderDialect::Nop).new(); + let dg = stage.digraph().node(member).new(); + let owner = stage + .statement() + .definition(BuilderDialect::OwnDiGraph(dg)) + .new(); + let q = stage.statement().definition(BuilderDialect::Nop).new(); + let _entry = stage.block().stmt(owner).stmt(q).new(); + let mut stage = stage.finalize().unwrap(); + + let mut rewriter = Rewriter::new(&mut stage); + assert_eq!( + rewriter.move_before(member, q), + Err(RewriteError::NotInBlockBody(member)) + ); + assert_eq!( + rewriter.move_before(q, member), + Err(RewriteError::NotInBlockBody(member)) + ); + assert!(rewriter.events().is_empty()); +} diff --git a/crates/kirin-ir/tests/walk.rs b/crates/kirin-ir/tests/walk.rs index bc1547e68..53f7cc4e9 100644 --- a/crates/kirin-ir/tests/walk.rs +++ b/crates/kirin-ir/tests/walk.rs @@ -766,15 +766,104 @@ fn reverse_bfs_orders_visit_statements_where_they_were_before_a_merge() { assert_eq!(result.unwrap(), without_br); } +/// Walk `O` over [`tree`] with `walk_mut`, running `edit` at each statement; +/// the statements visited, or the walk's error. +fn walk_mut_editing( + t: &mut Tree, + mut edit: impl FnMut(&mut Rewriter, Statement), +) -> Result, WalkError> { + let root = t.root; + let mut rewriter = Rewriter::new(&mut t.stage); + let mut visited = Vec::new(); + O::walk_mut(&mut rewriter, root, |rewriter, stmt| { + visited.push(stmt); + edit(rewriter, stmt); + WalkResult::Advance + })?; + Ok(visited) +} + +/// Moves `a` into the `then` arm, before `then_nop`, while visiting it. +fn move_a_while_visiting() -> (Statement, Result, WalkError>) { + let mut t = tree(); + let (a, then_nop) = (t.a, t.then_nop); + let result = walk_mut_editing::(&mut t, |rewriter, stmt| { + if stmt == a { + rewriter.move_before(a, then_nop).unwrap(); + } + }); + (a, result) +} + #[test] -fn splitting_the_visited_statements_block_keeps_the_walk() { - fn split_at_z(rewriter: &mut Rewriter, m: MergeCfg) { - rewriter.split_block(m.z).unwrap(); +fn moving_the_visited_statement_breaks_link_following_walks() { + for (a, result) in [ + move_a_while_visiting::(), + move_a_while_visiting::(), + move_a_while_visiting::(), + ] { + assert_eq!(result, Err(WalkError::VisitedMoved(a))); } +} + +#[test] +fn reverse_bfs_orders_go_on_after_the_visited_statement_is_moved() { + let t = tree(); + let (_, reverse) = move_a_while_visiting::(); + assert_eq!(reverse.unwrap(), ReverseBFS::schedule(&t.stage, t.root)); + let (_, keep_order) = move_a_while_visiting::(); + assert_eq!( + keep_order.unwrap(), + KeepOrderReverseBFS::schedule(&t.stage, t.root) + ); +} + +/// While visiting `a`, moves `else_nop` into the `then` arm, ahead of the +/// cursor in every link-following order. +fn move_else_nop_while_visiting_a() -> (Statement, Result, WalkError>) +{ + let mut t = tree(); + let (a, else_nop, then_nop) = (t.a, t.else_nop, t.then_nop); + let result = walk_mut_editing::(&mut t, |rewriter, stmt| { + if stmt == a { + rewriter.move_after(else_nop, then_nop).unwrap(); + } + }); + (else_nop, result) +} + +#[test] +fn moving_another_statement_keeps_the_walk() { + for (else_nop, result) in [ + move_else_nop_while_visiting_a::(), + move_else_nop_while_visiting_a::(), + move_else_nop_while_visiting_a::(), + ] { + assert!(result.unwrap().contains(&else_nop)); + } +} + +/// Splits the entry block at `a` while visiting `a`; the walk's schedule +/// before the split, and the statements it visited. +fn split_at_a_while_visiting() -> (Vec, Result, WalkError>) +{ + let mut t = tree(); + let a = t.a; + let schedule = O::schedule(&t.stage, t.root); + let result = walk_mut_editing::(&mut t, |rewriter, stmt| { + if stmt == a { + rewriter.split_block(a).unwrap(); + } + }); + (schedule, result) +} + +#[test] +fn splitting_the_visited_statements_block_keeps_the_walk() { for (schedule, result) in [ - edit_at_z::(split_at_z), - edit_at_z::(split_at_z), - edit_at_z::(split_at_z), + split_at_a_while_visiting::(), + split_at_a_while_visiting::(), + split_at_a_while_visiting::(), ] { assert_eq!(result.unwrap(), schedule); } diff --git a/crates/kirin-rewrite/src/driver/rewrite_driver.rs b/crates/kirin-rewrite/src/driver/rewrite_driver.rs index 9a3b62981..47f6c6c1a 100644 --- a/crates/kirin-rewrite/src/driver/rewrite_driver.rs +++ b/crates/kirin-rewrite/src/driver/rewrite_driver.rs @@ -118,7 +118,7 @@ impl<'r, L: Dialect> RewriteDriver<'r, L> { rewriter: &mut Rewriter<'_, L>, root: Statement, ) -> Result { - if !O::ALLOWS_ERASING_VISITED + if !O::ALLOWS_DISPLACING_VISITED && let Some(rule) = self.rules.rules().iter().find(|rule| rule.erases_root()) { let order = std::any::type_name::(); diff --git a/docs/design/rewrite-engine.md b/docs/design/rewrite-engine.md index 258b2d5df..3e0da8c80 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -1114,6 +1114,7 @@ pub enum MutationEvent { ChangedDiGraphYields { graph: DiGraph }, ChangedUnGraphTopology { graph: UnGraph }, ChangedBodyNesting { owner: Statement }, + MovedStatement { stmt: Statement }, } ``` @@ -1797,6 +1798,7 @@ stage remains valid. | `erase_statement` | Tombstone the statement, unlink it from its block list (or remove its graph node and the edges touching it, or drop an `UnGraph` edge statement from the member map), drop its operand uses, tombstone its (unused) result values. Statement-owned blocks, `Cfg`s, `DiGraph`s, and `UnGraph`s nested in it are tombstoned with it. Erasing a terminator clears the block's terminator and removes its outgoing edges, leaving the block unterminated. | `ErasedStatement` (covers the nested bodies; no CFG or graph events), then `ChangedTerminator` for a terminator, `ChangedDiGraphTopology` for a `DiGraph` member, or `ChangedUnGraphTopology` for an `UnGraph` member | parented by a block or graph; no live result uses (a `DiGraph` yield is a use, and so is a node reading an edge statement's wire) | | `new_result` | Allocate a typed result value for a statement not inserted yet; its arena slot stays empty (not live) until an insert binds it. | none | a value its `Rewriter` never binds is tombstoned when the `Rewriter` is dropped (`run_pass` drops it before verifying) | | `insert_before` / `insert_after` | Allocate a statement, splice it adjacent to `anchor`, bind its results, add operand uses; return the new id. | `InsertedStatement` | anchor is block-parented; only `insert_before` may anchor on the terminator, which appends to the end of the chain; `definition` is a non-terminator with no successors; operands must be live; results must be distinct, unbound `new_result` values | +| `move_before` / `move_after` | Unlink a statement and link it adjacent to `anchor`, in any block body; id, results, uses and nested bodies are unchanged. | `MovedStatement` | both are block-parented; the moved statement is not a terminator; only `move_before` may anchor on the terminator, which appends to the end of the chain; the anchor is not nested in the moved statement; moving next to itself is a no-op | | `insert_terminator` | Allocate `definition` as the terminator of a block that has none; add operand uses and predecessor entries. | `InsertedStatement`, `ChangedTerminator` | any block (`Cfg` member or statement-owned body); `definition` is a result-less terminator with live operands and successors | | `insert_node` | Allocate `definition` as a new `DiGraph` member and node; bind its results, add operand uses and an edge from each member producer of an operand. | `InsertedStatement`, `ChangedDiGraphTopology` | live graph; non-terminator with no successors; live operands (outside ones get no edge); results must be `new_result` values | | `replace_yield` | Rewrite one `DiGraph` yield slot; move its yield use `old → new`. Topology is unchanged. | `ChangedDiGraphYields` (none if slot unchanged) | live graph and value, index in range | @@ -2006,12 +2008,14 @@ axis — hence the column. | `BodyOwnedElsewhere(body)` | guard | a body the replacement gains already belongs to another statement or, for a block, to a CFG | `replace_statement` | | `BodyListedTwice(body)` | guard | the replacement names the same body in two fields | `replace_statement` | | `BodyEnclosesOwner(body)` | guard | the replacement names a body the statement is nested in, which would make it its own ancestor | `replace_statement` | +| `MovedIntoItself { stmt, anchor }` | guard | the anchor is nested in the moved statement, which would make it its own ancestor | `move_before`, `move_after` | The immediate `MutationEvent` vocabulary recorded by these methods is `ChangedOperands`, `ReplacedUses`, `InsertedStatement`, `ErasedStatement`, `ReplacedStatement`, `ChangedBlockList { cfg }`, `ChangedCfg { cfg }`, `ChangedTerminator { block }`, `ChangedDiGraphTopology { graph }`, -`ChangedDiGraphYields { graph }`, and `ChangedBodyNesting { owner }`. `ErasedStatement` also covers everything nested +`ChangedDiGraphYields { graph }`, `ChangedBodyNesting { owner }`, and +`MovedStatement { stmt }`. `ErasedStatement` also covers everything nested in the erased statement; no separate CFG or graph events are recorded for those bodies. The remaining documented event kinds (`UnGraph` and result-type events) grow in as @@ -2859,7 +2863,7 @@ sweep recording no `MutationEvent`. Decisions to revisit in later milestones: - A rule may erase its root, the statement being visited, only in a walk - order that allows it (`WalkOrder::ALLOWS_ERASING_VISITED`: the reverse BFS + order that allows it (`WalkOrder::ALLOWS_DISPLACING_VISITED`: the reverse BFS orders). A run in another order fails with `DriverError::ErasesRoot` before editing anything. A rule's other erase targets are producers of the root, never nested in or around it, so they meet the walk contract in @@ -3036,12 +3040,13 @@ implement the derive and its matching/remapping bodies. ### Gap 3 — MutationEvent vocabulary vs implementation — RESOLVED DELIVERY RULE The documented `MutationEvent` enum lists the full target vocabulary; the -implementation currently emits **twelve** — `ChangedOperands`, `ReplacedUses`, -`InsertedStatement`, `ErasedStatement`, `ReplacedStatement`, `ChangedBlockList`, -`ChangedCfg`, `ChangedTerminator`, `ChangedDiGraphTopology`, -`ChangedDiGraphYields`, `ChangedUnGraphTopology`, `ChangedBodyNesting` — -covering operand and yield rewriting, block-body statement and terminator -surgery (erase/insert/replace), `Cfg` block-list surgery, `DiGraph` and +implementation currently emits **thirteen** — `ChangedOperands`, +`ReplacedUses`, `InsertedStatement`, `ErasedStatement`, `ReplacedStatement`, +`ChangedBlockList`, `ChangedCfg`, `ChangedTerminator`, +`ChangedDiGraphTopology`, `ChangedDiGraphYields`, `ChangedUnGraphTopology`, +`ChangedBodyNesting`, `MovedStatement` — covering operand and yield +rewriting, block-body statement and terminator surgery +(erase/insert/replace/move), `Cfg` block-list surgery, `DiGraph` and `UnGraph` member surgery, and body changes through `replace_statement`. The remaining documented variants (result-type events) are not emitted because that surgery is not yet implemented, so an analysis-invalidation policy written against them would diff --git a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md new file mode 100644 index 000000000..3f79cf6ab --- /dev/null +++ b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md @@ -0,0 +1,252 @@ +# Remaining IR operations in the `kirin-ir` `Rewriter` + +Branch: `alex/remaining-ir-operations`, off `rust`. + +Adds the `Rewriter` operations needed for code motion and inlining: moving +statements and blocks, splicing whole bodies, copying bodies, and (as an +attempt) inlining a call. + +## Background + +### What exists + +| Body | Operations | +|---|---| +| block body | `erase_statement`, `insert_before` / `insert_after` (new statements only), `replace_statement`, `insert_terminator` | +| CFG | `append_block`, `insert_block_before` / `_after` (new empty blocks only), `erase_block`, `split_block`, `merge_blocks`, `retarget_successor`, `insert_block_arg`, `erase_block_arg` | +| values | `replace_operand`, `replace_all_uses`, `replace_results`, `new_result` | +| `DiGraph` | `insert_node`, `replace_yield` (removal through `erase_statement`) | +| `UnGraph` | `insert_ungraph_node`, `insert_ungraph_edge` (removal through `erase_statement`) | + +Nothing moves an existing statement or block, and nothing copies IR. The one +operation that moves IR is `merge_blocks`, which appends one block's +statements to the end of another block of the same CFG. + +### Cost model + +Each statement stores its parent block (`StatementParent::Block`), and each +block stores its parent CFG (`BlockParent::CFG`). Values store their defining +statement or block, and uses store the using statement. So: + +- moving a **statement** costs O(1): relink it, change its parent; +- moving a **block** costs O(1), and its statements don't change; +- moving the **statements of a block** into another block costs O(number + moved), one parent write each (this is why `split_block` costs O(statements + after the split point)); +- moving a **whole CFG's blocks** costs O(number of blocks). + +Ids never change when IR moves, so uses stay valid for free. A copy creates new +ids and must map every operand, which costs O(size of what is copied). + +### Comparison with MLIR + +| MLIR (`RewriterBase`, `Region`, `Operation`) | Kirin today | This plan | +|---|---|---| +| `moveOpBefore` / `moveOpAfter` | none | step 2 | +| `Block::moveBefore` | none | step 3 | +| `inlineBlockBefore(source, dest, before, args)` | `merge_blocks` (end of block, same CFG only) | step 4 | +| `inlineRegionBefore(region, parent, before)` | none | step 5 | +| `Operation::clone`, `Region::cloneInto` with `IRMapping` | none | step 6 | +| `inlineCall` + `DialectInlinerInterface` | none | step 7 (attempt) | +| `replaceUsesWithIf`, `replaceAllUsesExcept` | `replace_all_uses` only | step 1 | +| `Operation::isBeforeInBlock` | none | not in this plan (dominance comes later) | +| `Value::setType` | none | not in this plan | +| `replaceOpWithNewOp` | `replace_statement` (same result count) | not needed | +| `eraseOpResults` | `remove_result` deferred | not in this plan | + +## Steps (pause for review after each) + +Each step adds the operation to `Rewriter`, its rejections as `RewriteError` +variants where needed, its `MutationEvent`, tests in `kirin-ir`, and its row in +the operation table of `docs/design/rewrite-engine.md`. + +### 1. Replace only some uses + +```rust +rewriter.replace_uses_if(old, new, |stage: &StageInfo, site: Use| ...)?; +rewriter.replace_all_uses_except(old, new, stmt)?; +``` + +MLIR's `replaceUsesWithIf` / `replaceAllUsesExcept`. Like `replace_all_uses`, +but only the use sites the predicate accepts, or all but those of `stmt`. The +predicate gets the stage, so it can look up the using statement's parent or +neighbors; it runs on every use before anything is changed. CSE +and code motion need this when only some uses may change, e.g. only the uses +after a point. Same events as `replace_all_uses`, counting only the changed +slots. + +Cost: O(uses of `old`). + +### 2. Move a statement + +```rust +rewriter.move_before(stmt, anchor)?; +rewriter.move_after(stmt, anchor)?; +``` + +Same anchor rules as `insert_before` / `insert_after`: `move_before` a +terminator appends to the end of that block's statements; nothing can go after +a terminator. The anchor may be in any block, of any CFG or statement-owned +body. + +Example, hoisting out of an `scf.if` arm: + +``` +^entry: ^entry: + %r = if %c then { %t = neg %x // moved + %t = neg %x ==> %r = if %c then { + yield %t yield %t + } ... } ... +``` + +Rejects: +- `stmt` or `anchor` unknown; +- `stmt` is a terminator (terminators are not in the statement chain); +- `stmt` or `anchor` is a graph member; +- the anchor is inside `stmt`'s own nested bodies (moving an `scf.if` into its + own arm would make it contain itself). + +Not checked, like `insert_before`: that every use of `stmt`'s results still +comes after it. That is the caller's job. + +Records a new `MutationEvent::MovedStatement { stmt }`. `ErasedStatement` + +`InsertedStatement` would wrongly say the statement was tombstoned. + +**Walk check.** Moving the visited statement breaks pre-order, post-order and +BFS, which find the next statement through the visited one. Today the walk +only notices the `merge_blocks` case (the visited statement's old block was +erased). This step adds a check for moves, read from the `Rewriter`'s event +log, so nothing new is stored in the IR: + +- Every operation that moves IR records an event naming what moved: + `MovedStatement { stmt }` here, and the events of steps 3–5. +- `WalkCursor` gets a variant of `next` that takes the `Rewriter` instead of + the stage. At each visit the cursor remembers the visited statement, its + block, and the length of the event log. On its next step it scans only the + events recorded since, and reports `WalkError::VisitedMoved` if one moved the + visited statement or its block. When nothing was recorded, that is one + length comparison. +- `walk_mut` uses the new variant. The read-only `walk` and `schedule` keep + `next(stage)`, since nothing moves during them. +- The reverse BFS orders allow moves (their order is built at the start), so + they skip the check. +- `split_block` records no move event, so the split the contract allows still + passes. Moving an ancestor of the visited statement is not checked. + +Cost: O(1), plus O(nesting depth) for the cycle check. + +### 3. Move a block within its CFG + +```rust +rewriter.move_block_before(block, anchor)?; +rewriter.move_block_after(block, anchor)?; +``` + +Reorders a CFG's block list. Branches, arguments and statements are unchanged. +Both blocks must be in the same CFG. Records `ChangedBlockList`, whose docs +already cover "moved within `cfg`'s block list", and a new +`MovedBlock { block }` for the walk check (step 2). + +The entry block is the first block in the list, so moving a block to the +front, or moving the first block away, changes which block runs first. That is +allowed and documented, not rejected. + +Cost: O(1). + +### 4. Splice a block's statements into another block + +```rust +rewriter.splice_block_before(source, anchor, &arg_replacements)?; +``` + +MLIR's `inlineBlockBefore`. Moves `source`'s statements (not its terminator) +in front of `anchor`, and redirects uses of `source`'s arguments to +`arg_replacements`. `source` is left empty, keeping its arguments and +terminator, and is not erased: the caller decides what to do with it. +Generalizes `merge_blocks`: any position, in any block of any CFG or +statement-owned body. Records a new `SplicedBlock { source, dest }`, which the +walk check reads as "every statement of `source` moved". + +Example, folding `scf.if` with a known condition: splice the `then` arm in +front of the `if`, replace the `if`'s results with the arm's `yield` values, +erase the `if`. + +Cost: O(statements moved). + +### 5. Splice a whole CFG into another + +```rust +rewriter.splice_cfg_after(source, anchor_block)?; +``` + +MLIR's `inlineRegionBefore`. Moves every block of `source` into `anchor_block`'s +CFG, right after `anchor_block`. Branches between the moved blocks stay valid, +since they all move together. `source` is left empty and not erased; its +owner (e.g. the callee function) still names it, and the caller decides +whether to erase the owner. Records `ChangedBlockList` for both CFGs and a new +`SplicedCfg { source, dest }`, which the walk check reads as "every block of +`source` moved". + +Cost: O(blocks moved). This is the move-instead-of-copy path for inlining a +callee's last use. + +### 6. Copy a statement or a CFG + +```rust +let mut map = IrMap::default(); // MLIR's IRMapping +let copy = rewriter.clone_before(stmt, anchor, &mut map)?; +let blocks = rewriter.clone_cfg_after(source, anchor_block, &mut map)?; +``` + +A deep copy with fresh ids: statements, results, blocks, block arguments and +nested bodies, including `DiGraph` and `UnGraph` bodies. A copied graph gets +new ports, copied members and yields, and its topology is rebuilt from the +copied operands with the helpers `insert_node` / `insert_ungraph_node` +already use. Every operand and branch target inside the copy is mapped +through `map`. A value from outside the copied IR keeps pointing at the +original. The caller can pre-seed `map`, e.g. to map the callee's entry +arguments to the call's operands. `map` afterwards holds old → new for every +value and block copied, so the caller can find the copied `ret`s. + +Cost: O(size of what is copied). + +### 7. Inline a call (attempt) + +An inliner in `kirin-function`, which knows its own call and `ret` ops, built +from steps 4–6 and the existing `kirin-ir` operations. It takes a closure that +builds the branch op, since `kirin-function` doesn't know which branch op a +language uses. It inlines `CallSpecialized` calls within the same stage; the +other call kinds first need resolving to a specialization. For a call `c` in block `B`: + +1. `split_block` after `c`. The continuation block `C` gets the rest of `B`. +2. Copy the callee's CFG between `B` and `C`, mapping its entry arguments to + `c`'s operands (step 6). For the callee's last use, splice instead (step 5). +3. Splice the copied entry block into `B` (step 4), or branch to it. +4. Replace each copied `ret v` with a branch to `C` passing `v`. +5. Give `C` one argument per result of `c`, redirect `c`'s results to them, and + erase `c`. + +Cost per call: O(statements after `c` in `B`) for the split, plus O(callee +size) for the copy. Inlining a block's calls from last to first makes the +splits cost O(block size) in total, since each statement moves at most once. + +### 8. Docs + +`Rewriter` module docs, the design doc's operation table, and the +`MutationEvent` count in Gap 3. The `kirin_ir::walk` contract gets a line on +the new moves. + +## Decisions + +1. Moves record new events naming what moved: `MovedStatement`, + `MovedBlock`, `SplicedBlock`, `SplicedCfg`. +2. The walk checks for a moved visited statement or block by reading the + `Rewriter`'s event log (step 2). Ancestors are not checked. +3. Moving a block to or from the front of its CFG is allowed, not rejected. +4. A block's statements can be spliced into any block. +5. Splicing leaves the source block or CFG empty and does not erase it. +6. Copying includes graph bodies. +7. Inlining lives in `kirin-function`, with a closure building the branch. +8. Only `CallSpecialized` within the same stage is inlined. +9. Of the extra MLIR operations, only `replace_uses_if` / + `replace_all_uses_except` (step 1). Their predicate gets the stage. From 8c815c81b431ba5c61269650d6086f3c05e33df3 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Thu, 5 Nov 2026 17:01:09 -0500 Subject: [PATCH 04/14] feat(ir): splice a block's statements with Rewriter::splice_block_before/after Moves every non-terminator statement of a block next to an anchor and redirects its arguments; records SplicedBlock, which walk_mut checks. Drops the planned block move (step 3). Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/cfg/merge.rs | 42 +-- crates/kirin-ir/src/rewrite/moves.rs | 161 +++++++++++- crates/kirin-ir/src/rewrite/rewriter.rs | 13 + crates/kirin-ir/src/rewrite/walk/mod.rs | 4 + crates/kirin-ir/src/rewrite/walk/order.rs | 34 ++- crates/kirin-ir/tests/rewriter.rs | 246 ++++++++++++++++++ crates/kirin-ir/tests/walk.rs | 84 ++++++ docs/design/rewrite-engine.md | 25 +- .../2026-11-04-remaining-ir-operations.md | 43 ++- 9 files changed, 580 insertions(+), 72 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/cfg/merge.rs b/crates/kirin-ir/src/rewrite/cfg/merge.rs index da06fad96..369a1652c 100644 --- a/crates/kirin-ir/src/rewrite/cfg/merge.rs +++ b/crates/kirin-ir/src/rewrite/cfg/merge.rs @@ -85,44 +85,20 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.redirect_uses(&pairs); // Move succ's chain and terminator onto the end of pred. - let (succ_chain, succ_terminator) = { - let info = succ.expect_info(self.stage); - (info.statements, info.terminator) + let (succ_chain, succ_terminator) = match succ.get_info_mut(self.stage) { + Some(info) => (std::mem::take(&mut info.statements), info.terminator.take()), + None => (LinkedList::new(), None), }; - let moved = self.chain_from(succ_chain.head); - for &stmt in moved.iter().chain(succ_terminator.iter()) { - if let Some(info) = stmt.get_info_mut(self.stage) { - info.parent = Some(StatementParent::Block(pred)); - } - } - if let Some(succ_head) = succ_chain.head { - match pred.expect_info(self.stage).statements.tail { - Some(tail) => { - if let Some(info) = tail.get_info_mut(self.stage) { - info.node.next = Some(succ_head); - } - if let Some(info) = succ_head.get_info_mut(self.stage) { - info.node.prev = Some(tail); - } - } - None => { - if let Some(info) = pred.get_info_mut(self.stage) { - info.statements.head = Some(succ_head); - } - } - } - if let Some(info) = pred.get_info_mut(self.stage) { - info.statements.tail = succ_chain.tail; - } + let pred_tail = pred.expect_info(self.stage).statements.tail; + self.link_chain(succ_chain, pred, pred_tail, None); + if let Some(terminator) = succ_terminator + && let Some(info) = terminator.get_info_mut(self.stage) + { + info.parent = Some(StatementParent::Block(pred)); } if let Some(info) = pred.get_info_mut(self.stage) { - info.statements.len += succ_chain.len; info.terminator = succ_terminator; } - if let Some(info) = succ.get_info_mut(self.stage) { - info.statements = LinkedList::new(); - info.terminator = None; - } // succ is now empty: tombstoning it only removes the block and its // (already redirected) arguments. diff --git a/crates/kirin-ir/src/rewrite/moves.rs b/crates/kirin-ir/src/rewrite/moves.rs index 9eeda751b..5eced7ea9 100644 --- a/crates/kirin-ir/src/rewrite/moves.rs +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -1,7 +1,8 @@ -use crate::arena::GetInfo; +use crate::arena::{GetInfo, GetInfoMut}; +use crate::node::linked_list::LinkedList; use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter}; use crate::rewrite::walk::tree; -use crate::{Dialect, Statement, StatementParent}; +use crate::{Block, Dialect, SSAValue, Statement, StatementParent}; impl<'a, L: Dialect> Rewriter<'a, L> { /// Move `stmt` to immediately before `anchor`, which may be in any block @@ -76,4 +77,160 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.events.push(MutationEvent::MovedStatement { stmt }); Ok(()) } + + /// Move every statement of `source`, except its terminator, to + /// immediately before `anchor`, which may be in any block body, and + /// redirect every use of `source`'s `i`-th argument to + /// `arg_replacements[i]`. + /// + /// The statements keep their order, ids, results, uses and nested bodies. + /// `source` keeps its arguments and terminator and is not erased. `anchor` + /// follows the rules of [`Rewriter::move_before`]: before a terminator + /// means at the end of that block's statement chain. A replacement may be + /// one of `source`'s own arguments. + /// + /// Rejects, leaving the stage untouched, if `source` or `anchor` is + /// unknown, if `anchor` is not in a block body + /// ([`RewriteError::NotInBlockBody`]), if `anchor` is in `source` or + /// nested in one of its statements ([`RewriteError::SplicedIntoItself`]), + /// if the replacement count differs from `source`'s argument count + /// ([`RewriteError::BlockArgumentArityMismatch`]), or if a replacement is + /// not live. + /// + /// Records the [`Rewriter::replace_all_uses`] events for the argument + /// redirection, then [`MutationEvent::SplicedBlock`]. A walk visiting a + /// statement of `source` reports the splice (see [`crate::walk`]). + pub fn splice_block_before( + &mut self, + source: Block, + anchor: Statement, + arg_replacements: &[SSAValue], + ) -> Result<(), RewriteError> { + self.splice_relative(source, anchor, arg_replacements, false) + } + + /// Move every statement of `source`, except its terminator, to + /// immediately after `anchor`, and redirect `source`'s arguments. + /// + /// As [`Rewriter::splice_block_before`], except that `anchor` must not be + /// a terminator ([`RewriteError::AnchorIsTerminator`]): nothing follows + /// one. + pub fn splice_block_after( + &mut self, + source: Block, + anchor: Statement, + arg_replacements: &[SSAValue], + ) -> Result<(), RewriteError> { + self.splice_relative(source, anchor, arg_replacements, true) + } + + fn splice_relative( + &mut self, + source: Block, + anchor: Statement, + arg_replacements: &[SSAValue], + after: bool, + ) -> Result<(), RewriteError> { + let args: Vec = self + .live_block(source)? + .arguments + .iter() + .map(|&arg| arg.into()) + .collect(); + let (dest, prev, next) = self.slot_at(anchor, after)?; + // Check `anchor` and every statement whose body encloses it. If one + // of them is a statement of the block `source`, it would be spliced + // next to or inside itself. + let mut current = Some(anchor); + while let Some(stmt) = current { + if *stmt.parent(self.stage) == Some(StatementParent::Block(source)) { + return Err(RewriteError::SplicedIntoItself { source, anchor }); + } + current = tree::parent(self.stage, stmt); + } + if arg_replacements.len() != args.len() { + return Err(RewriteError::BlockArgumentArityMismatch { + block: source, + expected: args.len(), + found: arg_replacements.len(), + }); + } + for &value in arg_replacements { + if value.get_info(self.stage).is_none() { + return Err(RewriteError::UnknownValue(value)); + } + } + + let pairs: Vec<(SSAValue, SSAValue)> = args + .into_iter() + .zip(arg_replacements.iter().copied()) + .filter(|(old, new)| old != new) + .collect(); + self.redirect_uses(&pairs); + + let chain = match source.get_info_mut(self.stage) { + Some(info) => std::mem::take(&mut info.statements), + None => LinkedList::new(), + }; + self.link_chain(chain, dest, prev, next); + self.events + .push(MutationEvent::SplicedBlock { source, dest }); + Ok(()) + } + + /// Link `chain`, a whole statement chain already taken out of its block, + /// into `block` between `prev` and `next` (the block head or tail when one + /// is `None`), making `block` the parent of each statement. + pub(super) fn link_chain( + &mut self, + chain: LinkedList, + block: Block, + prev: Option, + next: Option, + ) { + let (Some(head), Some(tail)) = (chain.head, chain.tail) else { + return; + }; + let mut current = Some(head); + while let Some(stmt) = current { + let Some(info) = stmt.get_info_mut(self.stage) else { + break; + }; + info.parent = Some(StatementParent::Block(block)); + current = info.node.next; + } + if let Some(info) = head.get_info_mut(self.stage) { + info.node.prev = prev; + } + if let Some(info) = tail.get_info_mut(self.stage) { + info.node.next = next; + } + match prev { + Some(p) => { + if let Some(info) = p.get_info_mut(self.stage) { + info.node.next = Some(head); + } + } + None => { + if let Some(info) = block.get_info_mut(self.stage) { + info.statements.head = Some(head); + } + } + } + match next { + Some(n) => { + if let Some(info) = n.get_info_mut(self.stage) { + info.node.prev = Some(tail); + } + } + None => { + if let Some(info) = block.get_info_mut(self.stage) { + info.statements.tail = Some(tail); + } + } + } + if let Some(info) = block.get_info_mut(self.stage) { + info.statements.len += chain.len; + } + } } diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index f51b517ec..aad764dd7 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -143,6 +143,9 @@ pub enum MutationEvent { /// `stmt` was moved to another place in a block body, keeping its id, /// results, uses and nested bodies. MovedStatement { stmt: Statement }, + /// Every statement of `source`, except its terminator, was moved into + /// `dest`, keeping their ids, results, uses and nested bodies. + SplicedBlock { source: Block, dest: Block }, } /// An illegal mutation rejected by a [`Rewriter`]. @@ -235,6 +238,10 @@ pub enum RewriteError { /// statement's own bodies, which would make the statement its own /// ancestor. MovedIntoItself { stmt: Statement, anchor: Statement }, + /// `splice_block_before`/`splice_block_after` was given an anchor in + /// `source` or nested in one of its statements, which would place a + /// spliced statement next to or inside itself. + SplicedIntoItself { source: Block, anchor: Statement }, /// The block is not part of a CFG's block list (it is a statement-owned /// single-block body, or detached), so block-list surgery cannot apply. NotInCfgBody(Block), @@ -392,6 +399,12 @@ impl fmt::Display for RewriteError { "cannot move {stmt:?} next to {anchor:?}, which is nested in it" ) } + RewriteError::SplicedIntoItself { source, anchor } => { + write!( + f, + "cannot splice {source:?} next to {anchor:?}, which is in it" + ) + } RewriteError::YieldIndexOutOfRange { graph, index } => { write!(f, "yield index {index} is out of range for DiGraph {graph}") } diff --git a/crates/kirin-ir/src/rewrite/walk/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index 14f84d756..e62387ef5 100644 --- a/crates/kirin-ir/src/rewrite/walk/mod.rs +++ b/crates/kirin-ir/src/rewrite/walk/mod.rs @@ -78,6 +78,10 @@ //! or [`Rewriter::move_after`](crate::Rewriter::move_after), reported as //! [`WalkError::VisitedMoved`] by [`WalkOrder::walk_mut`] (a loop driving a //! cursor by hand checks with [`WalkOrder::check_not_moved_since`]); +//! - splicing its block's statements elsewhere with +//! [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) +//! or [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after), +//! reported the same way; //! - moving one of its ancestors, which is not checked; //! - reordering the bodies of one of its ancestors //! ([`Rewriter::replace_statement`](crate::Rewriter::replace_statement)), diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index e12bb3dd5..d7d3039b2 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -1,6 +1,6 @@ use std::fmt; -use crate::{Dialect, MutationEvent, Rewriter, StageInfo, Statement}; +use crate::{Dialect, MutationEvent, Rewriter, StageInfo, Statement, StatementParent}; /// A traversal order over the statements nested under a root statement. /// @@ -64,8 +64,9 @@ pub trait WalkOrder { let mut cursor = Self::start(rewriter.stage(), root); while let Some(stmt) = cursor.next(rewriter.stage())? { let seen = rewriter.events().len(); + let parent = *stmt.parent(rewriter.stage()); let result = visit(rewriter, stmt); - Self::check_not_moved_since(rewriter, stmt, seen)?; + Self::check_not_moved_since(rewriter, stmt, parent, seen)?; if let Some(status) = result.steer(&mut cursor) { return Ok(status); } @@ -76,24 +77,32 @@ pub trait WalkOrder { /// Fail with [`WalkError::VisitedMoved`] if an event recorded since the /// first `seen` events of `rewriter` moved `visited`, in an order that /// does not allow it - /// ([`ALLOWS_DISPLACING_VISITED`](Self::ALLOWS_DISPLACING_VISITED)). + /// ([`ALLOWS_DISPLACING_VISITED`](Self::ALLOWS_DISPLACING_VISITED)): + /// a move of `visited` itself, or a splice of `visited_parent`'s block. /// Panics if `rewriter` has recorded fewer than `seen` events. /// /// [`walk_mut`](Self::walk_mut) calls this after each visit. A loop driving - /// a cursor by hand calls it too: note `rewriter.events().len()` before - /// editing at a statement, and check before the next - /// [`WalkCursor::next`]. + /// a cursor by hand calls it too: note `rewriter.events().len()` and + /// `visited.parent(stage)` before editing at a statement, and check before + /// the next [`WalkCursor::next`]. The parent is needed because a splice + /// changes it: afterwards, a statement spliced into a block cannot be told + /// from one that was there already. fn check_not_moved_since( rewriter: &Rewriter<'_, L>, visited: Statement, + visited_parent: Option, seen: usize, ) -> Result<(), WalkError> { if Self::ALLOWS_DISPLACING_VISITED { return Ok(()); } - let moved = rewriter.events()[seen..].iter().any( - |event| matches!(event, MutationEvent::MovedStatement { stmt } if *stmt == visited), - ); + let moved = rewriter.events()[seen..].iter().any(|event| match *event { + MutationEvent::MovedStatement { stmt } => stmt == visited, + MutationEvent::SplicedBlock { source, .. } => { + visited_parent == Some(StatementParent::Block(source)) + } + _ => false, + }); if moved { Err(WalkError::VisitedMoved(visited)) } else { @@ -169,8 +178,11 @@ pub enum WalkError { /// ([`Rewriter::merge_blocks`](crate::Rewriter::merge_blocks) with its /// block as `succ`), or with /// [`Rewriter::move_before`](crate::Rewriter::move_before) / - /// [`Rewriter::move_after`](crate::Rewriter::move_after) (see - /// [`WalkOrder::check_not_moved_since`]). + /// [`Rewriter::move_after`](crate::Rewriter::move_after), or by splicing + /// its block with + /// [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) / + /// [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after) + /// (see [`WalkOrder::check_not_moved_since`]). VisitedMoved(Statement), } diff --git a/crates/kirin-ir/tests/rewriter.rs b/crates/kirin-ir/tests/rewriter.rs index 8f1da4691..1917971a1 100644 --- a/crates/kirin-ir/tests/rewriter.rs +++ b/crates/kirin-ir/tests/rewriter.rs @@ -1640,3 +1640,249 @@ fn move_rejects_graph_members() { ); assert!(rewriter.events().is_empty()); } + +/// ```text +/// entry(%x): a = nop; owner = own_blocks { arm(%t): u = use %t; w = nop; yret } { }; ret +/// ``` +struct SpliceFixture { + stage: StageInfo, + entry: Block, + arm: Block, + x: SSAValue, + t: SSAValue, + a: Statement, + u: Statement, + w: Statement, + yret: Statement, + owner: Statement, + ret: Statement, +} + +fn splice_fixture() -> SpliceFixture { + let mut stage = new_stage(); + let t = stage.block_argument().index(0); + let u = stage.statement().definition(BuilderDialect::Use(t)).new(); + let w = stage.statement().definition(BuilderDialect::Nop).new(); + let yret = stage.statement().definition(BuilderDialect::Return).new(); + let arm = stage + .block() + .argument(TestType::I32) + .stmt(u) + .stmt(w) + .terminator(yret) + .new(); + let empty = stage.block().new(); + let owner = stage + .statement() + .definition(BuilderDialect::OwnBlocks(arm, empty)) + .new(); + let a = stage.statement().definition(BuilderDialect::Nop).new(); + let ret = stage.statement().definition(BuilderDialect::Return).new(); + let entry = stage + .block() + .argument(TestType::I32) + .stmt(a) + .stmt(owner) + .terminator(ret) + .new(); + let stage = stage.finalize().unwrap(); + let x = arg_of(&stage, entry); + let t = arg_of(&stage, arm); + SpliceFixture { + stage, + entry, + arm, + x, + t, + a, + u, + w, + yret, + owner, + ret, + } +} + +fn arg_of(stage: &StageInfo, block: Block) -> SSAValue { + block.expect_info(stage).arguments()[0].into() +} + +#[test] +fn splice_block_before_moves_statements_and_redirects_arguments() { + let SpliceFixture { + mut stage, + entry, + arm, + x, + t, + a, + u, + w, + yret, + owner, + .. + } = splice_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + + rewriter.splice_block_before(arm, owner, &[x]).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [a, u, w, owner]); + for stmt in [u, w] { + assert_eq!( + *stmt.parent(rewriter.stage()), + Some(StatementParent::Block(entry)) + ); + } + // `arm` is left with its argument and terminator, and no statements. + assert_eq!(order(rewriter.stage(), arm), []); + assert_eq!(arm.terminator(rewriter.stage()), Some(yret)); + assert_eq!(arg_of(rewriter.stage(), arm), t); + // `u` now reads `%x`. + assert_eq!(uses_set(rewriter.stage(), x), HashSet::from([so(u, 0)])); + assert!(uses_set(rewriter.stage(), t).is_empty()); + assert_eq!( + rewriter.events().to_vec(), + vec![ + MutationEvent::ChangedOperands { stmt: u }, + MutationEvent::ReplacedUses { old: t, new: x }, + MutationEvent::SplicedBlock { + source: arm, + dest: entry, + }, + ] + ); + drop(rewriter); + verify_derived(&stage).unwrap(); +} + +#[test] +fn splice_block_after_and_before_a_terminator_place_the_chain() { + let SpliceFixture { + mut stage, + entry, + arm, + x, + a, + u, + w, + owner, + .. + } = splice_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + rewriter.splice_block_after(arm, a, &[x]).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [a, u, w, owner]); + drop(rewriter); + verify_derived(&stage).unwrap(); + + let SpliceFixture { + mut stage, + entry, + arm, + x, + a, + u, + w, + owner, + ret, + .. + } = splice_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + // Before the terminator: to the end of the chain. + rewriter.splice_block_before(arm, ret, &[x]).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [a, owner, u, w]); + drop(rewriter); + verify_derived(&stage).unwrap(); +} + +#[test] +fn splice_block_may_keep_its_own_arguments() { + let SpliceFixture { + mut stage, + entry, + arm, + t, + a, + u, + w, + owner, + .. + } = splice_fixture(); + let mut rewriter = Rewriter::new(&mut stage); + rewriter.splice_block_before(arm, owner, &[t]).unwrap(); + assert_eq!(order(rewriter.stage(), entry), [a, u, w, owner]); + assert_eq!(uses_set(rewriter.stage(), t), HashSet::from([so(u, 0)])); + assert_eq!( + rewriter.events().to_vec(), + vec![MutationEvent::SplicedBlock { + source: arm, + dest: entry, + }], + "nothing to redirect" + ); +} + +#[test] +fn splice_block_rejects_illegal_splices_leaving_the_stage_untouched() { + let SpliceFixture { + mut stage, + entry, + arm, + x, + a, + u, + w, + owner, + ret, + .. + } = splice_fixture(); + let dead_stmt = stage.statement_arena().next_id(); + let dead_block = Block::from(Id::from(TestSSAValue(9_999))); + let dead_value = SSAValue::from(Id::from(TestSSAValue(9_999))); + let mut rewriter = Rewriter::new(&mut stage); + + assert_eq!( + rewriter.splice_block_before(dead_block, a, &[]), + Err(RewriteError::UnknownBlock(dead_block)) + ); + assert_eq!( + rewriter.splice_block_before(arm, dead_stmt, &[x]), + Err(RewriteError::UnknownStatement(dead_stmt)) + ); + assert_eq!( + rewriter.splice_block_after(arm, ret, &[x]), + Err(RewriteError::AnchorIsTerminator(ret)) + ); + // The anchor is in the source block. + assert_eq!( + rewriter.splice_block_before(arm, w, &[x]), + Err(RewriteError::SplicedIntoItself { + source: arm, + anchor: w + }) + ); + // The anchor is nested in one of the source's statements (`owner`). + assert_eq!( + rewriter.splice_block_before(entry, u, &[x]), + Err(RewriteError::SplicedIntoItself { + source: entry, + anchor: u + }) + ); + assert_eq!( + rewriter.splice_block_before(arm, owner, &[]), + Err(RewriteError::BlockArgumentArityMismatch { + block: arm, + expected: 1, + found: 0, + }) + ); + assert_eq!( + rewriter.splice_block_before(arm, owner, &[dead_value]), + Err(RewriteError::UnknownValue(dead_value)) + ); + assert!(rewriter.events().is_empty()); + drop(rewriter); + + assert_eq!(order(&stage, entry), [a, owner]); + assert_eq!(order(&stage, arm), [u, w]); + verify_derived(&stage).unwrap(); +} diff --git a/crates/kirin-ir/tests/walk.rs b/crates/kirin-ir/tests/walk.rs index 53f7cc4e9..ee7d9da81 100644 --- a/crates/kirin-ir/tests/walk.rs +++ b/crates/kirin-ir/tests/walk.rs @@ -843,6 +843,90 @@ fn moving_another_statement_keeps_the_walk() { } } +/// The block `stmt` is in. +fn block_of(stage: &StageInfo, stmt: Statement) -> Block { + match *stmt.parent(stage) { + Some(StatementParent::Block(block)) => block, + other => panic!("{stmt:?} is not in a block: {other:?}"), + } +} + +/// Splices the `then` arm before `owner` while visiting `then_nop`, its only +/// statement. +fn splice_then_arm_while_visiting_it() +-> (Statement, Result, WalkError>) { + let mut t = tree(); + let (then_nop, owner) = (t.then_nop, t.owner); + let then_block = block_of(&t.stage, then_nop); + let result = walk_mut_editing::(&mut t, |rewriter, stmt| { + if stmt == then_nop { + rewriter + .splice_block_before(then_block, owner, &[]) + .unwrap(); + } + }); + (then_nop, result) +} + +#[test] +fn splicing_the_visited_statements_block_breaks_link_following_walks() { + for (then_nop, result) in [ + splice_then_arm_while_visiting_it::(), + splice_then_arm_while_visiting_it::(), + splice_then_arm_while_visiting_it::(), + ] { + assert_eq!(result, Err(WalkError::VisitedMoved(then_nop))); + } +} + +#[test] +fn reverse_bfs_orders_go_on_after_the_visited_statements_block_is_spliced() { + let t = tree(); + let (_, reverse) = splice_then_arm_while_visiting_it::(); + assert_eq!(reverse.unwrap(), ReverseBFS::schedule(&t.stage, t.root)); + let (_, keep_order) = splice_then_arm_while_visiting_it::(); + assert_eq!( + keep_order.unwrap(), + KeepOrderReverseBFS::schedule(&t.stage, t.root) + ); +} + +#[test] +fn splicing_another_block_into_the_visited_statements_block_keeps_the_walk() { + /// Splices the `else` arm after `a`, into the block `a` is in. + fn splice_else_arm_after_a(rewriter: &mut Rewriter, t: &Tree) { + let else_block = block_of(rewriter.stage(), t.else_nop); + rewriter.splice_block_after(else_block, t.a, &[]).unwrap(); + } + /// The schedule of `O` over [`tree`] with the splice already made. + fn spliced_schedule() -> Vec { + let mut t = tree(); + let ids = tree(); + splice_else_arm_after_a(&mut Rewriter::new(&mut t.stage), &ids); + O::schedule(&t.stage, t.root) + } + fn splice_while_visiting_a() -> Vec { + let mut t = tree(); + let ids = tree(); + walk_mut_editing::(&mut t, |rewriter, stmt| { + if stmt == ids.a { + splice_else_arm_after_a(rewriter, &ids); + } + }) + .unwrap() + } + // `else_nop` now follows `a`, and the walk visits it there. + assert_eq!( + splice_while_visiting_a::(), + spliced_schedule::() + ); + assert_eq!( + splice_while_visiting_a::(), + spliced_schedule::() + ); + assert_eq!(splice_while_visiting_a::(), spliced_schedule::()); +} + /// Splits the entry block at `a` while visiting `a`; the walk's schedule /// before the split, and the statements it visited. fn split_at_a_while_visiting() -> (Vec, Result, WalkError>) diff --git a/docs/design/rewrite-engine.md b/docs/design/rewrite-engine.md index 3e0da8c80..45332fcdb 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -1115,6 +1115,7 @@ pub enum MutationEvent { ChangedUnGraphTopology { graph: UnGraph }, ChangedBodyNesting { owner: Statement }, MovedStatement { stmt: Statement }, + SplicedBlock { source: Block, dest: Block }, } ``` @@ -1750,6 +1751,20 @@ impl<'a, L: Dialect> Rewriter<'a, L> { pub fn erase_block(&mut self, block: Block) -> Result<(), RewriteError>; + pub fn splice_block_before( + &mut self, + source: Block, + anchor: Statement, + arg_replacements: &[SSAValue], + ) -> Result<(), RewriteError>; + + pub fn splice_block_after( + &mut self, + source: Block, + anchor: Statement, + arg_replacements: &[SSAValue], + ) -> Result<(), RewriteError>; + pub fn split_block(&mut self, at: Statement) -> Result; pub fn merge_blocks( @@ -1799,6 +1814,7 @@ stage remains valid. | `new_result` | Allocate a typed result value for a statement not inserted yet; its arena slot stays empty (not live) until an insert binds it. | none | a value its `Rewriter` never binds is tombstoned when the `Rewriter` is dropped (`run_pass` drops it before verifying) | | `insert_before` / `insert_after` | Allocate a statement, splice it adjacent to `anchor`, bind its results, add operand uses; return the new id. | `InsertedStatement` | anchor is block-parented; only `insert_before` may anchor on the terminator, which appends to the end of the chain; `definition` is a non-terminator with no successors; operands must be live; results must be distinct, unbound `new_result` values | | `move_before` / `move_after` | Unlink a statement and link it adjacent to `anchor`, in any block body; id, results, uses and nested bodies are unchanged. | `MovedStatement` | both are block-parented; the moved statement is not a terminator; only `move_before` may anchor on the terminator, which appends to the end of the chain; the anchor is not nested in the moved statement; moving next to itself is a no-op | +| `splice_block_before` / `splice_block_after` | Move every statement of `source` but its terminator, in order, adjacent to `anchor` in any block body; redirect `source`'s arguments to `arg_replacements`. `source` keeps its arguments and terminator and is not erased. | the `replace_all_uses` events, `SplicedBlock` | the anchor is block-parented, not in `source` and not nested in its statements; only `splice_block_before` may anchor on the terminator, which appends to the end of the chain; exact argument arity; replacements must be live and may be `source`'s own arguments | | `insert_terminator` | Allocate `definition` as the terminator of a block that has none; add operand uses and predecessor entries. | `InsertedStatement`, `ChangedTerminator` | any block (`Cfg` member or statement-owned body); `definition` is a result-less terminator with live operands and successors | | `insert_node` | Allocate `definition` as a new `DiGraph` member and node; bind its results, add operand uses and an edge from each member producer of an operand. | `InsertedStatement`, `ChangedDiGraphTopology` | live graph; non-terminator with no successors; live operands (outside ones get no edge); results must be `new_result` values | | `replace_yield` | Rewrite one `DiGraph` yield slot; move its yield use `old → new`. Topology is unchanged. | `ChangedDiGraphYields` (none if slot unchanged) | live graph and value, index in range | @@ -1972,7 +1988,7 @@ axis — hence the column. | Variant | Category | Meaning | Raised by | |---|---|---|---| | `UnknownStatement(stmt)` | liveness | id does not resolve to a live (non-tombstoned) statement | all statement methods | -| `UnknownValue(value)` | liveness | SSA id does not resolve to a live value | `replace_operand`, `replace_all_uses`, `insert_*`, `replace_statement` | +| `UnknownValue(value)` | liveness | SSA id does not resolve to a live value | `replace_operand`, `replace_all_uses`, `insert_*`, `replace_statement`, `splice_block_*` | | `UnknownBlock(block)` | liveness | a block id (anchor, operand block, or successor target) does not resolve to a live block | `replace_statement`, CFG block-list methods | | `UnknownCfg(cfg)` | liveness | the `Cfg` id does not resolve to a live `Cfg` | `append_block` | | `UnknownDiGraph(graph)` | liveness | the `DiGraph` id does not resolve to a live `DiGraph` | `insert_node`, `replace_yield` | @@ -1986,7 +2002,7 @@ axis — hence the column. | `CannotMergeBlockIntoItself(block)` | guard | `merge_blocks` was given the same block twice | `merge_blocks` | | `ExpectedTerminator` | guard | the definition is not a terminator | `insert_terminator` | | `BlockHasTerminator(block)` | guard | the block already has a terminator | `insert_terminator` | -| `BlockArgumentArityMismatch { block, expected, found }` | guard | values do not line up one-for-one with the block's arguments | `merge_blocks` | +| `BlockArgumentArityMismatch { block, expected, found }` | guard | values do not line up one-for-one with the block's arguments | `merge_blocks`, `splice_block_*` | | `BlockArgumentIndexOutOfRange { block, index }` | liveness | argument index past the block's argument list (for insertion, past its end) | `insert_block_arg`, `erase_block_arg` | | `ReplacementIsMergedArgument { block, value }` | guard | a replacement is one of the merged block's own arguments, which the merge deletes | `merge_blocks` | | `BlockHasPredecessors { block, pred }` | guard | `pred` still branches to the block being removed; retarget its edges first | `erase_block`, `merge_blocks` | @@ -2009,13 +2025,14 @@ axis — hence the column. | `BodyListedTwice(body)` | guard | the replacement names the same body in two fields | `replace_statement` | | `BodyEnclosesOwner(body)` | guard | the replacement names a body the statement is nested in, which would make it its own ancestor | `replace_statement` | | `MovedIntoItself { stmt, anchor }` | guard | the anchor is nested in the moved statement, which would make it its own ancestor | `move_before`, `move_after` | +| `SplicedIntoItself { source, anchor }` | guard | the anchor is in `source` or nested in one of its statements, which would place a spliced statement next to or inside itself | `splice_block_before`, `splice_block_after` | The immediate `MutationEvent` vocabulary recorded by these methods is `ChangedOperands`, `ReplacedUses`, `InsertedStatement`, `ErasedStatement`, `ReplacedStatement`, `ChangedBlockList { cfg }`, `ChangedCfg { cfg }`, `ChangedTerminator { block }`, `ChangedDiGraphTopology { graph }`, -`ChangedDiGraphYields { graph }`, `ChangedBodyNesting { owner }`, and -`MovedStatement { stmt }`. `ErasedStatement` also covers everything nested +`ChangedDiGraphYields { graph }`, `ChangedBodyNesting { owner }`, +`MovedStatement { stmt }`, and `SplicedBlock { source, dest }`. `ErasedStatement` also covers everything nested in the erased statement; no separate CFG or graph events are recorded for those bodies. The remaining documented event kinds (`UnGraph` and result-type events) grow in as diff --git a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md index 3f79cf6ab..a4e3149f3 100644 --- a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md +++ b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md @@ -43,7 +43,7 @@ ids and must map every operand, which costs O(size of what is copied). | MLIR (`RewriterBase`, `Region`, `Operation`) | Kirin today | This plan | |---|---|---| | `moveOpBefore` / `moveOpAfter` | none | step 2 | -| `Block::moveBefore` | none | step 3 | +| `Block::moveBefore` | none | not in this plan (step 3, dropped) | | `inlineBlockBefore(source, dest, before, args)` | `merge_blocks` (end of block, same CFG only) | step 4 | | `inlineRegionBefore(region, parent, before)` | none | step 5 | | `Operation::clone`, `Region::cloneInto` with `IRMapping` | none | step 6 | @@ -135,37 +135,36 @@ log, so nothing new is stored in the IR: Cost: O(1), plus O(nesting depth) for the cycle check. -### 3. Move a block within its CFG +### 3. Move a block within its CFG (dropped) -```rust -rewriter.move_block_before(block, anchor)?; -rewriter.move_block_after(block, anchor)?; -``` - -Reorders a CFG's block list. Branches, arguments and statements are unchanged. -Both blocks must be in the same CFG. Records `ChangedBlockList`, whose docs -already cover "moved within `cfg`'s block list", and a new -`MovedBlock { block }` for the walk check (step 2). - -The entry block is the first block in the list, so moving a block to the -front, or moving the first block away, changes which block runs first. That is -allowed and documented, not rejected. - -Cost: O(1). +Not added. A CFG's block order has one semantic effect, choosing the entry +block (the first in the list); otherwise it only changes printing and +iteration order. A new entry block can already be made with +`insert_block_before` plus a branch, and steps 5 and 7 place blocks while +splicing, so nothing needs a separate move. Add it when block layout starts +to matter (e.g. a backend or an analysis seeded in reverse postorder). ### 4. Splice a block's statements into another block ```rust rewriter.splice_block_before(source, anchor, &arg_replacements)?; +rewriter.splice_block_after(source, anchor, &arg_replacements)?; ``` MLIR's `inlineBlockBefore`. Moves `source`'s statements (not its terminator) -in front of `anchor`, and redirects uses of `source`'s arguments to -`arg_replacements`. `source` is left empty, keeping its arguments and +next to `anchor`, and redirects uses of `source`'s arguments to +`arg_replacements`. The anchor rules are those of `move_before` / +`move_after`. `splice_block_after` is for step 7, where the callee's entry is +spliced after the call, the last statement of an unterminated block. `source` is left empty, keeping its arguments and terminator, and is not erased: the caller decides what to do with it. Generalizes `merge_blocks`: any position, in any block of any CFG or statement-owned body. Records a new `SplicedBlock { source, dest }`, which the -walk check reads as "every statement of `source` moved". +walk check reads as "every statement of `source` moved". Afterwards a spliced +statement's parent is `dest`, so the check also takes the visited statement's +parent from before the visit, and compares it with `source`. + +Rejects an anchor in `source` or nested in one of its statements +(`SplicedIntoItself`). Example, folding `scf.if` with a known condition: splice the `then` arm in front of the `if`, replace the `if`'s results with the arm's `yield` values, @@ -239,10 +238,10 @@ the new moves. ## Decisions 1. Moves record new events naming what moved: `MovedStatement`, - `MovedBlock`, `SplicedBlock`, `SplicedCfg`. + `SplicedBlock`, `SplicedCfg`. 2. The walk checks for a moved visited statement or block by reading the `Rewriter`'s event log (step 2). Ancestors are not checked. -3. Moving a block to or from the front of its CFG is allowed, not rejected. +3. Moving a block within its CFG is not added (step 3). 4. A block's statements can be spliced into any block. 5. Splicing leaves the source block or CFG empty and does not erase it. 6. Copying includes graph bodies. From 78babf43306f3ed62459b73496df9dd107b7731e Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Fri, 6 Nov 2026 09:08:24 -0500 Subject: [PATCH 05/14] feat(ir): splice a whole CFG with Rewriter::splice_cfg_after Moves every block of a CFG after an anchor block in another CFG; records SplicedCfg. walk_mut now records a VisitedPlace (block and CFG) before each visit to check splices against. Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/cfg/mod.rs | 2 + crates/kirin-ir/src/rewrite/cfg/splice.rs | 116 ++++++++++ crates/kirin-ir/src/rewrite/rewriter.rs | 14 ++ crates/kirin-ir/src/rewrite/walk/mod.rs | 5 +- crates/kirin-ir/src/rewrite/walk/order.rs | 59 ++++-- crates/kirin-ir/tests/rewriter_cfg.rs | 199 +++++++++++++++++- crates/kirin-ir/tests/walk.rs | 110 ++++++++++ docs/design/rewrite-engine.md | 10 +- .../2026-11-04-remaining-ir-operations.md | 10 +- 9 files changed, 505 insertions(+), 20 deletions(-) create mode 100644 crates/kirin-ir/src/rewrite/cfg/splice.rs diff --git a/crates/kirin-ir/src/rewrite/cfg/mod.rs b/crates/kirin-ir/src/rewrite/cfg/mod.rs index bfadba427..d3aeb845b 100644 --- a/crates/kirin-ir/src/rewrite/cfg/mod.rs +++ b/crates/kirin-ir/src/rewrite/cfg/mod.rs @@ -1,6 +1,7 @@ //! `CFG` block-list surgery for the [`Rewriter`](super::Rewriter). //! //! These methods add, erase, split, and merge blocks of a [`CFG`](crate::CFG), +//! splice one CFG's blocks into another, //! retarget control-flow edges, and insert or erase block arguments together //! with the incoming edges that forward them. Like statement surgery they //! maintain every derived mirror on each successful edit — @@ -31,4 +32,5 @@ mod blocks; mod edges; mod helpers; mod merge; +mod splice; mod split; diff --git a/crates/kirin-ir/src/rewrite/cfg/splice.rs b/crates/kirin-ir/src/rewrite/cfg/splice.rs new file mode 100644 index 000000000..86ee562c9 --- /dev/null +++ b/crates/kirin-ir/src/rewrite/cfg/splice.rs @@ -0,0 +1,116 @@ +use crate::arena::{GetInfo, GetInfoMut}; +use crate::node::linked_list::LinkedList; +use crate::rewrite::walk::tree; +use crate::{Block, BlockParent, CFG, Dialect, StatementParent}; + +use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter}; + +impl<'a, L: Dialect> Rewriter<'a, L> { + /// Move every block of `source`, in order, into `anchor`'s CFG, + /// immediately after `anchor`. + /// + /// The blocks keep their ids, arguments, statements and terminators, and + /// the edges between them stay valid since they all move together. + /// `source` is left empty and is not erased: its owner still names it. + /// `source`'s entry block becomes an ordinary block of the destination, + /// for the caller to branch to. + /// + /// Rejects, leaving the stage untouched, if `source` is unknown + /// ([`RewriteError::UnknownCfg`]), if `anchor` is unknown or not in a CFG + /// ([`RewriteError::NotInCfgBody`]), or if `anchor` is in `source` or + /// nested in one of its blocks ([`RewriteError::CfgSplicedIntoItself`]). + /// + /// Records [`MutationEvent::ChangedBlockList`] for `source`, then for the + /// destination, then [`MutationEvent::SplicedCfg`]. A walk visiting a + /// statement of `source` reports the splice (see [`crate::walk`]). + pub fn splice_cfg_after(&mut self, source: CFG, anchor: Block) -> Result<(), RewriteError> { + if source.get_info(self.stage).is_none() { + return Err(RewriteError::UnknownCfg(source)); + } + let dest = self.owning_cfg(anchor)?; + // Check `dest` and the CFG of every block enclosing it. If one of them + // is `source`, its blocks would be spliced into themselves. + if dest == source { + return Err(RewriteError::CfgSplicedIntoItself { source, anchor }); + } + let mut current = dest.get_info(self.stage).and_then(|info| info.parent); + while let Some(stmt) = current { + if let Some(StatementParent::Block(block)) = *stmt.parent(self.stage) + && self.live_block(block)?.parent == Some(BlockParent::CFG(source)) + { + return Err(RewriteError::CfgSplicedIntoItself { source, anchor }); + } + current = tree::parent(self.stage, stmt); + } + + let chain = match source.get_info_mut(self.stage) { + Some(info) => std::mem::take(&mut info.blocks), + None => LinkedList::new(), + }; + let next = anchor.expect_info(self.stage).node.next; + self.link_blocks(chain, dest, Some(anchor), next); + + self.events + .push(MutationEvent::ChangedBlockList { cfg: source }); + self.events + .push(MutationEvent::ChangedBlockList { cfg: dest }); + self.events.push(MutationEvent::SplicedCfg { source, dest }); + Ok(()) + } + + /// Link `chain`, a whole block chain already taken out of its CFG, into + /// `cfg` between `prev` and `next` (the list head or tail when one is + /// `None`), making `cfg` the parent of each block. + fn link_blocks( + &mut self, + chain: LinkedList, + cfg: CFG, + prev: Option, + next: Option, + ) { + let (Some(head), Some(tail)) = (chain.head, chain.tail) else { + return; + }; + let mut current = Some(head); + while let Some(block) = current { + let Some(info) = block.get_info_mut(self.stage) else { + break; + }; + info.parent = Some(BlockParent::CFG(cfg)); + current = info.node.next; + } + if let Some(info) = head.get_info_mut(self.stage) { + info.node.prev = prev; + } + if let Some(info) = tail.get_info_mut(self.stage) { + info.node.next = next; + } + match prev { + Some(p) => { + if let Some(info) = p.get_info_mut(self.stage) { + info.node.next = Some(head); + } + } + None => { + if let Some(info) = cfg.get_info_mut(self.stage) { + info.blocks.head = Some(head); + } + } + } + match next { + Some(n) => { + if let Some(info) = n.get_info_mut(self.stage) { + info.node.prev = Some(tail); + } + } + None => { + if let Some(info) = cfg.get_info_mut(self.stage) { + info.blocks.tail = Some(tail); + } + } + } + if let Some(info) = cfg.get_info_mut(self.stage) { + info.blocks.len += chain.len; + } + } +} diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index aad764dd7..8114b995f 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -146,6 +146,10 @@ pub enum MutationEvent { /// Every statement of `source`, except its terminator, was moved into /// `dest`, keeping their ids, results, uses and nested bodies. SplicedBlock { source: Block, dest: Block }, + /// Every block of `source` was moved into `dest`, keeping their ids, + /// arguments, statements and edges. Recorded after the + /// [`MutationEvent::ChangedBlockList`] of both CFGs. + SplicedCfg { source: CFG, dest: CFG }, } /// An illegal mutation rejected by a [`Rewriter`]. @@ -242,6 +246,10 @@ pub enum RewriteError { /// `source` or nested in one of its statements, which would place a /// spliced statement next to or inside itself. SplicedIntoItself { source: Block, anchor: Statement }, + /// `splice_cfg_after` was given an anchor block in `source` or nested in + /// one of its blocks, which would place `source`'s blocks inside + /// themselves. + CfgSplicedIntoItself { source: CFG, anchor: Block }, /// The block is not part of a CFG's block list (it is a statement-owned /// single-block body, or detached), so block-list surgery cannot apply. NotInCfgBody(Block), @@ -405,6 +413,12 @@ impl fmt::Display for RewriteError { "cannot splice {source:?} next to {anchor:?}, which is in it" ) } + RewriteError::CfgSplicedIntoItself { source, anchor } => { + write!( + f, + "cannot splice {source:?} after {anchor:?}, which is in it" + ) + } RewriteError::YieldIndexOutOfRange { graph, index } => { write!(f, "yield index {index} is out of range for DiGraph {graph}") } diff --git a/crates/kirin-ir/src/rewrite/walk/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index e62387ef5..931544c76 100644 --- a/crates/kirin-ir/src/rewrite/walk/mod.rs +++ b/crates/kirin-ir/src/rewrite/walk/mod.rs @@ -82,6 +82,9 @@ //! [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) //! or [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after), //! reported the same way; +//! - splicing the blocks of its CFG into another CFG with +//! [`Rewriter::splice_cfg_after`](crate::Rewriter::splice_cfg_after), +//! reported the same way; //! - moving one of its ancestors, which is not checked; //! - reordering the bodies of one of its ancestors //! ([`Rewriter::replace_statement`](crate::Rewriter::replace_statement)), @@ -106,4 +109,4 @@ pub(crate) mod tree; pub use breadth_first::{BFS, BFSCursor, BuiltOrderCursor, KeepOrderReverseBFS, ReverseBFS}; pub use depth_first::{PostOrder, PostOrderCursor, PreOrder, PreOrderCursor}; -pub use order::{WalkCursor, WalkError, WalkOrder, WalkResult, WalkStatus}; +pub use order::{VisitedPlace, WalkCursor, WalkError, WalkOrder, WalkResult, WalkStatus}; diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index d7d3039b2..80bfdcead 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -1,6 +1,10 @@ use std::fmt; -use crate::{Dialect, MutationEvent, Rewriter, StageInfo, Statement, StatementParent}; +use crate::arena::GetInfo; +use crate::{ + Block, BlockParent, CFG, Dialect, MutationEvent, Rewriter, StageInfo, Statement, + StatementParent, +}; /// A traversal order over the statements nested under a root statement. /// @@ -64,9 +68,9 @@ pub trait WalkOrder { let mut cursor = Self::start(rewriter.stage(), root); while let Some(stmt) = cursor.next(rewriter.stage())? { let seen = rewriter.events().len(); - let parent = *stmt.parent(rewriter.stage()); + let place = VisitedPlace::of(rewriter.stage(), stmt); let result = visit(rewriter, stmt); - Self::check_not_moved_since(rewriter, stmt, parent, seen)?; + Self::check_not_moved_since(rewriter, stmt, place, seen)?; if let Some(status) = result.steer(&mut cursor) { return Ok(status); } @@ -78,19 +82,17 @@ pub trait WalkOrder { /// first `seen` events of `rewriter` moved `visited`, in an order that /// does not allow it /// ([`ALLOWS_DISPLACING_VISITED`](Self::ALLOWS_DISPLACING_VISITED)): - /// a move of `visited` itself, or a splice of `visited_parent`'s block. - /// Panics if `rewriter` has recorded fewer than `seen` events. + /// a move of `visited` itself, or a splice of the block or CFG it was in + /// at `place`. Panics if `rewriter` has recorded fewer than `seen` events. /// /// [`walk_mut`](Self::walk_mut) calls this after each visit. A loop driving /// a cursor by hand calls it too: note `rewriter.events().len()` and - /// `visited.parent(stage)` before editing at a statement, and check before - /// the next [`WalkCursor::next`]. The parent is needed because a splice - /// changes it: afterwards, a statement spliced into a block cannot be told - /// from one that was there already. + /// [`VisitedPlace::of`] before editing at a statement, and check before + /// the next [`WalkCursor::next`]. fn check_not_moved_since( rewriter: &Rewriter<'_, L>, visited: Statement, - visited_parent: Option, + place: VisitedPlace, seen: usize, ) -> Result<(), WalkError> { if Self::ALLOWS_DISPLACING_VISITED { @@ -98,9 +100,8 @@ pub trait WalkOrder { } let moved = rewriter.events()[seen..].iter().any(|event| match *event { MutationEvent::MovedStatement { stmt } => stmt == visited, - MutationEvent::SplicedBlock { source, .. } => { - visited_parent == Some(StatementParent::Block(source)) - } + MutationEvent::SplicedBlock { source, .. } => place.block == Some(source), + MutationEvent::SplicedCfg { source, .. } => place.cfg == Some(source), _ => false, }); if moved { @@ -111,6 +112,32 @@ pub trait WalkOrder { } } +/// Where a visited statement was when its visit started: its block, and that +/// block's CFG. [`WalkOrder::check_not_moved_since`] compares splices against +/// it, since a splice changes both: afterwards, a statement spliced into a +/// block, or a block spliced into a CFG, cannot be told from one that was +/// there already. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct VisitedPlace { + block: Option, + cfg: Option, +} + +impl VisitedPlace { + /// Where `stmt` is now in `stage`. + pub fn of(stage: &StageInfo, stmt: Statement) -> Self { + let block = match stmt.get_info(stage).and_then(|info| info.parent) { + Some(StatementParent::Block(block)) => Some(block), + _ => None, + }; + let cfg = match block.and_then(|block| block.get_info(stage)?.parent()) { + Some(BlockParent::CFG(cfg)) => Some(cfg), + _ => None, + }; + Self { block, cfg } + } +} + /// What a walk does after visiting a statement. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] pub enum WalkResult { @@ -181,8 +208,10 @@ pub enum WalkError { /// [`Rewriter::move_after`](crate::Rewriter::move_after), or by splicing /// its block with /// [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) / - /// [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after) - /// (see [`WalkOrder::check_not_moved_since`]). + /// [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after), + /// or by splicing its CFG with + /// [`Rewriter::splice_cfg_after`](crate::Rewriter::splice_cfg_after) (see + /// [`WalkOrder::check_not_moved_since`]). VisitedMoved(Statement), } diff --git a/crates/kirin-ir/tests/rewriter_cfg.rs b/crates/kirin-ir/tests/rewriter_cfg.rs index bcfc213f9..1940c326a 100644 --- a/crates/kirin-ir/tests/rewriter_cfg.rs +++ b/crates/kirin-ir/tests/rewriter_cfg.rs @@ -70,7 +70,8 @@ fn append_and_insert_blocks_link_into_cfg() { "one ChangedBlockList per created block" ); - // The block list reads the same forwards and backwards. + // Following the `next` links from the head and the `prev` links from the + // tail visits the same blocks, in opposite orders. let order: Vec = cfg.blocks(&stage).collect(); assert_eq!(order, vec![before, entry, after, exit, tail]); let reversed: Vec = cfg.blocks(&stage).rev().collect(); @@ -1251,3 +1252,199 @@ fn block_arg_edits_reject_edges_without_successor_operands() { assert_eq!(target.expect_info(rw.stage()).arguments(), args_before); verify_derived(rw.stage()).unwrap(); } + +/// ```text +/// caller { ^b { ret } ^c { ret } } +/// callee { ^e(%a) { u = use %a; br ^f } ^f { inner_owner = own_cfg inner { ^i { ret } }; ret } } +/// ^lone { } // in no CFG +/// ``` +struct CfgSplice { + stage: StageInfo, + caller: CFG, + callee: CFG, + b: Block, + c: Block, + e: Block, + f: Block, + i: Block, + lone: Block, + a: SSAValue, + u: Statement, +} + +fn cfg_splice() -> CfgSplice { + let mut stage = new_stage(); + let b_ret = stage.statement().definition(BuilderDialect::Return).new(); + let b = stage.block().terminator(b_ret).new(); + let c_ret = stage.statement().definition(BuilderDialect::Return).new(); + let c = stage.block().terminator(c_ret).new(); + let caller = stage.cfg().add_block(b).add_block(c).new(); + + let i_ret = stage.statement().definition(BuilderDialect::Return).new(); + let i = stage.block().terminator(i_ret).new(); + let inner = stage.cfg().add_block(i).new(); + let inner_owner = stage + .statement() + .definition(BuilderDialect::OwnCFG(inner)) + .new(); + let f_ret = stage.statement().definition(BuilderDialect::Return).new(); + let f = stage.block().stmt(inner_owner).terminator(f_ret).new(); + let a = stage.block_argument().index(0); + let u = stage.statement().definition(BuilderDialect::Use(a)).new(); + let e_br = stage.statement().definition(br(f)).new(); + let e = stage + .block() + .argument(TestType::I32) + .stmt(u) + .terminator(e_br) + .new(); + let callee = stage.cfg().add_block(e).add_block(f).new(); + let lone = stage.block().new(); + + let stage = stage.finalize().unwrap(); + let a = arg(&stage, e, 0); + CfgSplice { + stage, + caller, + callee, + b, + c, + e, + f, + i, + lone, + a, + u, + } +} + +/// `cfg`'s blocks in order, checking that following the `prev` links from +/// the tail visits the same blocks as following the `next` links from the head. +fn block_order(stage: &StageInfo, cfg: CFG) -> Vec { + let forward: Vec = cfg.blocks(stage).collect(); + let mut backward: Vec = cfg.blocks(stage).rev().collect(); + backward.reverse(); + assert_eq!(forward, backward, "prev and next links agree"); + forward +} + +#[test] +fn splice_cfg_after_moves_every_block_and_keeps_their_edges() { + let CfgSplice { + mut stage, + caller, + callee, + b, + c, + e, + f, + a, + u, + .. + } = cfg_splice(); + let mut rw = Rewriter::new(&mut stage); + rw.splice_cfg_after(callee, b).unwrap(); + + assert_eq!(block_order(rw.stage(), caller), [b, e, f, c]); + assert_eq!(block_order(rw.stage(), callee), []); + for block in [e, f] { + assert_eq!( + block.expect_info(rw.stage()).parent(), + Some(BlockParent::CFG(caller)) + ); + } + // Arguments, uses and edges move with the blocks. + assert_eq!(arg(rw.stage(), e, 0), a); + assert_eq!(uses_set(rw.stage(), a), HashSet::from([so(u, 0)])); + assert_eq!(preds(rw.stage(), f), HashSet::from([e])); + assert_eq!( + rw.events().to_vec(), + vec![ + MutationEvent::ChangedBlockList { cfg: callee }, + MutationEvent::ChangedBlockList { cfg: caller }, + MutationEvent::SplicedCfg { + source: callee, + dest: caller, + }, + ] + ); + verify_derived(rw.stage()).unwrap(); +} + +#[test] +fn splice_cfg_after_the_last_block_and_from_an_empty_cfg() { + let CfgSplice { + mut stage, + caller, + callee, + b, + c, + e, + f, + .. + } = cfg_splice(); + let mut rw = Rewriter::new(&mut stage); + // After the last block: the moved blocks end the list. + rw.splice_cfg_after(callee, c).unwrap(); + assert_eq!(block_order(rw.stage(), caller), [b, c, e, f]); + + // `callee` is now empty: splicing it again moves nothing. + let mut log = EventCursor::default(); + log.advance(&rw); + rw.splice_cfg_after(callee, b).unwrap(); + assert_eq!(block_order(rw.stage(), caller), [b, c, e, f]); + assert_eq!(log.read_and_advance(&rw).len(), 3); + verify_derived(rw.stage()).unwrap(); +} + +#[test] +fn splice_cfg_rejects_illegal_splices_leaving_the_stage_untouched() { + let CfgSplice { + mut stage, + caller, + callee, + b, + c, + e, + f, + i, + lone, + .. + } = cfg_splice(); + let dead_cfg = CFG::from(Id::from(TestSSAValue(9_999))); + let dead_block = Block::from(Id::from(TestSSAValue(9_999))); + let mut rw = Rewriter::new(&mut stage); + + assert_eq!( + rw.splice_cfg_after(dead_cfg, b), + Err(RewriteError::UnknownCfg(dead_cfg)) + ); + assert_eq!( + rw.splice_cfg_after(callee, dead_block), + Err(RewriteError::UnknownBlock(dead_block)) + ); + assert_eq!( + rw.splice_cfg_after(callee, lone), + Err(RewriteError::NotInCfgBody(lone)) + ); + // The anchor is in the source. + assert_eq!( + rw.splice_cfg_after(callee, e), + Err(RewriteError::CfgSplicedIntoItself { + source: callee, + anchor: e + }) + ); + // The anchor is nested in one of the source's blocks (`^f` owns `inner`). + assert_eq!( + rw.splice_cfg_after(callee, i), + Err(RewriteError::CfgSplicedIntoItself { + source: callee, + anchor: i + }) + ); + assert!(rw.events().is_empty()); + assert_eq!(block_order(rw.stage(), caller), [b, c]); + assert_eq!(block_order(rw.stage(), callee), [e, f]); + verify_derived(rw.stage()).unwrap(); +} diff --git a/crates/kirin-ir/tests/walk.rs b/crates/kirin-ir/tests/walk.rs index ee7d9da81..da0d84be6 100644 --- a/crates/kirin-ir/tests/walk.rs +++ b/crates/kirin-ir/tests/walk.rs @@ -927,6 +927,116 @@ fn splicing_another_block_into_the_visited_statements_block_keeps_the_walk() { assert_eq!(splice_while_visiting_a::(), spliced_schedule::()); } +/// ```text +/// root = own_cfg outer { +/// ^o { p = own_cfg inner { ^i { q; ret } }; ret } +/// } +/// ``` +#[derive(Clone, Copy)] +struct NestedCfg { + root: Statement, + o: Block, + p: Statement, + q: Statement, + inner: CFG, +} + +fn nested_cfg() -> (StageInfo, NestedCfg) { + let mut stage = new_stage(); + let q = stage.statement().definition(BuilderDialect::Nop).new(); + let i_ret = stage.statement().definition(BuilderDialect::Return).new(); + let i = stage.block().stmt(q).terminator(i_ret).new(); + let inner = stage.cfg().add_block(i).new(); + let p = stage + .statement() + .definition(BuilderDialect::OwnCFG(inner)) + .new(); + let o_ret = stage.statement().definition(BuilderDialect::Return).new(); + let o = stage.block().stmt(p).terminator(o_ret).new(); + let outer = stage.cfg().add_block(o).new(); + let root = stage + .statement() + .definition(BuilderDialect::OwnCFG(outer)) + .new(); + let ids = NestedCfg { + root, + o, + p, + q, + inner, + }; + (stage.finalize().unwrap(), ids) +} + +/// Splices `inner` after `^o`, into `outer`. +fn splice_inner_after_o(rewriter: &mut Rewriter, n: NestedCfg) { + rewriter.splice_cfg_after(n.inner, n.o).unwrap(); +} + +/// Walk `O` over [`nested_cfg`] with `walk_mut`, splicing `inner` into +/// `outer` when visiting `at`; the statements visited. +fn splice_inner_while_visiting( + at: impl Fn(NestedCfg) -> Statement, +) -> Result, WalkError> { + let (mut stage, n) = nested_cfg(); + let mut rewriter = Rewriter::new(&mut stage); + let mut visited = Vec::new(); + O::walk_mut(&mut rewriter, n.root, |rewriter, stmt| { + visited.push(stmt); + if stmt == at(n) { + splice_inner_after_o(rewriter, n); + } + WalkResult::Advance + }) + .map(|_| visited) +} + +#[test] +fn splicing_the_visited_statements_cfg_breaks_link_following_walks() { + let (_, n) = nested_cfg(); + for result in [ + splice_inner_while_visiting::(|n| n.q), + splice_inner_while_visiting::(|n| n.q), + splice_inner_while_visiting::(|n| n.q), + ] { + assert_eq!(result, Err(WalkError::VisitedMoved(n.q))); + } +} + +#[test] +fn reverse_bfs_orders_go_on_after_the_visited_statements_cfg_is_spliced() { + let (stage, n) = nested_cfg(); + assert_eq!( + splice_inner_while_visiting::(|n| n.q).unwrap(), + ReverseBFS::schedule(&stage, n.root) + ); + assert_eq!( + splice_inner_while_visiting::(|n| n.q).unwrap(), + KeepOrderReverseBFS::schedule(&stage, n.root) + ); +} + +#[test] +fn splicing_a_cfg_into_the_visited_statements_cfg_keeps_the_walk() { + /// The schedule of `O` over [`nested_cfg`] with the splice already made. + fn spliced_schedule() -> Vec { + let (mut stage, n) = nested_cfg(); + splice_inner_after_o(&mut Rewriter::new(&mut stage), n); + O::schedule(&stage, n.root) + } + // While visiting `p`, `inner`'s block moves into `outer`, after `^o`; + // the walk visits it there. Not `PostOrder`: it visits `p` after `q`, so + // the splice would put `q` ahead of the cursor a second time. + assert_eq!( + splice_inner_while_visiting::(|n| n.p).unwrap(), + spliced_schedule::() + ); + assert_eq!( + splice_inner_while_visiting::(|n| n.p).unwrap(), + spliced_schedule::() + ); +} + /// Splits the entry block at `a` while visiting `a`; the walk's schedule /// before the split, and the statements it visited. fn split_at_a_while_visiting() -> (Vec, Result, WalkError>) diff --git a/docs/design/rewrite-engine.md b/docs/design/rewrite-engine.md index 45332fcdb..83162f757 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -1116,6 +1116,7 @@ pub enum MutationEvent { ChangedBodyNesting { owner: Statement }, MovedStatement { stmt: Statement }, SplicedBlock { source: Block, dest: Block }, + SplicedCfg { source: Cfg, dest: Cfg }, } ``` @@ -1767,6 +1768,8 @@ impl<'a, L: Dialect> Rewriter<'a, L> { pub fn split_block(&mut self, at: Statement) -> Result; + pub fn splice_cfg_after(&mut self, source: CFG, anchor: Block) -> Result<(), RewriteError>; + pub fn merge_blocks( &mut self, pred: Block, @@ -1825,6 +1828,7 @@ stage remains valid. | `append_block` / `insert_block_before` / `insert_block_after` | Allocate an empty block (no terminator) with fresh typed arguments and link it into the `Cfg` block list. | `ChangedBlockList` | anchor must be a live block of a `Cfg` list (not a statement-owned body) | | `erase_block` | Unlink the block and tombstone it with its arguments, statements, results, and nested blocks/`Cfg`s; drop their uses and outgoing edges. | `ChangedBlockList`, `ChangedCfg` | rejects if another block still branches to it (a self-loop is fine) or an argument or top-level result is used outside the block; retarget those edges and replace those uses first | | `split_block` | Move `at`, its successors in the chain, and the terminator into a new block placed after the original; the original is left without a terminator, for the caller to close with `insert_terminator`. | `ChangedBlockList`, `ChangedCfg` | new block has no arguments | +| `splice_cfg_after` | Move every block of `source`, in order, into `anchor`'s `Cfg` right after `anchor`; ids, arguments, statements, terminators and the edges between them are unchanged. `source` is left empty and not erased. | `ChangedBlockList` for `source`, then for the destination, then `SplicedCfg` | `source` live; `anchor` a live block of a `Cfg` list, not in `source` and not nested in its blocks; an empty `source` moves nothing | | `merge_blocks` | Erase `pred`'s terminator, redirect `succ`'s arguments to `arg_replacements`, splice `succ`'s statements and terminator onto `pred`, tombstone `succ`. | the `erase_statement` events, the `replace_all_uses` events, `ChangedBlockList`, `ChangedCfg` | same `Cfg`, distinct blocks, exact argument arity; replacements must be live and not `succ`'s own arguments; rejects if a block other than `pred` still branches to `succ`; `pred` need not branch to `succ` | | `retarget_successor` | Point every successor slot of `stmt` naming `old` at `new`. | `ChangedCfg` (none if nothing changed) | returns the count rewritten; `old` need not be live, `new` must be live | | `insert_block_arg` / `erase_block_arg` | Insert or erase one block argument, renumbering later ones, and insert or remove the matching operand on each incoming edge through `BranchOp`. `insert_block_arg` asks `forwarded(stmt, successor)` for each edge's new operand. | `ReplacedStatement` and `ChangedTerminator` per edge, then `ChangedCfg` | every rewritten edge is checked before any edit; an erased argument must be unused; every incoming edge needs an editable operand list (`NoSuccessorOperands` otherwise) | @@ -1990,7 +1994,7 @@ axis — hence the column. | `UnknownStatement(stmt)` | liveness | id does not resolve to a live (non-tombstoned) statement | all statement methods | | `UnknownValue(value)` | liveness | SSA id does not resolve to a live value | `replace_operand`, `replace_all_uses`, `insert_*`, `replace_statement`, `splice_block_*` | | `UnknownBlock(block)` | liveness | a block id (anchor, operand block, or successor target) does not resolve to a live block | `replace_statement`, CFG block-list methods | -| `UnknownCfg(cfg)` | liveness | the `Cfg` id does not resolve to a live `Cfg` | `append_block` | +| `UnknownCfg(cfg)` | liveness | the `Cfg` id does not resolve to a live `Cfg` | `append_block`, `splice_cfg_after` | | `UnknownDiGraph(graph)` | liveness | the `DiGraph` id does not resolve to a live `DiGraph` | `insert_node`, `replace_yield` | | `YieldIndexOutOfRange { graph, index }` | liveness | yield index past the graph's yield list | `replace_yield` | | `UnknownUnGraph(graph)` | liveness | the `UnGraph` id does not resolve to a live `UnGraph` | `insert_ungraph_node`, `insert_ungraph_edge` | @@ -2025,6 +2029,7 @@ axis — hence the column. | `BodyListedTwice(body)` | guard | the replacement names the same body in two fields | `replace_statement` | | `BodyEnclosesOwner(body)` | guard | the replacement names a body the statement is nested in, which would make it its own ancestor | `replace_statement` | | `MovedIntoItself { stmt, anchor }` | guard | the anchor is nested in the moved statement, which would make it its own ancestor | `move_before`, `move_after` | +| `CfgSplicedIntoItself { source, anchor }` | guard | the anchor block is in `source` or nested in one of its blocks, which would place `source`'s blocks inside themselves | `splice_cfg_after` | | `SplicedIntoItself { source, anchor }` | guard | the anchor is in `source` or nested in one of its statements, which would place a spliced statement next to or inside itself | `splice_block_before`, `splice_block_after` | The immediate `MutationEvent` vocabulary recorded by these methods is @@ -2032,7 +2037,8 @@ The immediate `MutationEvent` vocabulary recorded by these methods is `ReplacedStatement`, `ChangedBlockList { cfg }`, `ChangedCfg { cfg }`, `ChangedTerminator { block }`, `ChangedDiGraphTopology { graph }`, `ChangedDiGraphYields { graph }`, `ChangedBodyNesting { owner }`, -`MovedStatement { stmt }`, and `SplicedBlock { source, dest }`. `ErasedStatement` also covers everything nested +`MovedStatement { stmt }`, `SplicedBlock { source, dest }`, and +`SplicedCfg { source, dest }`. `ErasedStatement` also covers everything nested in the erased statement; no separate CFG or graph events are recorded for those bodies. The remaining documented event kinds (`UnGraph` and result-type events) grow in as diff --git a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md index a4e3149f3..981b05017 100644 --- a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md +++ b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md @@ -184,7 +184,15 @@ since they all move together. `source` is left empty and not erased; its owner (e.g. the callee function) still names it, and the caller decides whether to erase the owner. Records `ChangedBlockList` for both CFGs and a new `SplicedCfg { source, dest }`, which the walk check reads as "every block of -`source` moved". +`source` moved". Afterwards a moved block's parent is `dest`, so, as in step 4, +the check compares against where the visited statement was before the visit: +a `VisitedPlace` holding its block and that block's CFG. +`ChangedBlockList` alone could not drive the check, since `split_block`, which +a walk allows, records it too. + +Rejects an anchor block in `source` or nested in one of its blocks +(`CfgSplicedIntoItself`). There is no `_before` variant and no way to splice +into an empty CFG, which has no anchor; inlining needs neither. Cost: O(blocks moved). This is the move-instead-of-copy path for inlining a callee's last use. From 9fa436476c23fa0c476afdd1c07af3a1da60caa3 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 11:34:42 -0500 Subject: [PATCH 06/14] fix(ir): skip moves that leave a statement where it already is Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/moves.rs | 10 +++++++--- crates/kirin-ir/tests/rewriter.rs | 4 ++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/moves.rs b/crates/kirin-ir/src/rewrite/moves.rs index 5eced7ea9..b2e711d64 100644 --- a/crates/kirin-ir/src/rewrite/moves.rs +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -11,7 +11,9 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// The statement keeps its id, results, uses and nested bodies; only its /// place changes. As with [`Rewriter::insert_before`], `anchor` may be a /// block's terminator: `stmt` then goes to the end of that block's - /// statement chain. Moving a statement before itself does nothing. + /// statement chain. Moving a statement to where it already is (before + /// itself, or before the statement that already follows it) does nothing + /// and records no event. /// /// Not checked, as with [`Rewriter::insert_before`]: that every use of /// `stmt`'s results still comes after it. @@ -54,8 +56,10 @@ impl<'a, L: Dialect> Rewriter<'a, L> { }; // Checks the anchor; the place itself is read again once `stmt` is // unlinked, since `stmt` may be one of its neighbours. - self.slot_at(anchor, after)?; - if stmt == anchor { + let (_, prev, next) = self.slot_at(anchor, after)?; + // `stmt` already borders the place (this includes `stmt == anchor`), + // so unlinking it and linking it back would put it where it is. + if prev == Some(stmt) || next == Some(stmt) { return Ok(()); } // Walk up the statements whose bodies enclose `anchor`. If `stmt` is diff --git a/crates/kirin-ir/tests/rewriter.rs b/crates/kirin-ir/tests/rewriter.rs index 1917971a1..5777d9024 100644 --- a/crates/kirin-ir/tests/rewriter.rs +++ b/crates/kirin-ir/tests/rewriter.rs @@ -1558,11 +1558,15 @@ fn moving_a_statement_next_to_itself_does_nothing() { q, r, owner, + ret, .. } = move_fixture(); let mut rewriter = Rewriter::new(&mut stage); rewriter.move_before(q, q).unwrap(); rewriter.move_after(q, q).unwrap(); + rewriter.move_before(p, q).unwrap(); + rewriter.move_after(r, q).unwrap(); + rewriter.move_before(owner, ret).unwrap(); assert!(rewriter.events().is_empty()); assert_eq!(order(rewriter.stage(), entry), [p, q, r, owner]); } From 2f69cc4139d83e0155cb577ecd4bf5f2d712b917 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 13:49:23 -0500 Subject: [PATCH 07/14] fix(walk): check every walk contract break, in walks and the driver - run the rewrite driver's sweep through walk_mut - check erases and merges of the visited statement on interrupt too - no longer exempt the walk root from the move checks - make walk cursors crate-private - rename ALLOWS_DISPLACING_VISITED to ALLOWS_ERASING_OR_MOVING_VISITED Co-Authored-By: Claude Opus 5.5 --- .../src/rewrite/walk/breadth_first.rs | 128 +++++++----- .../kirin-ir/src/rewrite/walk/depth_first.rs | 52 +++-- crates/kirin-ir/src/rewrite/walk/mod.rs | 29 +-- crates/kirin-ir/src/rewrite/walk/order.rs | 137 +++++++------ crates/kirin-ir/tests/walk.rs | 182 +++++++++--------- .../src/driver/rewrite_driver.rs | 80 +++++--- crates/kirin-rewrite/tests/walk_result.rs | 63 +++++- docs/design/rewrite-engine.md | 4 +- 8 files changed, 420 insertions(+), 255 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs index acfc6d4b4..4feaf0e28 100644 --- a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs @@ -9,14 +9,33 @@ use std::collections::VecDeque; use super::depth_first::{check_not_moved, live, visited_block}; use super::order::Position; +use super::order::{CursorOrder, WalkCursor}; use super::tree::{append_children, first_child, next_sibling, owns_bodies}; -use super::{WalkCursor, WalkError, WalkOrder}; +use super::{WalkError, WalkOrder}; use crate::{Block, Dialect, GetInfo, StageInfo, Statement}; /// Breadth-first by nesting depth: the root, then every statement nested /// directly in it, then every statement nested two levels down, and so on. /// Each level is in program order. -/// See [`BFSCursor`] for how edits affect discovery during a walk. +/// +/// When the walk moves on from a visited statement, it queues that statement +/// as a parent for the next level only if it owns bodies and was not skipped. +/// At the next level, it walks each queued parent's bodies in turn, reading +/// their current links. Inserts ahead of the walk in those bodies are +/// visited. +/// +/// The walk does not revisit bodies it has passed or reconsider statements +/// it did not queue as parents. For example, if a leaf gains a body after the +/// walk has moved on from it, that body and its descendants are not visited. +/// Likewise, a subtree inserted behind the walk is not discovered, even if +/// its descendants lie at depths the walk has not reached. Such edits are +/// allowed by the walk contract; visiting them requires another walk. +/// +/// A visit may erase an ancestor of the visited statement, and with it the +/// visited statement and its siblings: the walk goes on with the next queued +/// parent outside the erased subtree. Erasing the visited statement alone is +/// an error, because its siblings are still live but can no longer be +/// reached. #[derive(Clone, Copy, Debug, Default)] pub struct BFS; @@ -32,28 +51,9 @@ pub struct ReverseBFS; #[derive(Clone, Copy, Debug, Default)] pub struct KeepOrderReverseBFS; -/// A [`BFS`] walk in progress. -/// -/// When the cursor moves on from a visited statement, it queues that statement -/// as a parent for the next level only if it owns bodies and was not skipped. -/// At the next level, the cursor walks each queued parent's bodies in turn, -/// reading their current links. Inserts ahead of the cursor in those bodies -/// are visited. -/// -/// The cursor does not revisit bodies it has passed or reconsider statements -/// it did not queue as parents. For example, if a leaf gains a body after the -/// cursor has moved on from it, that body and its descendants are not visited. -/// Likewise, a subtree inserted behind the cursor is not discovered, even if -/// its descendants lie at depths the walk has not reached. Such edits are -/// allowed by the walk contract; visiting them requires another walk. -/// -/// A visit may erase an ancestor of the visited statement, and with it the -/// visited statement and its siblings: the walk goes on with the next queued -/// parent outside the erased subtree. Erasing the visited statement alone is -/// an error, because its siblings are still live but can no longer be -/// reached. +/// A [`BFS`] walk in progress; see [`BFS`] for how edits affect discovery. #[derive(Clone, Debug)] -pub struct BFSCursor { +pub(crate) struct BFSCursor { root: Statement, position: Position, /// The visited statements that own bodies, oldest first. The front is the @@ -73,15 +73,18 @@ pub struct BFSCursor { /// visited statement and its ancestors: the cursor never reads the visited /// statement's links, so a visit may erase them. #[derive(Clone, Debug)] -pub struct BuiltOrderCursor { +pub(crate) struct BuiltOrderCursor { ordered_nodes: Vec, /// The index of the next statement to consider. next: usize, } impl WalkOrder for BFS { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = false; +} + +impl CursorOrder for BFS { type Cursor = BFSCursor; - const ALLOWS_DISPLACING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { BFSCursor { @@ -95,13 +98,14 @@ impl WalkOrder for BFS { impl WalkCursor for BFSCursor { fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + self.check_visited(stage)?; let next = match self.position { Position::Start => live(stage, self.root), Position::At(visited) | Position::Skipped(visited) => { + // Erased with an ancestor, which `check_visited` allowed. if visited.get_info(stage).is_none() { - return self.after_erased(stage, visited); + return Ok(self.after_parent_erased(stage)); } - check_not_moved(stage, visited, self.visited_block)?; // Queue it as a parent for the next level, whose statements // are its children. Statements without bodies have none, // and a skipped statement's are not walked. @@ -128,39 +132,53 @@ impl WalkCursor for BFSCursor { Position::Done => None, }; self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, self.root, next); + self.visited_block = visited_block(stage, next); Ok(next) } - fn skip(&mut self) { - self.position.skip(); - } -} - -impl BFSCursor { - /// The next statement after the visit erased `visited`. Allowed only if - /// it erased an ancestor too, killing the front parent: its remaining - /// children are then dead, and the walk goes on with the next live - /// parent's. - fn after_erased( - &mut self, - stage: &StageInfo, - visited: Statement, - ) -> Result, WalkError> { + /// A live visited statement must still be in its block: merging that + /// block into another moves it, and the walk would go on from its new + /// place. Moves and splices leave its block alive, so only + /// [`WalkOrder::walk_mut`](super::WalkOrder::walk_mut) catches them. + /// + /// An erased one is allowed only if erased together with an ancestor, + /// which kills the front parent: its remaining children are then dead + /// too. Erasing it alone would leave its siblings live but unreachable. + fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { + let Some(visited) = self.position.visited() else { + return Ok(()); + }; + if visited.get_info(stage).is_some() { + return check_not_moved(stage, visited, self.visited_block); + } let parent_erased = visited != self.root && self .parents .front() .is_some_and(|&parent| parent.get_info(stage).is_none()); - if !parent_erased { - return Err(WalkError::VisitedErased(visited)); + if parent_erased { + Ok(()) + } else { + Err(WalkError::VisitedErased(visited)) } + } + + fn skip(&mut self) { + self.position.skip(); + } +} + +impl BFSCursor { + /// The next statement after the visit erased the visited statement with + /// its front parent: the walk goes on with the next live parent's + /// children. + fn after_parent_erased(&mut self, stage: &StageInfo) -> Option { // Dead parents, the front one included, have no children, so this // drops them. let next = self.first_child_of_next_parent(stage); self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, self.root, next); - Ok(next) + self.visited_block = visited_block(stage, next); + next } /// The first child of the front parent or a later one, dropping parents @@ -182,8 +200,11 @@ impl BFSCursor { } impl WalkOrder for ReverseBFS { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = true; +} + +impl CursorOrder for ReverseBFS { type Cursor = BuiltOrderCursor; - const ALLOWS_DISPLACING_VISITED: bool = true; fn start(stage: &StageInfo, root: Statement) -> Self::Cursor { let mut ordered_nodes = breadth_first(stage, root, |_| {}); @@ -193,8 +214,11 @@ impl WalkOrder for ReverseBFS { } impl WalkOrder for KeepOrderReverseBFS { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = true; +} + +impl CursorOrder for KeepOrderReverseBFS { type Cursor = BuiltOrderCursor; - const ALLOWS_DISPLACING_VISITED: bool = true; fn start(stage: &StageInfo, root: Statement) -> Self::Cursor { let mut level_starts = Vec::new(); @@ -234,6 +258,12 @@ impl WalkCursor for BuiltOrderCursor { Ok(None) } + /// Nothing to check: the order was built at the start, so a visit may + /// erase or move the visited statement. + fn check_visited(&self, _stage: &StageInfo) -> Result<(), WalkError> { + Ok(()) + } + /// Nothing to skip: the order visits a statement after the statements /// nested in it. fn skip(&mut self) {} diff --git a/crates/kirin-ir/src/rewrite/walk/depth_first.rs b/crates/kirin-ir/src/rewrite/walk/depth_first.rs index c294bdb5e..a8e5078f8 100644 --- a/crates/kirin-ir/src/rewrite/walk/depth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/depth_first.rs @@ -1,8 +1,9 @@ //! Depth-first orders, walked through the IR's own links in O(1) extra space. use super::order::Position; +use super::order::{CursorOrder, WalkCursor}; use super::tree::{first_child, next_sibling, parent}; -use super::{WalkCursor, WalkError, WalkOrder}; +use super::{WalkError, WalkOrder}; use crate::{Block, Dialect, GetInfo, StageInfo, Statement, StatementParent}; /// Depth-first, each statement before the statements nested in it. @@ -15,7 +16,7 @@ pub struct PostOrder; /// A [`PreOrder`] walk in progress. #[derive(Clone, Copy, Debug)] -pub struct PreOrderCursor { +pub(crate) struct PreOrderCursor { root: Statement, position: Position, /// The block the visited statement was in when visited. @@ -24,7 +25,7 @@ pub struct PreOrderCursor { /// A [`PostOrder`] walk in progress. #[derive(Clone, Copy, Debug)] -pub struct PostOrderCursor { +pub(crate) struct PostOrderCursor { root: Statement, position: Position, /// The block the visited statement was in when visited. @@ -32,8 +33,11 @@ pub struct PostOrderCursor { } impl WalkOrder for PreOrder { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = false; +} + +impl CursorOrder for PreOrder { type Cursor = PreOrderCursor; - const ALLOWS_DISPLACING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { PreOrderCursor { @@ -45,8 +49,11 @@ impl WalkOrder for PreOrder { } impl WalkOrder for PostOrder { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = false; +} + +impl CursorOrder for PostOrder { type Cursor = PostOrderCursor; - const ALLOWS_DISPLACING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { PostOrderCursor { @@ -59,25 +66,29 @@ impl WalkOrder for PostOrder { impl WalkCursor for PreOrderCursor { fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + self.check_visited(stage)?; let next = match self.position { Position::Start => live(stage, self.root), Position::At(visited) => { - check_in_place(stage, visited, self.visited_block)?; // The children are read now, so edits made while visiting // `visited` are walked. first_child(stage, visited).or_else(|| after_subtree(stage, self.root, visited)) } - Position::Skipped(visited) => { - check_in_place(stage, visited, self.visited_block)?; - after_subtree(stage, self.root, visited) - } + Position::Skipped(visited) => after_subtree(stage, self.root, visited), Position::Done => None, }; self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, self.root, next); + self.visited_block = visited_block(stage, next); Ok(next) } + fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { + match self.position.visited() { + Some(visited) => check_in_place(stage, visited, self.visited_block), + None => Ok(()), + } + } + fn skip(&mut self) { self.position.skip(); } @@ -85,11 +96,11 @@ impl WalkCursor for PreOrderCursor { impl WalkCursor for PostOrderCursor { fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + self.check_visited(stage)?; let next = match self.position { Position::Start => live(stage, self.root).map(|root| deepest_first(stage, root)), // Never `Skipped`: `skip` does nothing in this order. Position::At(visited) | Position::Skipped(visited) => { - check_in_place(stage, visited, self.visited_block)?; if visited == self.root { None } else if let Some(sibling) = next_sibling(stage, visited) { @@ -101,10 +112,17 @@ impl WalkCursor for PostOrderCursor { Position::Done => None, }; self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, self.root, next); + self.visited_block = visited_block(stage, next); Ok(next) } + fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { + match self.position.visited() { + Some(visited) => check_in_place(stage, visited, self.visited_block), + None => Ok(()), + } + } + /// Nothing to skip: the children were visited before the statement. fn skip(&mut self) {} } @@ -147,16 +165,12 @@ pub(super) fn check_not_moved( } } -/// The block `next` is in, for [`check_not_moved`] at the next step. `None` -/// for the root, whose place outside its subtree the walk never reads, so -/// moving it is harmless. +/// The block `next` is in, for [`check_not_moved`] at the next step. pub(super) fn visited_block( stage: &StageInfo, - root: Statement, next: Option, ) -> Option { - let stmt = next.filter(|&stmt| stmt != root)?; - match stmt.get_info(stage)?.parent? { + match next?.get_info(stage)?.parent? { StatementParent::Block(block) => Some(block), StatementParent::DiGraph(_) | StatementParent::UnGraph(_) => None, } diff --git a/crates/kirin-ir/src/rewrite/walk/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index 931544c76..2a00d6e46 100644 --- a/crates/kirin-ir/src/rewrite/walk/mod.rs +++ b/crates/kirin-ir/src/rewrite/walk/mod.rs @@ -12,8 +12,8 @@ //! id order. Rewrites whose result depends on visit order therefore depend //! on these orders across blocks and graph members. //! -//! A walk is a [`WalkCursor`] stepped one statement at a time, so the IR can -//! be edited at each statement before the walk moves on: +//! A walk visits one statement at a time, so the IR can be edited at each +//! statement before the walk moves on: //! //! - [`PreOrder`] and [`PostOrder`] follow the IR's own links in O(1) extra //! space, reading them at every step, so they visit statements inserted @@ -27,7 +27,7 @@ //! as their bodies are walked. Inserts ahead of the cursor in those //! bodies are visited; edits to bodies already passed, or to statements //! not queued as parents, may be missed together with their descendants, -//! even at depths the walk has not reached. See [`BFSCursor`] for details. +//! even at depths the walk has not reached. See [`BFS`] for details. //! - [`ReverseBFS`] (each level right to left) and [`KeepOrderReverseBFS`] //! (each level left to right) build their order when the walk starts, in //! O(n). Statements inserted later are not visited; statements erased @@ -40,15 +40,12 @@ //! //! - [`WalkResult::Advance`] goes on to the next statement. //! - [`WalkResult::Skip`] goes on without walking the statements nested in -//! this one, through [`WalkCursor::skip`]. [`PreOrder`] steps past them and +//! this one. [`PreOrder`] steps past them and //! [`BFS`] never queues them. The other orders have already visited them, //! so there it is the same as `Advance`. //! - [`WalkResult::Interrupt`] ends the walk, which reports //! [`WalkStatus::Interrupted`]. Edits already made stay. //! -//! A cursor driven by hand gets the same control: calling -//! [`WalkCursor::skip`] skips, and no longer calling `next` interrupts. -//! //! # Contract //! //! While a walk is visiting a statement, edits may add or erase statements @@ -76,8 +73,7 @@ //! block as `succ`), reported as [`WalkError::VisitedMoved`]; //! - moving it with [`Rewriter::move_before`](crate::Rewriter::move_before) //! or [`Rewriter::move_after`](crate::Rewriter::move_after), reported as -//! [`WalkError::VisitedMoved`] by [`WalkOrder::walk_mut`] (a loop driving a -//! cursor by hand checks with [`WalkOrder::check_not_moved_since`]); +//! [`WalkError::VisitedMoved`]; //! - splicing its block's statements elsewhere with //! [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) //! or [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after), @@ -99,14 +95,19 @@ //! visited in this walk. Discovery follows each order's rules above; nesting //! depth alone does not determine whether an insertion is ahead of the cursor. //! -//! The cursor reports an erase or merge of the visited statement that its -//! order does not allow as a [`WalkError`] on its next step. +//! These limits hold for the root like any other visited statement, and +//! whatever the visit returns: an edit that breaks them fails the walk even +//! if the visit interrupts it. +//! +//! [`WalkOrder::walk_mut`] reports an edit breaking these limits as a +//! [`WalkError`] once the visit returns. [`WalkOrder::walk`] cannot edit, so +//! it cannot fail. mod breadth_first; mod depth_first; mod order; pub(crate) mod tree; -pub use breadth_first::{BFS, BFSCursor, BuiltOrderCursor, KeepOrderReverseBFS, ReverseBFS}; -pub use depth_first::{PostOrder, PostOrderCursor, PreOrder, PreOrderCursor}; -pub use order::{VisitedPlace, WalkCursor, WalkError, WalkOrder, WalkResult, WalkStatus}; +pub use breadth_first::{BFS, KeepOrderReverseBFS, ReverseBFS}; +pub use depth_first::{PostOrder, PreOrder}; +pub use order::{WalkError, WalkOrder, WalkResult, WalkStatus}; diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index 80bfdcead..f8f510ff3 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -13,19 +13,17 @@ use crate::{ /// [`ReverseBFS`](super::ReverseBFS) and /// [`KeepOrderReverseBFS`](super::KeepOrderReverseBFS), so a driver selects /// its order as a type parameter. -pub trait WalkOrder { - /// The cursor walking this order. - type Cursor: WalkCursor; - +#[expect( + private_bounds, + reason = "`CursorOrder` is crate-private so other crates cannot step a cursor by hand: \ + that would skip `walk_mut`'s check that the visited statement was not moved" +)] +pub trait WalkOrder: CursorOrder { /// Whether an edit may erase or move the statement being visited (see /// the [contract](super#contract)). Only the reverse BFS orders allow it: /// they build their order at the start, while the others find the next /// statement through the visited one. - const ALLOWS_DISPLACING_VISITED: bool; - - /// A cursor at the start of the walk of the tree under `root`, `root` - /// included. - fn start(stage: &StageInfo, root: Statement) -> Self::Cursor; + const ALLOWS_ERASING_OR_MOVING_VISITED: bool; /// Every statement of the walk under `root`, in order, for a walk that /// edits nothing. Empty if `root` is not live. @@ -57,9 +55,9 @@ pub trait WalkOrder { /// [`walk`](Self::walk) with a `visit` that may edit through `rewriter`, /// within the walk's contract (see the [module docs](super)). Fails if - /// `visit` erases the statement it is visiting where the order does not - /// allow it, or moves it where the order does not allow it (see - /// [`check_not_moved_since`](Self::check_not_moved_since)). + /// `visit` erases or moves the statement it is visiting where the order + /// does not allow it, whatever `visit` returns: an interrupted walk is + /// checked too. fn walk_mut<'r, L: Dialect>( rewriter: &mut Rewriter<'r, L>, root: Statement, @@ -70,62 +68,60 @@ pub trait WalkOrder { let seen = rewriter.events().len(); let place = VisitedPlace::of(rewriter.stage(), stmt); let result = visit(rewriter, stmt); - Self::check_not_moved_since(rewriter, stmt, place, seen)?; + check_not_moved_since::(rewriter, stmt, place, seen)?; + // `None` on `Advance` and `Skip`, the final status on `Interrupt`. if let Some(status) = result.steer(&mut cursor) { + // The cursor checks erases and merges of the visited + // statement only when it steps; the walk ends here, so check + // them now. + cursor.check_visited(rewriter.stage())?; return Ok(status); } } Ok(WalkStatus::Completed) } +} - /// Fail with [`WalkError::VisitedMoved`] if an event recorded since the - /// first `seen` events of `rewriter` moved `visited`, in an order that - /// does not allow it - /// ([`ALLOWS_DISPLACING_VISITED`](Self::ALLOWS_DISPLACING_VISITED)): - /// a move of `visited` itself, or a splice of the block or CFG it was in - /// at `place`. Panics if `rewriter` has recorded fewer than `seen` events. - /// - /// [`walk_mut`](Self::walk_mut) calls this after each visit. A loop driving - /// a cursor by hand calls it too: note `rewriter.events().len()` and - /// [`VisitedPlace::of`] before editing at a statement, and check before - /// the next [`WalkCursor::next`]. - fn check_not_moved_since( - rewriter: &Rewriter<'_, L>, - visited: Statement, - place: VisitedPlace, - seen: usize, - ) -> Result<(), WalkError> { - if Self::ALLOWS_DISPLACING_VISITED { - return Ok(()); - } - let moved = rewriter.events()[seen..].iter().any(|event| match *event { - MutationEvent::MovedStatement { stmt } => stmt == visited, - MutationEvent::SplicedBlock { source, .. } => place.block == Some(source), - MutationEvent::SplicedCfg { source, .. } => place.cfg == Some(source), - _ => false, - }); - if moved { - Err(WalkError::VisitedMoved(visited)) - } else { - Ok(()) - } +/// Fail with [`WalkError::VisitedMoved`] if an event recorded since the first +/// `seen` events of `rewriter` moved `visited`, in an order `O` that does not +/// allow it ([`WalkOrder::ALLOWS_ERASING_OR_MOVING_VISITED`]): a move of +/// `visited` itself, or a splice of the block or CFG it was in at `place`. +/// Panics if `rewriter` has recorded fewer than `seen` events. +fn check_not_moved_since( + rewriter: &Rewriter<'_, L>, + visited: Statement, + place: VisitedPlace, + seen: usize, +) -> Result<(), WalkError> { + if O::ALLOWS_ERASING_OR_MOVING_VISITED { + return Ok(()); + } + let moved = rewriter.events()[seen..].iter().any(|event| match *event { + MutationEvent::MovedStatement { stmt } => stmt == visited, + MutationEvent::SplicedBlock { source, .. } => place.block == Some(source), + MutationEvent::SplicedCfg { source, .. } => place.cfg == Some(source), + _ => false, + }); + if moved { + Err(WalkError::VisitedMoved(visited)) + } else { + Ok(()) } } /// Where a visited statement was when its visit started: its block, and that -/// block's CFG. [`WalkOrder::check_not_moved_since`] compares splices against -/// it, since a splice changes both: afterwards, a statement spliced into a -/// block, or a block spliced into a CFG, cannot be told from one that was -/// there already. +/// block's CFG. [`check_not_moved_since`] compares splices against it, since +/// a splice changes both: afterwards, a statement spliced into a block, or a +/// block spliced into a CFG, cannot be told from one that was there already. #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct VisitedPlace { +struct VisitedPlace { block: Option, cfg: Option, } impl VisitedPlace { /// Where `stmt` is now in `stage`. - pub fn of(stage: &StageInfo, stmt: Statement) -> Self { + fn of(stage: &StageInfo, stmt: Statement) -> Self { let block = match stmt.get_info(stage).and_then(|info| info.parent) { Some(StatementParent::Block(block)) => Some(block), _ => None, @@ -144,8 +140,11 @@ pub enum WalkResult { /// Go on to the next statement. #[default] Advance, - /// Go on, without walking the statements nested in this one. See - /// [`WalkCursor::skip`] for the orders where this has an effect. + /// Go on, without walking the statements nested in this one. Has no + /// effect in an order that visits a statement after the statements + /// nested in it ([`PostOrder`](super::PostOrder), + /// [`ReverseBFS`](super::ReverseBFS), + /// [`KeepOrderReverseBFS`](super::KeepOrderReverseBFS)). Skip, /// End the walk. Interrupt, @@ -174,9 +173,24 @@ pub enum WalkStatus { Interrupted, } +/// The cursor an order walks with, and how to start one at a root. +/// +/// Crate-private, so that outside this crate a walk is only +/// [`WalkOrder::walk`], [`WalkOrder::walk_mut`] or [`WalkOrder::schedule`]: +/// an editing walk then always goes through `walk_mut`, which checks the +/// whole contract, while a cursor stepped by hand would miss moves. +pub(crate) trait CursorOrder { + /// The cursor walking this order. + type Cursor: WalkCursor; + + /// A cursor at the start of the walk of the tree under `root`, `root` + /// included. + fn start(stage: &StageInfo, root: Statement) -> Self::Cursor; +} + /// A position in a walk. It holds only ids, so the IR can be edited between /// steps. -pub trait WalkCursor { +pub(crate) trait WalkCursor { /// The next statement to visit, or `None` once the walk is done. /// /// The IR may be edited between calls within the walk's contract (see @@ -185,6 +199,12 @@ pub trait WalkCursor { /// it, or moved by merging its block into another. fn next(&mut self, stage: &StageInfo) -> Result, WalkError>; + /// Fail if the statement returned by the previous [`next`](Self::next) + /// has since been erased where the order does not allow it, or moved by + /// merging its block into another. Only reads; the cursor does not move. + /// [`next`](Self::next) runs this check first. + fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError>; + /// Do not walk the statements nested in the statement the previous /// [`next`](Self::next) returned, including any added while visiting it. /// Has no effect before the first step, after the last, or in an order @@ -210,8 +230,7 @@ pub enum WalkError { /// [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) / /// [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after), /// or by splicing its CFG with - /// [`Rewriter::splice_cfg_after`](crate::Rewriter::splice_cfg_after) (see - /// [`WalkOrder::check_not_moved_since`]). + /// [`Rewriter::splice_cfg_after`](crate::Rewriter::splice_cfg_after). VisitedMoved(Statement), } @@ -247,6 +266,14 @@ pub(super) enum Position { } impl Position { + /// The statement the cursor is at, if any. + pub(super) fn visited(self) -> Option { + match self { + Position::At(stmt) | Position::Skipped(stmt) => Some(stmt), + Position::Start | Position::Done => None, + } + } + /// Mark the statement the cursor is at as skipped; nothing to mark /// before the first statement or past the last. pub(super) fn skip(&mut self) { diff --git a/crates/kirin-ir/tests/walk.rs b/crates/kirin-ir/tests/walk.rs index da0d84be6..3eead8492 100644 --- a/crates/kirin-ir/tests/walk.rs +++ b/crates/kirin-ir/tests/walk.rs @@ -4,8 +4,8 @@ mod common; use common::{BuilderDialect, new_stage}; use kirin_ir::walk::{ - BFS, KeepOrderReverseBFS, PostOrder, PreOrder, ReverseBFS, WalkCursor, WalkError, WalkOrder, - WalkResult, WalkStatus, + BFS, KeepOrderReverseBFS, PostOrder, PreOrder, ReverseBFS, WalkError, WalkOrder, WalkResult, + WalkStatus, }; use kirin_ir::*; @@ -219,19 +219,20 @@ fn a_nested_root_walks_only_its_subtree() { assert_eq!(PostOrder::schedule(&t.stage, t.owner), post); } -/// Walk `O` under `t.root`, running `edit` at each statement, and return the -/// statements visited. +/// Walk `O` under `root` with `walk_mut`, running `edit` at each statement; +/// the statements visited, or the walk's error. fn walk_editing( - t: &mut Tree, + stage: &mut StageInfo, + root: Statement, mut edit: impl FnMut(&mut Rewriter, Statement), ) -> Result, WalkError> { - let mut rewriter = Rewriter::new(&mut t.stage); - let mut cursor = O::start(rewriter.stage(), t.root); + let mut rewriter = Rewriter::new(stage); let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(rewriter.stage())? { + O::walk_mut(&mut rewriter, root, |rewriter, stmt| { visited.push(stmt); - edit(&mut rewriter, stmt); - } + edit(rewriter, stmt); + WalkResult::Advance + })?; Ok(visited) } @@ -240,7 +241,7 @@ fn pre_order_walks_children_inserted_while_visiting() { let mut t = tree(); let (owner, then_ret) = (t.owner, t.then_ret); let mut inserted = None; - let visited = walk_editing::(&mut t, |rewriter, stmt| { + let visited = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == owner { inserted = Some( rewriter @@ -264,7 +265,7 @@ fn post_order_walks_siblings_inserted_after_the_visited_statement() { let mut t = tree(); let a = t.a; let mut inserted = None; - let visited = walk_editing::(&mut t, |rewriter, stmt| { + let visited = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == a { inserted = Some(rewriter.insert_after(a, BuilderDialect::Nop).unwrap()); } @@ -278,7 +279,7 @@ fn post_order_walks_siblings_inserted_after_the_visited_statement() { fn visits_with_insert_after_a() -> usize { let mut t = tree(); let a = t.a; - walk_editing::(&mut t, |rewriter, stmt| { + walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == a { rewriter.insert_after(a, BuilderDialect::Nop).unwrap(); } @@ -324,12 +325,8 @@ fn two_owners() -> (StageInfo, [Statement; 6]) { #[test] fn bfs_walks_statements_inserted_at_the_current_level_and_below() { let (mut stage, [root, a, b, c, x, y]) = two_owners(); - let mut rewriter = Rewriter::new(&mut stage); - let mut cursor = BFS::start(rewriter.stage(), root); let (mut n, mut z, mut h) = (None, None, None); - let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(rewriter.stage()).unwrap() { - visited.push(stmt); + let visited = walk_editing::(&mut stage, root, |rewriter, stmt| { if stmt == a { // Ahead of the cursor at the current level, and one level down. n = Some(rewriter.insert_after(b, BuilderDialect::Nop).unwrap()); @@ -339,7 +336,8 @@ fn bfs_walks_statements_inserted_at_the_current_level_and_below() { // At a level that is already finished. h = Some(rewriter.insert_after(a, BuilderDialect::Nop).unwrap()); } - } + }) + .unwrap(); let (n, z, h) = (n.unwrap(), z.unwrap(), h.unwrap()); assert_eq!(visited, [root, a, b, n, c, x, z, y]); assert!(!visited.contains(&h)); @@ -348,17 +346,14 @@ fn bfs_walks_statements_inserted_at_the_current_level_and_below() { #[test] fn bfs_skips_statements_erased_ahead_of_it() { let (mut stage, [root, a, b, c, x, _y]) = two_owners(); - let mut rewriter = Rewriter::new(&mut stage); - let mut cursor = BFS::start(rewriter.stage(), root); - let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(rewriter.stage()).unwrap() { - visited.push(stmt); + let visited = walk_editing::(&mut stage, root, |rewriter, stmt| { if stmt == a { // A sibling ahead, and a whole subtree waiting to be walked. rewriter.erase_statement(c).unwrap(); rewriter.erase_statement(b).unwrap(); } - } + }) + .unwrap(); assert_eq!(visited, [root, a, x]); } @@ -366,7 +361,7 @@ fn bfs_skips_statements_erased_ahead_of_it() { fn erase_while_visiting() -> (Statement, Result, WalkError>) { let mut t = tree(); let a = t.a; - let result = walk_editing::(&mut t, |rewriter, stmt| { + let result = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == a { rewriter.erase_statement(a).unwrap(); } @@ -401,7 +396,7 @@ fn reverse_bfs_orders_go_on_after_the_visited_statement_is_erased() { fn reverse_bfs_orders_skip_the_rest_of_an_ancestor_erased_while_visiting() { let mut t = tree(); let (owner, then_nop, then_ret) = (t.owner, t.then_nop, t.then_ret); - let reverse = walk_editing::(&mut t, |rewriter, stmt| { + let reverse = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == then_ret { rewriter.erase_statement(owner).unwrap(); } @@ -412,7 +407,7 @@ fn reverse_bfs_orders_skip_the_rest_of_an_ancestor_erased_while_visiting() { ); let mut t = tree(); - let keep_order = walk_editing::(&mut t, |rewriter, stmt| { + let keep_order = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == then_nop { rewriter.erase_statement(owner).unwrap(); } @@ -423,15 +418,12 @@ fn reverse_bfs_orders_skip_the_rest_of_an_ancestor_erased_while_visiting() { #[test] fn bfs_goes_on_with_the_next_parent_after_the_parent_is_erased() { let (mut stage, [root, a, b, c, x, y]) = two_owners(); - let mut rewriter = Rewriter::new(&mut stage); - let mut cursor = BFS::start(rewriter.stage(), root); - let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(rewriter.stage()).unwrap() { - visited.push(stmt); + let visited = walk_editing::(&mut stage, root, |rewriter, stmt| { if stmt == x { rewriter.erase_statement(a).unwrap(); } - } + }) + .unwrap(); assert_eq!(visited, [root, a, b, c, x, y]); } @@ -470,33 +462,27 @@ fn bfs_three_levels() -> (StageInfo, [Statement; 7]) { #[test] fn bfs_goes_on_outside_a_grandparent_erased_while_visiting() { let (mut stage, [root, g, h, p, q, a, d]) = bfs_three_levels(); - let mut rewriter = Rewriter::new(&mut stage); - let mut cursor = BFS::start(rewriter.stage(), root); - let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(rewriter.stage()).unwrap() { - visited.push(stmt); + let visited = walk_editing::(&mut stage, root, |rewriter, stmt| { if stmt == a { // Erases `p` and `a` with it; `p` is the front parent. rewriter.erase_statement(g).unwrap(); } - } + }) + .unwrap(); assert_eq!(visited, [root, g, h, p, q, a, d]); } #[test] fn pre_order_skips_children_and_siblings_erased_while_visiting() { let (mut stage, [root, a, b, c, x, y]) = two_owners(); - let mut rewriter = Rewriter::new(&mut stage); - let mut cursor = PreOrder::start(rewriter.stage(), root); - let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(rewriter.stage()).unwrap() { - visited.push(stmt); + let visited = walk_editing::(&mut stage, root, |rewriter, stmt| { if stmt == a { // Its only child, and a sibling ahead of the cursor. rewriter.erase_statement(x).unwrap(); rewriter.erase_statement(c).unwrap(); } - } + }) + .unwrap(); assert_eq!(visited, [root, a, b, y]); } @@ -654,16 +640,36 @@ fn walk_mut_fails_when_the_visited_statement_is_erased() { assert_eq!(result, Err(WalkError::VisitedErased(a))); } +/// Erases `a` while visiting it, then interrupts the walk. +fn erase_a_then_interrupt() -> (Statement, Result) { + let mut t = tree(); + let (root, a) = (t.root, t.a); + let mut rewriter = Rewriter::new(&mut t.stage); + let result = O::walk_mut(&mut rewriter, root, |rewriter, stmt| { + if stmt == a { + rewriter.erase_statement(a).unwrap(); + return WalkResult::Interrupt; + } + WalkResult::Advance + }); + (a, result) +} + #[test] -fn skip_before_the_first_step_has_no_effect() { - let t = tree(); - let mut cursor = PreOrder::start(&t.stage, t.root); - cursor.skip(); - let mut visited = Vec::new(); - while let Some(stmt) = cursor.next(&t.stage).unwrap() { - visited.push(stmt); +fn walk_mut_fails_when_the_visited_statement_is_erased_before_an_interrupt() { + for (a, result) in [ + erase_a_then_interrupt::(), + erase_a_then_interrupt::(), + erase_a_then_interrupt::(), + ] { + assert_eq!(result, Err(WalkError::VisitedErased(a))); + } + for (_, result) in [ + erase_a_then_interrupt::(), + erase_a_then_interrupt::(), + ] { + assert_eq!(result, Ok(WalkStatus::Interrupted)); } - assert_eq!(visited, PreOrder::schedule(&t.stage, t.root)); } /// The ids of [`merge_cfg`]'s blocks and statements. @@ -719,21 +725,11 @@ fn edit_at_z( ) -> (Vec, Result, WalkError>) { let (mut stage, m) = merge_cfg(); let schedule = O::schedule(&stage, m.root); - let mut rewriter = Rewriter::new(&mut stage); - let mut cursor = O::start(rewriter.stage(), m.root); - let mut visited = Vec::new(); - let result = loop { - match cursor.next(rewriter.stage()) { - Ok(Some(stmt)) => { - visited.push(stmt); - if stmt == m.z { - edit(&mut rewriter, m); - } - } - Ok(None) => break Ok(visited), - Err(error) => break Err(error), + let result = walk_editing::(&mut stage, m.root, |rewriter, stmt| { + if stmt == m.z { + edit(rewriter, m); } - }; + }); (schedule, result) } @@ -753,6 +749,29 @@ fn merging_the_visited_statements_block_away_breaks_link_following_walks() { } } +/// Walks `O` from `z` as the root, merging `z`'s block away while visiting +/// it. +fn merge_the_roots_block() -> (Statement, Result) { + let (mut stage, m) = merge_cfg(); + let mut rewriter = Rewriter::new(&mut stage); + let result = O::walk_mut(&mut rewriter, m.z, |rewriter, _| { + merge_b2_into_b1(rewriter, m); + WalkResult::Advance + }); + (m.z, result) +} + +#[test] +fn merging_the_roots_block_away_breaks_link_following_walks() { + for (z, result) in [ + merge_the_roots_block::(), + merge_the_roots_block::(), + merge_the_roots_block::(), + ] { + assert_eq!(result, Err(WalkError::VisitedMoved(z))); + } +} + #[test] fn reverse_bfs_orders_visit_statements_where_they_were_before_a_merge() { let (_, m) = merge_cfg(); @@ -766,28 +785,11 @@ fn reverse_bfs_orders_visit_statements_where_they_were_before_a_merge() { assert_eq!(result.unwrap(), without_br); } -/// Walk `O` over [`tree`] with `walk_mut`, running `edit` at each statement; -/// the statements visited, or the walk's error. -fn walk_mut_editing( - t: &mut Tree, - mut edit: impl FnMut(&mut Rewriter, Statement), -) -> Result, WalkError> { - let root = t.root; - let mut rewriter = Rewriter::new(&mut t.stage); - let mut visited = Vec::new(); - O::walk_mut(&mut rewriter, root, |rewriter, stmt| { - visited.push(stmt); - edit(rewriter, stmt); - WalkResult::Advance - })?; - Ok(visited) -} - /// Moves `a` into the `then` arm, before `then_nop`, while visiting it. fn move_a_while_visiting() -> (Statement, Result, WalkError>) { let mut t = tree(); let (a, then_nop) = (t.a, t.then_nop); - let result = walk_mut_editing::(&mut t, |rewriter, stmt| { + let result = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == a { rewriter.move_before(a, then_nop).unwrap(); } @@ -824,7 +826,7 @@ fn move_else_nop_while_visiting_a() -> (Statement, Result(&mut t, |rewriter, stmt| { + let result = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == a { rewriter.move_after(else_nop, then_nop).unwrap(); } @@ -858,7 +860,7 @@ fn splice_then_arm_while_visiting_it() let mut t = tree(); let (then_nop, owner) = (t.then_nop, t.owner); let then_block = block_of(&t.stage, then_nop); - let result = walk_mut_editing::(&mut t, |rewriter, stmt| { + let result = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == then_nop { rewriter .splice_block_before(then_block, owner, &[]) @@ -908,7 +910,7 @@ fn splicing_another_block_into_the_visited_statements_block_keeps_the_walk() { fn splice_while_visiting_a() -> Vec { let mut t = tree(); let ids = tree(); - walk_mut_editing::(&mut t, |rewriter, stmt| { + walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == ids.a { splice_else_arm_after_a(rewriter, &ids); } @@ -1044,7 +1046,7 @@ fn split_at_a_while_visiting() -> (Vec, Result(&mut t, |rewriter, stmt| { + let result = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { if stmt == a { rewriter.split_block(a).unwrap(); } diff --git a/crates/kirin-rewrite/src/driver/rewrite_driver.rs b/crates/kirin-rewrite/src/driver/rewrite_driver.rs index 47f6c6c1a..78b629afc 100644 --- a/crates/kirin-rewrite/src/driver/rewrite_driver.rs +++ b/crates/kirin-rewrite/src/driver/rewrite_driver.rs @@ -1,4 +1,4 @@ -use kirin_ir::walk::{WalkCursor, WalkOrder, WalkResult}; +use kirin_ir::walk::{WalkOrder, WalkResult}; use kirin_ir::{Dialect, GetInfo, Rewriter, Statement}; use super::{DriverError, Match}; @@ -118,7 +118,7 @@ impl<'r, L: Dialect> RewriteDriver<'r, L> { rewriter: &mut Rewriter<'_, L>, root: Statement, ) -> Result { - if !O::ALLOWS_DISPLACING_VISITED + if !O::ALLOWS_ERASING_OR_MOVING_VISITED && let Some(rule) = self.rules.rules().iter().find(|rule| rule.erases_root()) { let order = std::any::type_name::(); @@ -153,38 +153,70 @@ impl<'r, L: Dialect> RewriteDriver<'r, L> { /// remaining rules still run at the statement, then its nested /// statements are skipped; an `interrupt` stops trying rules and ends /// the sweep. + /// + /// The sweep is a [`WalkOrder::walk_mut`], so it checks the walk's + /// contract like any other walk. A rule failing at a statement interrupts + /// it, and that failure is reported before any broken contract. fn sweep( &self, rewriter: &mut Rewriter<'_, L>, root: Statement, ) -> Result { let mut applied = 0; - let mut cursor = O::start(rewriter.stage(), root); - while let Some(stmt) = cursor.next(rewriter.stage())? { - let mut skip = false; - for rule in self.rules.rules() { - let Some(found) = Match::find(self.rules, rule, rewriter.stage(), stmt)? else { - continue; - }; - let before = rewriter.events().len(); - let walk_result = found.apply(self.rules, rule, rewriter, stmt)?; - if rewriter.events().len() > before { - applied += 1; - } - match walk_result { - WalkResult::Advance => {} - WalkResult::Skip => skip = true, - WalkResult::Interrupt => return Ok(applied), + let mut failure = None; + let walked = O::walk_mut(rewriter, root, |rewriter, stmt| { + match self.visit(rewriter, stmt) { + Ok((edits, walk_result)) => { + applied += edits; + walk_result } - // The rule erased the statement, so no later rule can match it. - if stmt.get_info(rewriter.stage()).is_none() { - break; + Err(error) => { + failure = Some(error); + WalkResult::Interrupt } } - if skip { - cursor.skip(); - } + }); + if let Some(error) = failure { + return Err(error); } + walked?; Ok(applied) } + + /// Try every rule at `stmt`, applying each that matches. Returns how many + /// applications edited the IR, and how the walk goes on: `Skip` if any + /// applied rule skipped, `Interrupt` as soon as one interrupts. + fn visit( + &self, + rewriter: &mut Rewriter<'_, L>, + stmt: Statement, + ) -> Result<(usize, WalkResult), DriverError> { + let mut applied = 0; + let mut skip = false; + for rule in self.rules.rules() { + let Some(found) = Match::find(self.rules, rule, rewriter.stage(), stmt)? else { + continue; + }; + let before = rewriter.events().len(); + let walk_result = found.apply(self.rules, rule, rewriter, stmt)?; + if rewriter.events().len() > before { + applied += 1; + } + match walk_result { + WalkResult::Advance => {} + WalkResult::Skip => skip = true, + WalkResult::Interrupt => return Ok((applied, WalkResult::Interrupt)), + } + // The rule erased the statement, so no later rule can match it. + if stmt.get_info(rewriter.stage()).is_none() { + break; + } + } + let walk_result = if skip { + WalkResult::Skip + } else { + WalkResult::Advance + }; + Ok((applied, walk_result)) + } } diff --git a/crates/kirin-rewrite/tests/walk_result.rs b/crates/kirin-rewrite/tests/walk_result.rs index e11b9d428..2a03fa0a9 100644 --- a/crates/kirin-rewrite/tests/walk_result.rs +++ b/crates/kirin-rewrite/tests/walk_result.rs @@ -3,8 +3,8 @@ use kirin::prelude::*; use kirin_arith::{Arith, ArithType}; -use kirin_ir::walk::{PreOrder, WalkResult}; -use kirin_ir::{BuilderSSAKind, Placeholder}; +use kirin_ir::walk::{PreOrder, WalkError, WalkResult}; +use kirin_ir::{BuilderSSAKind, GetInfo, Placeholder, SSAKind, StatementParent}; use kirin_rewrite::{ DriverError, NativeValue, PatternRule, RewriteDriver, RewriteSummary, RuleErrorKind, RuleSet, }; @@ -67,6 +67,51 @@ fn neg_unchanged_then_interrupt() { }) } +/// Moves the statement defining `r` to the front of its block. Returns +/// whether it moved. +fn hoist(rewriter: &mut Rewriter<'_, L>, r: SSAValue) -> bool { + let stage = rewriter.stage(); + let SSAKind::Result(stmt, _) = *r.get_info(stage).expect("`r` is live").kind() else { + panic!("`r` is a result"); + }; + let Some(StatementParent::Block(block)) = *stmt.parent(stage) else { + panic!("`r`'s statement is in a block"); + }; + let first = block + .first_statement(stage) + .expect("the block is not empty"); + rewriter.move_before(stmt, first).unwrap(); + first != stmt +} + +/// Moves any `neg` to the front of its block. +#[rewrite_rule(language = L)] +fn hoist_neg() { + let r = op![Arith::neg(x)]; + rewrite!(r, { + let _ = hoist(r); + }) +} + +/// Ends the sweep if `moved`. +fn interrupt_if(_rewriter: &mut Rewriter<'_, L>, moved: bool) -> WalkResult { + if moved { + WalkResult::Interrupt + } else { + WalkResult::Advance + } +} + +/// Moves any `neg` to the front of its block, ending the sweep if it moved. +#[rewrite_rule(language = L)] +fn hoist_neg_then_interrupt() { + let r = op![Arith::neg(x)]; + rewrite!(r, { + let moved = hoist(r); + interrupt_if(moved) + }) +} + /// `-(-x)` is `x`, going on as `decide` says. #[rewrite_rule(language = L)] fn double_neg_then_decide() { @@ -349,3 +394,17 @@ fn skip_does_not_change_the_run_yet() { let skipping = run(&rules(double_neg_then_skip)).unwrap(); assert_eq!(skipping, advancing); } + +/// Moving the statement being visited breaks the pre-order walk's +/// contract, which the sweep checks like any walk, even when the rule then +/// interrupts it. +#[test] +fn moving_the_visited_statement_fails_the_run() { + for rules in [rules(hoist_neg), rules(hoist_neg_then_interrupt)] { + let error = run(&rules).err(); + assert!( + matches!(error, Some(DriverError::Walk(WalkError::VisitedMoved(_)))), + "{error:?}" + ); + } +} diff --git a/docs/design/rewrite-engine.md b/docs/design/rewrite-engine.md index 83162f757..c4fdf3c60 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -2886,8 +2886,8 @@ sweep recording no `MutationEvent`. Decisions to revisit in later milestones: - A rule may erase its root, the statement being visited, only in a walk - order that allows it (`WalkOrder::ALLOWS_DISPLACING_VISITED`: the reverse BFS - orders). A run in another order fails with `DriverError::ErasesRoot` + order that allows it (`WalkOrder::ALLOWS_ERASING_OR_MOVING_VISITED`: the + reverse BFS orders). A run in another order fails with `DriverError::ErasesRoot` before editing anything. A rule's other erase targets are producers of the root, never nested in or around it, so they meet the walk contract in every order. Native actions are checked by the walk as it runs. From 146e994691aa4bb78843b979f18e40468383be52 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 14:09:00 -0500 Subject: [PATCH 08/14] fix(walk): count a splice as moving the visited statement only if it moved Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/walk/mod.rs | 3 +- crates/kirin-ir/src/rewrite/walk/order.rs | 25 +++++++++---- crates/kirin-ir/tests/walk.rs | 44 +++++++++++++++++++++++ 3 files changed, 65 insertions(+), 7 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/walk/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index 2a00d6e46..2216f0ee3 100644 --- a/crates/kirin-ir/src/rewrite/walk/mod.rs +++ b/crates/kirin-ir/src/rewrite/walk/mod.rs @@ -77,7 +77,8 @@ //! - splicing its block's statements elsewhere with //! [`Rewriter::splice_block_before`](crate::Rewriter::splice_block_before) //! or [`Rewriter::splice_block_after`](crate::Rewriter::splice_block_after), -//! reported the same way; +//! reported the same way (a splice leaves the block's terminator behind, so +//! it is allowed while visiting the terminator); //! - splicing the blocks of its CFG into another CFG with //! [`Rewriter::splice_cfg_after`](crate::Rewriter::splice_cfg_after), //! reported the same way; diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index f8f510ff3..968474a90 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -85,7 +85,9 @@ pub trait WalkOrder: CursorOrder { /// Fail with [`WalkError::VisitedMoved`] if an event recorded since the first /// `seen` events of `rewriter` moved `visited`, in an order `O` that does not /// allow it ([`WalkOrder::ALLOWS_ERASING_OR_MOVING_VISITED`]): a move of -/// `visited` itself, or a splice of the block or CFG it was in at `place`. +/// `visited` itself, or a splice of the block or CFG it was in at `place` +/// that took it along. A splice that left it behind, such as the terminator +/// of a spliced block, or one that moved nothing, is not a move. /// Panics if `rewriter` has recorded fewer than `seen` events. fn check_not_moved_since( rewriter: &Rewriter<'_, L>, @@ -96,10 +98,21 @@ fn check_not_moved_since( if O::ALLOWS_ERASING_OR_MOVING_VISITED { return Ok(()); } + // Where `visited` is now. An erased one has no place; the cursor + // reports its erasure. + let now = || { + visited + .get_info(rewriter.stage()) + .map(|_| VisitedPlace::of(rewriter.stage(), visited)) + }; let moved = rewriter.events()[seen..].iter().any(|event| match *event { MutationEvent::MovedStatement { stmt } => stmt == visited, - MutationEvent::SplicedBlock { source, .. } => place.block == Some(source), - MutationEvent::SplicedCfg { source, .. } => place.cfg == Some(source), + MutationEvent::SplicedBlock { source, .. } => { + place.block == Some(source) && now().is_some_and(|now| now.block != place.block) + } + MutationEvent::SplicedCfg { source, .. } => { + place.cfg == Some(source) && now().is_some_and(|now| now.cfg != place.cfg) + } _ => false, }); if moved { @@ -110,9 +123,9 @@ fn check_not_moved_since( } /// Where a visited statement was when its visit started: its block, and that -/// block's CFG. [`check_not_moved_since`] compares splices against it, since -/// a splice changes both: afterwards, a statement spliced into a block, or a -/// block spliced into a CFG, cannot be told from one that was there already. +/// block's CFG. [`check_not_moved_since`] needs it for splices: it names the +/// splices that could have moved the statement, and comparing it with where +/// the statement is afterwards tells whether one did. #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct VisitedPlace { block: Option, diff --git a/crates/kirin-ir/tests/walk.rs b/crates/kirin-ir/tests/walk.rs index 3eead8492..d6d506b07 100644 --- a/crates/kirin-ir/tests/walk.rs +++ b/crates/kirin-ir/tests/walk.rs @@ -881,6 +881,50 @@ fn splicing_the_visited_statements_block_breaks_link_following_walks() { } } +/// Walks `O` over [`tree`] and, while visiting `then_ret`, splices the `then` +/// block's statements before `owner`. `then_ret` is that block's terminator, +/// so the splice does not move it. +/// +/// If `empty_first` is true, the block's statements are spliced away once +/// before the walk, so the splice during the walk moves nothing. +/// +/// Returns the walk's schedule taken before the walk starts, and what the +/// walk returned: the statements it visited, or its error. +fn splice_then_arm_while_visiting_its_terminator( + empty_first: bool, +) -> (Vec, Result, WalkError>) { + let mut t = tree(); + let (then_ret, owner) = (t.then_ret, t.owner); + let then_block = block_of(&t.stage, then_ret); + if empty_first { + Rewriter::new(&mut t.stage) + .splice_block_before(then_block, owner, &[]) + .unwrap(); + } + let schedule = O::schedule(&t.stage, t.root); + let result = walk_editing::(&mut t.stage, t.root, |rewriter, stmt| { + if stmt == then_ret { + rewriter + .splice_block_before(then_block, owner, &[]) + .unwrap(); + } + }); + (schedule, result) +} + +#[test] +fn splicing_the_visited_terminators_block_keeps_the_walk() { + for empty_first in [false, true] { + for (schedule, result) in [ + splice_then_arm_while_visiting_its_terminator::(empty_first), + splice_then_arm_while_visiting_its_terminator::(empty_first), + splice_then_arm_while_visiting_its_terminator::(empty_first), + ] { + assert_eq!(result.unwrap(), schedule); + } + } +} + #[test] fn reverse_bfs_orders_go_on_after_the_visited_statements_block_is_spliced() { let t = tree(); From 7f2d20e23ffc43cc2a4c44a9beec5b3f24120257 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 14:31:24 -0500 Subject: [PATCH 09/14] refactor(ir): link and unlink statements and blocks through shared list helpers Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/node/linked_list.rs | 9 + crates/kirin-ir/src/rewrite/cfg/blocks.rs | 2 +- crates/kirin-ir/src/rewrite/cfg/helpers.rs | 73 +---- crates/kirin-ir/src/rewrite/cfg/merge.rs | 4 +- crates/kirin-ir/src/rewrite/cfg/splice.rs | 58 +--- crates/kirin-ir/src/rewrite/cfg/split.rs | 46 +-- crates/kirin-ir/src/rewrite/list.rs | 313 +++++++++++++++++++++ crates/kirin-ir/src/rewrite/mod.rs | 1 + crates/kirin-ir/src/rewrite/moves.rs | 62 +--- crates/kirin-ir/src/rewrite/rewriter.rs | 92 +----- 10 files changed, 350 insertions(+), 310 deletions(-) create mode 100644 crates/kirin-ir/src/rewrite/list.rs diff --git a/crates/kirin-ir/src/node/linked_list.rs b/crates/kirin-ir/src/node/linked_list.rs index f6ca939a6..1e38ccdb7 100644 --- a/crates/kirin-ir/src/node/linked_list.rs +++ b/crates/kirin-ir/src/node/linked_list.rs @@ -48,6 +48,15 @@ impl LinkedList { } } + /// A list of the one item `ptr`. + pub(crate) fn single(ptr: Ptr) -> Self { + LinkedList { + head: Some(ptr), + tail: Some(ptr), + len: 1, + } + } + pub fn head(&self) -> Option<&Ptr> { self.head.as_ref() } diff --git a/crates/kirin-ir/src/rewrite/cfg/blocks.rs b/crates/kirin-ir/src/rewrite/cfg/blocks.rs index 4d6e8e13e..90b77baad 100644 --- a/crates/kirin-ir/src/rewrite/cfg/blocks.rs +++ b/crates/kirin-ir/src/rewrite/cfg/blocks.rs @@ -70,7 +70,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.check_no_outside_predecessors(block, &[block])?; self.check_no_outside_uses(block)?; - self.unlink_block(cfg, block); + self.unlink(block, cfg); self.tombstone_block(block); self.events.push(MutationEvent::ChangedBlockList { cfg }); diff --git a/crates/kirin-ir/src/rewrite/cfg/helpers.rs b/crates/kirin-ir/src/rewrite/cfg/helpers.rs index fd5d84193..2252e77d6 100644 --- a/crates/kirin-ir/src/rewrite/cfg/helpers.rs +++ b/crates/kirin-ir/src/rewrite/cfg/helpers.rs @@ -3,7 +3,7 @@ use std::collections::HashSet; use smallvec::SmallVec; use crate::arena::{GetInfo, GetInfoMut}; -use crate::node::linked_list::LinkedListNode; +use crate::node::linked_list::{LinkedList, LinkedListNode}; use crate::{ Block, BlockArgument, BlockInfo, BlockParent, CFG, Dialect, SSAInfo, SSAKind, SSAValue, Statement, StatementParent, Use, @@ -203,11 +203,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { let block = self.stage.nodes.blocks.alloc_with_id(|id| { BlockInfo::builder() .parent(BlockParent::CFG(cfg)) - .node(LinkedListNode { - ptr: id, - next, - prev, - }) + .node(LinkedListNode::new(id)) .arguments(Vec::new()) .predecessors(SmallVec::new()) .new() @@ -221,72 +217,9 @@ impl<'a, L: Dialect> Rewriter<'a, L> { info.arguments = args; } - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = Some(block); - } - } - None => { - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.head = Some(block); - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = Some(block); - } - } - None => { - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.tail = Some(block); - } - } - } - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.len += 1; - } + self.link(LinkedList::single(block), cfg, prev, next); self.events.push(MutationEvent::ChangedBlockList { cfg }); block } - - /// Unlink `block` from `cfg`'s block list, clearing its own links. - pub(super) fn unlink_block(&mut self, cfg: CFG, block: Block) { - let node = block.expect_info(self.stage).node; - let (prev, next) = (node.prev, node.next); - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = next; - } - } - None => { - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.head = next; - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = prev; - } - } - None => { - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.tail = prev; - } - } - } - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.len -= 1; - } - if let Some(info) = block.get_info_mut(self.stage) { - info.node.prev = None; - info.node.next = None; - } - } } diff --git a/crates/kirin-ir/src/rewrite/cfg/merge.rs b/crates/kirin-ir/src/rewrite/cfg/merge.rs index 369a1652c..b93796ad2 100644 --- a/crates/kirin-ir/src/rewrite/cfg/merge.rs +++ b/crates/kirin-ir/src/rewrite/cfg/merge.rs @@ -90,7 +90,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { None => (LinkedList::new(), None), }; let pred_tail = pred.expect_info(self.stage).statements.tail; - self.link_chain(succ_chain, pred, pred_tail, None); + self.link(succ_chain, pred, pred_tail, None); if let Some(terminator) = succ_terminator && let Some(info) = terminator.get_info_mut(self.stage) { @@ -102,7 +102,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { // succ is now empty: tombstoning it only removes the block and its // (already redirected) arguments. - self.unlink_block(cfg, succ); + self.unlink(succ, cfg); self.tombstone_block(succ); let pred_successors_after = self.block_successors(pred); diff --git a/crates/kirin-ir/src/rewrite/cfg/splice.rs b/crates/kirin-ir/src/rewrite/cfg/splice.rs index 86ee562c9..aba99f89b 100644 --- a/crates/kirin-ir/src/rewrite/cfg/splice.rs +++ b/crates/kirin-ir/src/rewrite/cfg/splice.rs @@ -48,7 +48,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { None => LinkedList::new(), }; let next = anchor.expect_info(self.stage).node.next; - self.link_blocks(chain, dest, Some(anchor), next); + self.link(chain, dest, Some(anchor), next); self.events .push(MutationEvent::ChangedBlockList { cfg: source }); @@ -57,60 +57,4 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.events.push(MutationEvent::SplicedCfg { source, dest }); Ok(()) } - - /// Link `chain`, a whole block chain already taken out of its CFG, into - /// `cfg` between `prev` and `next` (the list head or tail when one is - /// `None`), making `cfg` the parent of each block. - fn link_blocks( - &mut self, - chain: LinkedList, - cfg: CFG, - prev: Option, - next: Option, - ) { - let (Some(head), Some(tail)) = (chain.head, chain.tail) else { - return; - }; - let mut current = Some(head); - while let Some(block) = current { - let Some(info) = block.get_info_mut(self.stage) else { - break; - }; - info.parent = Some(BlockParent::CFG(cfg)); - current = info.node.next; - } - if let Some(info) = head.get_info_mut(self.stage) { - info.node.prev = prev; - } - if let Some(info) = tail.get_info_mut(self.stage) { - info.node.next = next; - } - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = Some(head); - } - } - None => { - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.head = Some(head); - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = Some(tail); - } - } - None => { - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.tail = Some(tail); - } - } - } - if let Some(info) = cfg.get_info_mut(self.stage) { - info.blocks.len += chain.len; - } - } } diff --git a/crates/kirin-ir/src/rewrite/cfg/split.rs b/crates/kirin-ir/src/rewrite/cfg/split.rs index feb667694..e80328b54 100644 --- a/crates/kirin-ir/src/rewrite/cfg/split.rs +++ b/crates/kirin-ir/src/rewrite/cfg/split.rs @@ -21,12 +21,12 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// Records [`MutationEvent::ChangedBlockList`] then /// [`MutationEvent::ChangedCfg`]. pub fn split_block(&mut self, at: Statement) -> Result { - let (block, at_is_terminator, at_prev) = { + let (block, at_is_terminator) = { let item = self.live_statement(at)?; let Some(StatementParent::Block(block)) = item.parent else { return Err(RewriteError::NotInBlockBody(at)); }; - (block, item.definition.is_terminator(), item.node.prev) + (block, item.definition.is_terminator()) }; let cfg = self.owning_cfg(block)?; let (next_block, terminator) = { @@ -35,43 +35,25 @@ impl<'a, L: Dialect> Rewriter<'a, L> { }; let successors_before = self.block_successors(block); + // Cut `at` and the statements after it off the original block's + // chain. Splitting at the terminator cuts nothing: it is not in the + // chain. let moved = if at_is_terminator { - Vec::new() + LinkedList::new() } else { - self.chain_from(Some(at)) + self.split_off(block, at) }; + // Give the cut statements and the terminator to a new block right + // after the original. let new_block = self.create_block(cfg, Some(block), next_block, []); - - // Cut the moved suffix off the original block's statement chain. - if !moved.is_empty() { - if let Some(prev) = at_prev - && let Some(info) = prev.get_info_mut(self.stage) - { - info.node.next = None; - } - if let Some(info) = at.get_info_mut(self.stage) { - info.node.prev = None; - } - if let Some(info) = block.get_info_mut(self.stage) { - info.statements.tail = at_prev; - if at_prev.is_none() { - info.statements.head = None; - } - info.statements.len -= moved.len(); - } - } - for &stmt in moved.iter().chain(terminator.iter()) { - if let Some(info) = stmt.get_info_mut(self.stage) { - info.parent = Some(StatementParent::Block(new_block)); - } + self.link(moved, new_block, None, None); + if let Some(terminator) = terminator + && let Some(info) = terminator.get_info_mut(self.stage) + { + info.parent = Some(StatementParent::Block(new_block)); } if let Some(info) = new_block.get_info_mut(self.stage) { - info.statements = LinkedList { - head: moved.first().copied(), - tail: moved.last().copied(), - len: moved.len(), - }; info.terminator = terminator; } if let Some(info) = block.get_info_mut(self.stage) { diff --git a/crates/kirin-ir/src/rewrite/list.rs b/crates/kirin-ir/src/rewrite/list.rs new file mode 100644 index 000000000..80bee755b --- /dev/null +++ b/crates/kirin-ir/src/rewrite/list.rs @@ -0,0 +1,313 @@ +//! Linking and unlinking the IR's two doubly linked lists: a block's +//! statement chain and a CFG's block list. Each operation is written once, +//! for any [`ListItem`]. + +use crate::arena::{GetInfo, GetInfoMut}; +use crate::node::linked_list::{LinkedList, LinkedListNode}; +use crate::rewrite::rewriter::Rewriter; +use crate::{Block, BlockParent, CFG, Dialect, Statement, StatementParent}; + +/// An item of one of the IR's linked lists, kept in a list its `Owner` holds: +/// a [`Statement`] in its block's statement chain, or a [`Block`] in its +/// CFG's block list. +pub(super) trait ListItem: GetInfoMut + Copy + PartialEq { + /// What holds the list. + type Owner: GetInfoMut + Copy; + + /// The item's links to its neighbours. + fn node(info: &mut Self::Info) -> &mut LinkedListNode; + + /// Make `owner` the item's parent. + fn set_owner(info: &mut Self::Info, owner: Self::Owner); + + /// The list `owner` holds. + fn list(info: &mut >::Info) -> &mut LinkedList; +} + +impl ListItem for Statement { + type Owner = Block; + + fn node(info: &mut Self::Info) -> &mut LinkedListNode { + &mut info.node + } + + fn set_owner(info: &mut Self::Info, owner: Block) { + info.parent = Some(StatementParent::Block(owner)); + } + + fn list(info: &mut >::Info) -> &mut LinkedList { + &mut info.statements + } +} + +impl ListItem for Block { + type Owner = CFG; + + fn node(info: &mut Self::Info) -> &mut LinkedListNode { + &mut info.node + } + + fn set_owner(info: &mut Self::Info, owner: CFG) { + info.parent = Some(BlockParent::CFG(owner)); + } + + fn list(info: &mut >::Info) -> &mut LinkedList { + &mut info.blocks + } +} + +// Items and owners that are no longer live are skipped rather than +// panicking on, as the rewriter's edits check liveness before linking. +impl<'a, L: Dialect> Rewriter<'a, L> { + /// Link `chain`, a chain of items not in any list, into `owner`'s list + /// between `prev` and `next` (the list's head or tail when one is + /// `None`), making `owner` the parent of each item. For a single item, + /// pass [`LinkedList::single`]. + pub(super) fn link>( + &mut self, + chain: LinkedList, + owner: T::Owner, + prev: Option, + next: Option, + ) { + let (Some(head), Some(tail)) = (chain.head, chain.tail) else { + return; + }; + // Stop at `tail`: a single item may still carry links from before + // it was unlinked. + let mut current = Some(head); + while let Some(item) = current + && let Some(info) = item.get_info_mut(self.stage) + { + T::set_owner(info, owner); + current = if item == tail { + None + } else { + T::node(info).next + }; + } + if let Some(info) = head.get_info_mut(self.stage) { + T::node(info).prev = prev; + } + if let Some(info) = tail.get_info_mut(self.stage) { + T::node(info).next = next; + } + match prev { + Some(prev) => { + if let Some(info) = prev.get_info_mut(self.stage) { + T::node(info).next = Some(head); + } + } + None => { + if let Some(info) = owner.get_info_mut(self.stage) { + T::list(info).head = Some(head); + } + } + } + match next { + Some(next) => { + if let Some(info) = next.get_info_mut(self.stage) { + T::node(info).prev = Some(tail); + } + } + None => { + if let Some(info) = owner.get_info_mut(self.stage) { + T::list(info).tail = Some(tail); + } + } + } + if let Some(info) = owner.get_info_mut(self.stage) { + T::list(info).len += chain.len; + } + } + + /// Unlink `item` from `owner`'s list, joining its neighbours (or moving + /// the list's head or tail when it is at an end) and clearing its own + /// links. Its parent is left for the caller. + pub(super) fn unlink>(&mut self, item: T, owner: T::Owner) { + let Some(info) = item.get_info_mut(self.stage) else { + return; + }; + let node = T::node(info); + let (prev, next) = (node.prev, node.next); + node.prev = None; + node.next = None; + match prev { + Some(prev) => { + if let Some(info) = prev.get_info_mut(self.stage) { + T::node(info).next = next; + } + } + None => { + if let Some(info) = owner.get_info_mut(self.stage) { + T::list(info).head = next; + } + } + } + match next { + Some(next) => { + if let Some(info) = next.get_info_mut(self.stage) { + T::node(info).prev = prev; + } + } + None => { + if let Some(info) = owner.get_info_mut(self.stage) { + T::list(info).tail = prev; + } + } + } + if let Some(info) = owner.get_info_mut(self.stage) { + let list = T::list(info); + list.len = list.len.saturating_sub(1); + } + } + + /// Cut `at` and every item after it off `owner`'s list, returning them as + /// a chain, ready for [`link`](Self::link). Their parent is left for the + /// caller. + pub(super) fn split_off>(&mut self, owner: T::Owner, at: T) -> LinkedList { + let mut len = 0; + let mut current = Some(at); + while let Some(item) = current { + len += 1; + current = item + .get_info_mut(self.stage) + .and_then(|info| T::node(info).next); + } + let prev = match at.get_info_mut(self.stage) { + Some(info) => T::node(info).prev.take(), + None => return LinkedList::new(), + }; + if let Some(prev) = prev + && let Some(info) = prev.get_info_mut(self.stage) + { + T::node(info).next = None; + } + let Some(info) = owner.get_info_mut(self.stage) else { + return LinkedList::new(); + }; + let list = T::list(info); + let tail = list.tail; + if prev.is_none() { + list.head = None; + } + list.tail = prev; + list.len = list.len.saturating_sub(len); + LinkedList { + head: Some(at), + tail, + len, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::testing::dialect::{TestLang, new_stage}; + use crate::{GetInfo, StageInfo}; + + /// A block holding `count` `nop`s, and the `nop`s in order. + fn block_of_nops(count: usize) -> (StageInfo, Block, Vec) { + let mut stage = new_stage(); + let nops: Vec = (0..count) + .map(|_| stage.statement().definition(TestLang::Nop).new()) + .collect(); + let block = nops + .iter() + .fold(stage.block(), |block, &nop| block.stmt(nop)) + .new(); + (stage.finalize().unwrap(), block, nops) + } + + /// `block`'s statements in order, after checking that walking the list + /// backwards gives the same order, that its length is right, and that + /// every statement names `block` as its parent. + fn statements(stage: &StageInfo, block: Block) -> Vec { + let list = block.expect_info(stage).statements; + let mut forward = Vec::new(); + let mut current = list.head; + while let Some(stmt) = current { + forward.push(stmt); + current = stmt.expect_info(stage).node.next; + } + let mut backward = Vec::new(); + let mut current = list.tail; + while let Some(stmt) = current { + backward.push(stmt); + current = stmt.expect_info(stage).node.prev; + } + backward.reverse(); + assert_eq!(forward, backward); + assert_eq!(list.len, forward.len()); + for stmt in &forward { + assert_eq!( + stmt.expect_info(stage).parent, + Some(StatementParent::Block(block)) + ); + } + forward + } + + #[test] + fn link_places_a_single_item_at_the_head_middle_and_tail() { + let (mut stage, block, nops) = block_of_nops(4); + let [a, b, c, d] = nops[..] else { + unreachable!() + }; + let mut rewriter = Rewriter::new(&mut stage); + rewriter.unlink(a, block); + rewriter.unlink(c, block); + rewriter.unlink(d, block); + assert_eq!(statements(rewriter.stage(), block), [b]); + + rewriter.link(LinkedList::single(d), block, Some(b), None); + rewriter.link(LinkedList::single(a), block, None, Some(b)); + rewriter.link(LinkedList::single(c), block, Some(b), Some(d)); + assert_eq!(statements(rewriter.stage(), block), [a, b, c, d]); + } + + #[test] + fn unlinking_the_only_item_empties_the_list() { + let (mut stage, block, nops) = block_of_nops(1); + let mut rewriter = Rewriter::new(&mut stage); + rewriter.unlink(nops[0], block); + assert_eq!(statements(rewriter.stage(), block), []); + let list = block.expect_info(rewriter.stage()).statements; + assert_eq!((list.head, list.tail), (None, None)); + let node = nops[0].expect_info(rewriter.stage()).node; + assert_eq!((node.prev, node.next), (None, None)); + } + + #[test] + fn split_off_in_the_middle_keeps_the_front_and_returns_the_rest() { + let (mut stage, block, nops) = block_of_nops(4); + let [a, b, c, d] = nops[..] else { + unreachable!() + }; + let mut rewriter = Rewriter::new(&mut stage); + let rest = rewriter.split_off(block, c); + assert_eq!(statements(rewriter.stage(), block), [a, b]); + assert_eq!((rest.head, rest.tail, rest.len), (Some(c), Some(d), 2)); + + // Linked back at the end, the block is as it was. + rewriter.link(rest, block, Some(b), None); + assert_eq!(statements(rewriter.stage(), block), [a, b, c, d]); + } + + #[test] + fn split_off_at_the_head_takes_the_whole_list() { + let (mut stage, block, nops) = block_of_nops(3); + let mut rewriter = Rewriter::new(&mut stage); + let all = rewriter.split_off(block, nops[0]); + assert_eq!(statements(rewriter.stage(), block), []); + assert_eq!( + (all.head, all.tail, all.len), + (Some(nops[0]), Some(nops[2]), 3) + ); + + // Linked into an empty list, they become the whole list again. + rewriter.link(all, block, None, None); + assert_eq!(statements(rewriter.stage(), block), nops); + } +} diff --git a/crates/kirin-ir/src/rewrite/mod.rs b/crates/kirin-ir/src/rewrite/mod.rs index 9973e5405..ea8ab6c2e 100644 --- a/crates/kirin-ir/src/rewrite/mod.rs +++ b/crates/kirin-ir/src/rewrite/mod.rs @@ -23,6 +23,7 @@ mod alloc; mod bodies; pub(crate) mod cfg; pub(crate) mod digraph; +mod list; mod moves; pub(crate) mod pass; #[cfg(test)] diff --git a/crates/kirin-ir/src/rewrite/moves.rs b/crates/kirin-ir/src/rewrite/moves.rs index b2e711d64..1975acd02 100644 --- a/crates/kirin-ir/src/rewrite/moves.rs +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -72,12 +72,12 @@ impl<'a, L: Dialect> Rewriter<'a, L> { owner = tree::parent(self.stage, current); } - self.unlink_from_block(stmt, block); + self.unlink(stmt, block); let (block, prev, next) = self .slot_at(anchor, after) .expect("the anchor was checked above and moving `stmt` does not touch it"); debug_assert!(stmt.get_info(self.stage).is_some()); - self.link_into_block(stmt, block, prev, next); + self.link(LinkedList::single(stmt), block, prev, next); self.events.push(MutationEvent::MovedStatement { stmt }); Ok(()) } @@ -176,65 +176,9 @@ impl<'a, L: Dialect> Rewriter<'a, L> { Some(info) => std::mem::take(&mut info.statements), None => LinkedList::new(), }; - self.link_chain(chain, dest, prev, next); + self.link(chain, dest, prev, next); self.events .push(MutationEvent::SplicedBlock { source, dest }); Ok(()) } - - /// Link `chain`, a whole statement chain already taken out of its block, - /// into `block` between `prev` and `next` (the block head or tail when one - /// is `None`), making `block` the parent of each statement. - pub(super) fn link_chain( - &mut self, - chain: LinkedList, - block: Block, - prev: Option, - next: Option, - ) { - let (Some(head), Some(tail)) = (chain.head, chain.tail) else { - return; - }; - let mut current = Some(head); - while let Some(stmt) = current { - let Some(info) = stmt.get_info_mut(self.stage) else { - break; - }; - info.parent = Some(StatementParent::Block(block)); - current = info.node.next; - } - if let Some(info) = head.get_info_mut(self.stage) { - info.node.prev = prev; - } - if let Some(info) = tail.get_info_mut(self.stage) { - info.node.next = next; - } - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = Some(head); - } - } - None => { - if let Some(info) = block.get_info_mut(self.stage) { - info.statements.head = Some(head); - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = Some(tail); - } - } - None => { - if let Some(info) = block.get_info_mut(self.stage) { - info.statements.tail = Some(tail); - } - } - } - if let Some(info) = block.get_info_mut(self.stage) { - info.statements.len += chain.len; - } - } } diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index 8114b995f..986f10d93 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -84,6 +84,7 @@ use indexmap::{IndexMap, IndexSet}; use rustc_hash::{FxBuildHasher, FxHashMap}; use crate::arena::{GetInfo, GetInfoMut}; +use crate::node::linked_list::LinkedList; use smallvec::SmallVec; use super::bodies::bodies_of; @@ -1011,7 +1012,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { remove_predecessor(self.stage, target, block); } } else { - self.unlink_from_block(stmt, block); + self.unlink(stmt, block); } // Drop the def-use edges this statement contributed, erase the bodies @@ -1190,93 +1191,6 @@ impl<'a, L: Dialect> Rewriter<'a, L> { } } - /// Link the live, unlinked `stmt` into `block`'s chain between `prev` and - /// `next` (the block head or tail when one is `None`), making `block` its - /// parent. - pub(super) fn link_into_block( - &mut self, - stmt: Statement, - block: Block, - prev: Option, - next: Option, - ) { - if let Some(info) = stmt.get_info_mut(self.stage) { - info.parent = Some(StatementParent::Block(block)); - info.node.prev = prev; - info.node.next = next; - } - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = Some(stmt); - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.head = Some(stmt); - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = Some(stmt); - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.tail = Some(stmt); - } - } - } - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.len += 1; - } - } - - /// Unlink the non-terminator `stmt` from `block`'s chain, fixing its - /// neighbours (or the block head/tail at an end) and clearing its own - /// links. Its parent is left for the caller. - pub(super) fn unlink_from_block(&mut self, stmt: Statement, block: Block) { - let Some(node) = stmt.get_info(self.stage).map(|item| item.node) else { - return; - }; - let (prev, next) = (node.prev, node.next); - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = next; - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.head = next; - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = prev; - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.tail = prev; - } - } - } - if let Some(block_info) = block.get_info_mut(self.stage) - && block_info.statements.len > 0 - { - block_info.statements.len -= 1; - } - if let Some(info) = stmt.get_info_mut(self.stage) { - info.node.prev = None; - info.node.next = None; - } - } - fn insert_relative( &mut self, anchor: Statement, @@ -1301,7 +1215,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { &operands, ); - self.link_into_block(new_stmt, block, prev, next); + self.link(LinkedList::single(new_stmt), block, prev, next); // Predecessor mirror does not need to be updated since we rejected terminators earlier. self.events From 59cd5038e4858281090d4a874d3c892e794da4a4 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 14:38:42 -0500 Subject: [PATCH 10/14] docs(ir): name the errors moves and splices return for unknown ids Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/cfg/splice.rs | 3 ++- crates/kirin-ir/src/rewrite/moves.rs | 14 ++++++++------ 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/cfg/splice.rs b/crates/kirin-ir/src/rewrite/cfg/splice.rs index aba99f89b..4409d5c98 100644 --- a/crates/kirin-ir/src/rewrite/cfg/splice.rs +++ b/crates/kirin-ir/src/rewrite/cfg/splice.rs @@ -16,7 +16,8 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// for the caller to branch to. /// /// Rejects, leaving the stage untouched, if `source` is unknown - /// ([`RewriteError::UnknownCfg`]), if `anchor` is unknown or not in a CFG + /// ([`RewriteError::UnknownCfg`]), if `anchor` is unknown + /// ([`RewriteError::UnknownBlock`]) or not in a CFG /// ([`RewriteError::NotInCfgBody`]), or if `anchor` is in `source` or /// nested in one of its blocks ([`RewriteError::CfgSplicedIntoItself`]). /// diff --git a/crates/kirin-ir/src/rewrite/moves.rs b/crates/kirin-ir/src/rewrite/moves.rs index 1975acd02..08b70cf01 100644 --- a/crates/kirin-ir/src/rewrite/moves.rs +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -19,9 +19,10 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// `stmt`'s results still comes after it. /// /// Rejects, leaving the stage untouched, if `stmt` or `anchor` is unknown - /// or not in a block body ([`RewriteError::NotInBlockBody`]), if `stmt` is - /// a terminator ([`RewriteError::UnexpectedTerminator`]), or if `anchor` is - /// nested in `stmt`'s own bodies ([`RewriteError::MovedIntoItself`]). + /// ([`RewriteError::UnknownStatement`]) or not in a block body + /// ([`RewriteError::NotInBlockBody`]), if `stmt` is a terminator + /// ([`RewriteError::UnexpectedTerminator`]), or if `anchor` is nested in + /// `stmt`'s own bodies ([`RewriteError::MovedIntoItself`]). /// /// Records [`MutationEvent::MovedStatement`]. A walk visiting `stmt` /// reports the move (see [`crate::walk`]). @@ -93,13 +94,14 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// means at the end of that block's statement chain. A replacement may be /// one of `source`'s own arguments. /// - /// Rejects, leaving the stage untouched, if `source` or `anchor` is - /// unknown, if `anchor` is not in a block body + /// Rejects, leaving the stage untouched, if `source` is unknown + /// ([`RewriteError::UnknownBlock`]), if `anchor` is unknown + /// ([`RewriteError::UnknownStatement`]) or not in a block body /// ([`RewriteError::NotInBlockBody`]), if `anchor` is in `source` or /// nested in one of its statements ([`RewriteError::SplicedIntoItself`]), /// if the replacement count differs from `source`'s argument count /// ([`RewriteError::BlockArgumentArityMismatch`]), or if a replacement is - /// not live. + /// not live ([`RewriteError::UnknownValue`]). /// /// Records the [`Rewriter::replace_all_uses`] events for the argument /// redirection, then [`MutationEvent::SplicedBlock`]. A walk visiting a From 43ef3f6f275e782864975d0535c8a55d8a9db3c3 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 14:51:56 -0500 Subject: [PATCH 11/14] refactor(ir): tidy walk checks and assert unlink's preconditions in debug - share the depth-first cursors' visited-statement check - merge split imports and rewrap the Skip docs - debug-assert that unlink gets a live item from a non-empty list Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/list.rs | 2 ++ .../src/rewrite/walk/breadth_first.rs | 3 +-- .../kirin-ir/src/rewrite/walk/depth_first.rs | 23 ++++++++----------- crates/kirin-ir/src/rewrite/walk/mod.rs | 6 ++--- 4 files changed, 16 insertions(+), 18 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/list.rs b/crates/kirin-ir/src/rewrite/list.rs index 80bee755b..7f84143a0 100644 --- a/crates/kirin-ir/src/rewrite/list.rs +++ b/crates/kirin-ir/src/rewrite/list.rs @@ -126,6 +126,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// links. Its parent is left for the caller. pub(super) fn unlink>(&mut self, item: T, owner: T::Owner) { let Some(info) = item.get_info_mut(self.stage) else { + debug_assert!(false, "unlinking an item that is not live"); return; }; let node = T::node(info); @@ -158,6 +159,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { } if let Some(info) = owner.get_info_mut(self.stage) { let list = T::list(info); + debug_assert!(list.len > 0, "unlinking from an empty list"); list.len = list.len.saturating_sub(1); } } diff --git a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs index 4feaf0e28..f46a97ce2 100644 --- a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs @@ -8,8 +8,7 @@ use std::collections::VecDeque; use super::depth_first::{check_not_moved, live, visited_block}; -use super::order::Position; -use super::order::{CursorOrder, WalkCursor}; +use super::order::{CursorOrder, Position, WalkCursor}; use super::tree::{append_children, first_child, next_sibling, owns_bodies}; use super::{WalkError, WalkOrder}; use crate::{Block, Dialect, GetInfo, StageInfo, Statement}; diff --git a/crates/kirin-ir/src/rewrite/walk/depth_first.rs b/crates/kirin-ir/src/rewrite/walk/depth_first.rs index a8e5078f8..83d3faab6 100644 --- a/crates/kirin-ir/src/rewrite/walk/depth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/depth_first.rs @@ -1,7 +1,6 @@ //! Depth-first orders, walked through the IR's own links in O(1) extra space. -use super::order::Position; -use super::order::{CursorOrder, WalkCursor}; +use super::order::{CursorOrder, Position, WalkCursor}; use super::tree::{first_child, next_sibling, parent}; use super::{WalkError, WalkOrder}; use crate::{Block, Dialect, GetInfo, StageInfo, Statement, StatementParent}; @@ -83,10 +82,7 @@ impl WalkCursor for PreOrderCursor { } fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { - match self.position.visited() { - Some(visited) => check_in_place(stage, visited, self.visited_block), - None => Ok(()), - } + check_in_place(stage, self.position, self.visited_block) } fn skip(&mut self) { @@ -117,10 +113,7 @@ impl WalkCursor for PostOrderCursor { } fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { - match self.position.visited() { - Some(visited) => check_in_place(stage, visited, self.visited_block), - None => Ok(()), - } + check_in_place(stage, self.position, self.visited_block) } /// Nothing to skip: the children were visited before the statement. @@ -139,13 +132,17 @@ fn check_live(stage: &StageInfo, visited: Statement) -> Result<() .ok_or(WalkError::VisitedErased(visited)) } -/// Fail if the statement just visited was erased, or moved by merging -/// `block`, the block it was in when visited, into another block. +/// Fail if the statement the cursor is at was erased, or moved by merging +/// `block`, the block it was in when visited, into another block. Nothing to +/// check before the first step or after the last. fn check_in_place( stage: &StageInfo, - visited: Statement, + position: Position, block: Option, ) -> Result<(), WalkError> { + let Some(visited) = position.visited() else { + return Ok(()); + }; check_live(stage, visited)?; check_not_moved(stage, visited, block) } diff --git a/crates/kirin-ir/src/rewrite/walk/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index 2216f0ee3..317011966 100644 --- a/crates/kirin-ir/src/rewrite/walk/mod.rs +++ b/crates/kirin-ir/src/rewrite/walk/mod.rs @@ -40,9 +40,9 @@ //! //! - [`WalkResult::Advance`] goes on to the next statement. //! - [`WalkResult::Skip`] goes on without walking the statements nested in -//! this one. [`PreOrder`] steps past them and -//! [`BFS`] never queues them. The other orders have already visited them, -//! so there it is the same as `Advance`. +//! this one. [`PreOrder`] steps past them and [`BFS`] never queues them. +//! The other orders have already visited them, so there it is the same as +//! `Advance`. //! - [`WalkResult::Interrupt`] ends the walk, which reports //! [`WalkStatus::Interrupted`]. Edits already made stay. //! From 4000d9532f4e17813af53794ce7e4bda9f6fd84c Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 15:33:25 -0500 Subject: [PATCH 12/14] refactor(ir): name placement sides and share terminator handover - Side enum replaces the before/after bool in slot_at and its callers - Rewriter::set_terminator replaces duplicated code in merge and split - check_not_moved_since no longer generic over the walk order Co-Authored-By: Claude Opus 5.5 --- crates/kirin-ir/src/rewrite/cfg/helpers.rs | 13 ++++++++ crates/kirin-ir/src/rewrite/cfg/merge.rs | 11 ++---- crates/kirin-ir/src/rewrite/cfg/split.rs | 14 ++------ crates/kirin-ir/src/rewrite/moves.rs | 20 +++++------ crates/kirin-ir/src/rewrite/rewriter.rs | 39 +++++++++++++--------- crates/kirin-ir/src/rewrite/walk/order.rs | 13 +++----- 6 files changed, 56 insertions(+), 54 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/cfg/helpers.rs b/crates/kirin-ir/src/rewrite/cfg/helpers.rs index 2252e77d6..bd393846b 100644 --- a/crates/kirin-ir/src/rewrite/cfg/helpers.rs +++ b/crates/kirin-ir/src/rewrite/cfg/helpers.rs @@ -24,6 +24,19 @@ impl<'a, L: Dialect> Rewriter<'a, L> { } } + /// Make `terminator` the terminator of `block`, with `block` as its + /// parent, or leave `block` with none. + pub(super) fn set_terminator(&mut self, block: Block, terminator: Option) { + if let Some(terminator) = terminator + && let Some(info) = terminator.get_info_mut(self.stage) + { + info.parent = Some(StatementParent::Block(block)); + } + if let Some(info) = block.get_info_mut(self.stage) { + info.terminator = terminator; + } + } + /// The statement chain starting at `head`, following `next` links. pub(super) fn chain_from(&self, head: Option) -> Vec { let mut out = Vec::new(); diff --git a/crates/kirin-ir/src/rewrite/cfg/merge.rs b/crates/kirin-ir/src/rewrite/cfg/merge.rs index b93796ad2..452fd356a 100644 --- a/crates/kirin-ir/src/rewrite/cfg/merge.rs +++ b/crates/kirin-ir/src/rewrite/cfg/merge.rs @@ -1,6 +1,6 @@ use crate::arena::{GetInfo, GetInfoMut}; use crate::node::linked_list::LinkedList; -use crate::{Block, Dialect, SSAValue, StatementParent}; +use crate::{Block, Dialect, SSAValue}; use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter, remove_predecessor}; @@ -91,14 +91,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { }; let pred_tail = pred.expect_info(self.stage).statements.tail; self.link(succ_chain, pred, pred_tail, None); - if let Some(terminator) = succ_terminator - && let Some(info) = terminator.get_info_mut(self.stage) - { - info.parent = Some(StatementParent::Block(pred)); - } - if let Some(info) = pred.get_info_mut(self.stage) { - info.terminator = succ_terminator; - } + self.set_terminator(pred, succ_terminator); // succ is now empty: tombstoning it only removes the block and its // (already redirected) arguments. diff --git a/crates/kirin-ir/src/rewrite/cfg/split.rs b/crates/kirin-ir/src/rewrite/cfg/split.rs index e80328b54..15589f777 100644 --- a/crates/kirin-ir/src/rewrite/cfg/split.rs +++ b/crates/kirin-ir/src/rewrite/cfg/split.rs @@ -1,4 +1,3 @@ -use crate::arena::GetInfoMut; use crate::node::linked_list::LinkedList; use crate::{Block, Dialect, Statement, StatementParent}; @@ -48,17 +47,8 @@ impl<'a, L: Dialect> Rewriter<'a, L> { // after the original. let new_block = self.create_block(cfg, Some(block), next_block, []); self.link(moved, new_block, None, None); - if let Some(terminator) = terminator - && let Some(info) = terminator.get_info_mut(self.stage) - { - info.parent = Some(StatementParent::Block(new_block)); - } - if let Some(info) = new_block.get_info_mut(self.stage) { - info.terminator = terminator; - } - if let Some(info) = block.get_info_mut(self.stage) { - info.terminator = None; - } + self.set_terminator(new_block, terminator); + self.set_terminator(block, None); // The moved edges now leave from the new block; whatever stays in the // original block keeps its edges. diff --git a/crates/kirin-ir/src/rewrite/moves.rs b/crates/kirin-ir/src/rewrite/moves.rs index 08b70cf01..951158b78 100644 --- a/crates/kirin-ir/src/rewrite/moves.rs +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -1,6 +1,6 @@ use crate::arena::{GetInfo, GetInfoMut}; use crate::node::linked_list::LinkedList; -use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter}; +use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter, Side}; use crate::rewrite::walk::tree; use crate::{Block, Dialect, SSAValue, Statement, StatementParent}; @@ -27,7 +27,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// Records [`MutationEvent::MovedStatement`]. A walk visiting `stmt` /// reports the move (see [`crate::walk`]). pub fn move_before(&mut self, stmt: Statement, anchor: Statement) -> Result<(), RewriteError> { - self.move_relative(stmt, anchor, false) + self.move_relative(stmt, anchor, Side::Before) } /// Move `stmt` to immediately after `anchor`, which may be in any block @@ -36,14 +36,14 @@ impl<'a, L: Dialect> Rewriter<'a, L> { /// As [`Rewriter::move_before`], except that `anchor` must not be a /// terminator ([`RewriteError::AnchorIsTerminator`]): nothing follows one. pub fn move_after(&mut self, stmt: Statement, anchor: Statement) -> Result<(), RewriteError> { - self.move_relative(stmt, anchor, true) + self.move_relative(stmt, anchor, Side::After) } fn move_relative( &mut self, stmt: Statement, anchor: Statement, - after: bool, + side: Side, ) -> Result<(), RewriteError> { let block = { let item = self.live_statement(stmt)?; @@ -57,7 +57,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { }; // Checks the anchor; the place itself is read again once `stmt` is // unlinked, since `stmt` may be one of its neighbours. - let (_, prev, next) = self.slot_at(anchor, after)?; + let (_, prev, next) = self.slot_at(anchor, side)?; // `stmt` already borders the place (this includes `stmt == anchor`), // so unlinking it and linking it back would put it where it is. if prev == Some(stmt) || next == Some(stmt) { @@ -75,7 +75,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { self.unlink(stmt, block); let (block, prev, next) = self - .slot_at(anchor, after) + .slot_at(anchor, side) .expect("the anchor was checked above and moving `stmt` does not touch it"); debug_assert!(stmt.get_info(self.stage).is_some()); self.link(LinkedList::single(stmt), block, prev, next); @@ -112,7 +112,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { anchor: Statement, arg_replacements: &[SSAValue], ) -> Result<(), RewriteError> { - self.splice_relative(source, anchor, arg_replacements, false) + self.splice_relative(source, anchor, arg_replacements, Side::Before) } /// Move every statement of `source`, except its terminator, to @@ -127,7 +127,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { anchor: Statement, arg_replacements: &[SSAValue], ) -> Result<(), RewriteError> { - self.splice_relative(source, anchor, arg_replacements, true) + self.splice_relative(source, anchor, arg_replacements, Side::After) } fn splice_relative( @@ -135,7 +135,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { source: Block, anchor: Statement, arg_replacements: &[SSAValue], - after: bool, + side: Side, ) -> Result<(), RewriteError> { let args: Vec = self .live_block(source)? @@ -143,7 +143,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { .iter() .map(|&arg| arg.into()) .collect(); - let (dest, prev, next) = self.slot_at(anchor, after)?; + let (dest, prev, next) = self.slot_at(anchor, side)?; // Check `anchor` and every statement whose body encloses it. If one // of them is a statement of the block `source`, it would be spliced // next to or inside itself. diff --git a/crates/kirin-ir/src/rewrite/rewriter.rs b/crates/kirin-ir/src/rewrite/rewriter.rs index 986f10d93..b59dbbc20 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -895,6 +895,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { continue; }; let uses = info.uses_mut(); + // `sites` is a subset of `uses`, so equal lengths mean every use was selected. if sites.len() == uses.len() { uses.clear(); } else if !sites.is_empty() { @@ -1140,7 +1141,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { anchor: Statement, definition: L, ) -> Result { - self.insert_relative(anchor, definition, false) + self.insert_relative(anchor, definition, Side::Before) } /// Insert `definition` as a new statement immediately after `anchor`. @@ -1156,12 +1157,12 @@ impl<'a, L: Dialect> Rewriter<'a, L> { anchor: Statement, definition: L, ) -> Result { - self.insert_relative(anchor, definition, true) + self.insert_relative(anchor, definition, Side::After) } - /// Where a statement goes when placed before (`after == false`) or after - /// `anchor`: its block, and the chain statements that will sit directly - /// before and after it (`None` at either end of the chain). + /// Where a statement goes when placed on `side` of `anchor`: its block, + /// and the chain statements that will sit directly before and after it + /// (`None` at either end of the chain). /// /// The terminator sits outside the chain, so a place before it is the end /// of the chain, and nothing can follow it ([`RewriteError::AnchorIsTerminator`]). @@ -1169,35 +1170,36 @@ impl<'a, L: Dialect> Rewriter<'a, L> { pub(super) fn slot_at( &self, anchor: Statement, - after: bool, + side: Side, ) -> Result<(Block, Option, Option), RewriteError> { let item = self.live_statement(anchor)?; let block = match item.parent { Some(StatementParent::Block(block)) => block, _ => return Err(RewriteError::NotInBlockBody(anchor)), }; - match (item.definition.is_terminator(), after) { - (true, true) => Err(RewriteError::AnchorIsTerminator(anchor)), - (true, false) => { + let (prev, next) = match (item.definition.is_terminator(), side) { + (true, Side::After) => return Err(RewriteError::AnchorIsTerminator(anchor)), + (true, Side::Before) => { let tail = block .get_info(self.stage) .and_then(|b| b.statements.tail().copied()); // The tail ends the chain, so the statement becomes the tail. debug_assert!(tail.is_none_or(|t| t.next(self.stage).is_none())); - Ok((block, tail, None)) + (tail, None) } - (false, true) => Ok((block, Some(anchor), item.node.next)), - (false, false) => Ok((block, item.node.prev, Some(anchor))), - } + (false, Side::After) => (Some(anchor), item.node.next), + (false, Side::Before) => (item.node.prev, Some(anchor)), + }; + Ok((block, prev, next)) } fn insert_relative( &mut self, anchor: Statement, definition: L, - after: bool, + side: Side, ) -> Result { - let (block, prev, next) = self.slot_at(anchor, after)?; + let (block, prev, next) = self.slot_at(anchor, side)?; if definition.is_terminator() { return Err(RewriteError::UnexpectedTerminator); @@ -1577,6 +1579,13 @@ impl Rewriter<'_, L> { } } +/// Which side of an anchor statement an edit places a statement on. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum Side { + Before, + After, +} + /// Old-state snapshot taken by [`Rewriter::prepare_replacement`]. pub(super) struct PreparedReplacement { old_operands: Vec, diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index 968474a90..5bbf428ac 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -68,7 +68,9 @@ pub trait WalkOrder: CursorOrder { let seen = rewriter.events().len(); let place = VisitedPlace::of(rewriter.stage(), stmt); let result = visit(rewriter, stmt); - check_not_moved_since::(rewriter, stmt, place, seen)?; + if !Self::ALLOWS_ERASING_OR_MOVING_VISITED { + check_not_moved_since(rewriter, stmt, place, seen)?; + } // `None` on `Advance` and `Skip`, the final status on `Interrupt`. if let Some(status) = result.steer(&mut cursor) { // The cursor checks erases and merges of the visited @@ -83,21 +85,16 @@ pub trait WalkOrder: CursorOrder { } /// Fail with [`WalkError::VisitedMoved`] if an event recorded since the first -/// `seen` events of `rewriter` moved `visited`, in an order `O` that does not -/// allow it ([`WalkOrder::ALLOWS_ERASING_OR_MOVING_VISITED`]): a move of -/// `visited` itself, or a splice of the block or CFG it was in at `place` +/// `seen` events of `rewriter` moved `visited`: a move of `visited` itself, or a splice of the block or CFG it was in at `place` /// that took it along. A splice that left it behind, such as the terminator /// of a spliced block, or one that moved nothing, is not a move. /// Panics if `rewriter` has recorded fewer than `seen` events. -fn check_not_moved_since( +fn check_not_moved_since( rewriter: &Rewriter<'_, L>, visited: Statement, place: VisitedPlace, seen: usize, ) -> Result<(), WalkError> { - if O::ALLOWS_ERASING_OR_MOVING_VISITED { - return Ok(()); - } // Where `visited` is now. An erased one has no place; the cursor // reports its erasure. let now = || { From 711398a6f9fc2243cb1446df077c169111ff38b2 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 16:05:02 -0500 Subject: [PATCH 13/14] refactor(walk): check the walk contract in one place in walk_mut - walk_mut checks erasure and moves (including merges) after each visit - cursors no longer track the visited block; next() is infallible - rename check_visited to check_erasure Co-Authored-By: Claude Opus 5.5 --- .../src/rewrite/walk/breadth_first.rs | 42 ++++----- .../kirin-ir/src/rewrite/walk/depth_first.rs | 86 ++++-------------- crates/kirin-ir/src/rewrite/walk/order.rs | 87 ++++++++++--------- 3 files changed, 81 insertions(+), 134 deletions(-) diff --git a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs index f46a97ce2..b76025310 100644 --- a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs @@ -7,11 +7,11 @@ use std::collections::VecDeque; -use super::depth_first::{check_not_moved, live, visited_block}; +use super::depth_first::live; use super::order::{CursorOrder, Position, WalkCursor}; use super::tree::{append_children, first_child, next_sibling, owns_bodies}; use super::{WalkError, WalkOrder}; -use crate::{Block, Dialect, GetInfo, StageInfo, Statement}; +use crate::{Dialect, GetInfo, StageInfo, Statement}; /// Breadth-first by nesting depth: the root, then every statement nested /// directly in it, then every statement nested two levels down, and so on. @@ -60,8 +60,6 @@ pub(crate) struct BFSCursor { /// its own parent is at the front, so every parent of one level is queued /// before any of the next. parents: VecDeque, - /// The block the visited statement was in when visited. - visited_block: Option, } /// A [`ReverseBFS`] or [`KeepOrderReverseBFS`] walk in progress: the order @@ -90,20 +88,18 @@ impl CursorOrder for BFS { root, position: Position::Start, parents: VecDeque::new(), - visited_block: None, } } } impl WalkCursor for BFSCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { - self.check_visited(stage)?; + fn next(&mut self, stage: &StageInfo) -> Option { let next = match self.position { Position::Start => live(stage, self.root), Position::At(visited) | Position::Skipped(visited) => { - // Erased with an ancestor, which `check_visited` allowed. + // Erased with an ancestor, which `check_erasure` allowed. if visited.get_info(stage).is_none() { - return Ok(self.after_parent_erased(stage)); + return self.after_parent_erased(stage); } // Queue it as a parent for the next level, whose statements // are its children. Statements without bodies have none, @@ -131,24 +127,19 @@ impl WalkCursor for BFSCursor { Position::Done => None, }; self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, next); - Ok(next) + next } - /// A live visited statement must still be in its block: merging that - /// block into another moves it, and the walk would go on from its new - /// place. Moves and splices leave its block alive, so only - /// [`WalkOrder::walk_mut`](super::WalkOrder::walk_mut) catches them. - /// - /// An erased one is allowed only if erased together with an ancestor, - /// which kills the front parent: its remaining children are then dead - /// too. Erasing it alone would leave its siblings live but unreachable. - fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { + /// An erased visited statement is allowed only if erased together with + /// an ancestor, which kills the front parent: its remaining children are + /// then dead too. Erasing it alone would leave its siblings live but + /// unreachable. + fn check_erasure(&self, stage: &StageInfo) -> Result<(), WalkError> { let Some(visited) = self.position.visited() else { return Ok(()); }; if visited.get_info(stage).is_some() { - return check_not_moved(stage, visited, self.visited_block); + return Ok(()); } let parent_erased = visited != self.root && self @@ -176,7 +167,6 @@ impl BFSCursor { // drops them. let next = self.first_child_of_next_parent(stage); self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, next); next } @@ -245,21 +235,21 @@ impl BuiltOrderCursor { } impl WalkCursor for BuiltOrderCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + fn next(&mut self, stage: &StageInfo) -> Option { // Return the next statement still live, skipping any erased since the // order was built. while let Some(&stmt) = self.ordered_nodes.get(self.next) { self.next += 1; if stmt.get_info(stage).is_some() { - return Ok(Some(stmt)); + return Some(stmt); } } - Ok(None) + None } /// Nothing to check: the order was built at the start, so a visit may /// erase or move the visited statement. - fn check_visited(&self, _stage: &StageInfo) -> Result<(), WalkError> { + fn check_erasure(&self, _stage: &StageInfo) -> Result<(), WalkError> { Ok(()) } diff --git a/crates/kirin-ir/src/rewrite/walk/depth_first.rs b/crates/kirin-ir/src/rewrite/walk/depth_first.rs index 83d3faab6..a1056c8ea 100644 --- a/crates/kirin-ir/src/rewrite/walk/depth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/depth_first.rs @@ -3,7 +3,7 @@ use super::order::{CursorOrder, Position, WalkCursor}; use super::tree::{first_child, next_sibling, parent}; use super::{WalkError, WalkOrder}; -use crate::{Block, Dialect, GetInfo, StageInfo, Statement, StatementParent}; +use crate::{Dialect, GetInfo, StageInfo, Statement}; /// Depth-first, each statement before the statements nested in it. #[derive(Clone, Copy, Debug, Default)] @@ -18,8 +18,6 @@ pub struct PostOrder; pub(crate) struct PreOrderCursor { root: Statement, position: Position, - /// The block the visited statement was in when visited. - visited_block: Option, } /// A [`PostOrder`] walk in progress. @@ -27,8 +25,6 @@ pub(crate) struct PreOrderCursor { pub(crate) struct PostOrderCursor { root: Statement, position: Position, - /// The block the visited statement was in when visited. - visited_block: Option, } impl WalkOrder for PreOrder { @@ -42,7 +38,6 @@ impl CursorOrder for PreOrder { PreOrderCursor { root, position: Position::Start, - visited_block: None, } } } @@ -58,14 +53,12 @@ impl CursorOrder for PostOrder { PostOrderCursor { root, position: Position::Start, - visited_block: None, } } } impl WalkCursor for PreOrderCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { - self.check_visited(stage)?; + fn next(&mut self, stage: &StageInfo) -> Option { let next = match self.position { Position::Start => live(stage, self.root), Position::At(visited) => { @@ -77,12 +70,16 @@ impl WalkCursor for PreOrderCursor { Position::Done => None, }; self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, next); - Ok(next) + next } - fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { - check_in_place(stage, self.position, self.visited_block) + fn check_erasure(&self, stage: &StageInfo) -> Result<(), WalkError> { + match self.position.visited() { + Some(visited) if visited.get_info(stage).is_none() => { + Err(WalkError::VisitedErased(visited)) + } + _ => Ok(()), + } } fn skip(&mut self) { @@ -91,8 +88,7 @@ impl WalkCursor for PreOrderCursor { } impl WalkCursor for PostOrderCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { - self.check_visited(stage)?; + fn next(&mut self, stage: &StageInfo) -> Option { let next = match self.position { Position::Start => live(stage, self.root).map(|root| deepest_first(stage, root)), // Never `Skipped`: `skip` does nothing in this order. @@ -108,12 +104,16 @@ impl WalkCursor for PostOrderCursor { Position::Done => None, }; self.position = next.map_or(Position::Done, Position::At); - self.visited_block = visited_block(stage, next); - Ok(next) + next } - fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError> { - check_in_place(stage, self.position, self.visited_block) + fn check_erasure(&self, stage: &StageInfo) -> Result<(), WalkError> { + match self.position.visited() { + Some(visited) if visited.get_info(stage).is_none() => { + Err(WalkError::VisitedErased(visited)) + } + _ => Ok(()), + } } /// Nothing to skip: the children were visited before the statement. @@ -125,54 +125,6 @@ pub(super) fn live(stage: &StageInfo, stmt: Statement) -> Option< stmt.get_info(stage).map(|_| stmt) } -/// Fail if the statement just visited was erased. -fn check_live(stage: &StageInfo, visited: Statement) -> Result<(), WalkError> { - live(stage, visited) - .map(|_| ()) - .ok_or(WalkError::VisitedErased(visited)) -} - -/// Fail if the statement the cursor is at was erased, or moved by merging -/// `block`, the block it was in when visited, into another block. Nothing to -/// check before the first step or after the last. -fn check_in_place( - stage: &StageInfo, - position: Position, - block: Option, -) -> Result<(), WalkError> { - let Some(visited) = position.visited() else { - return Ok(()); - }; - check_live(stage, visited)?; - check_not_moved(stage, visited, block) -} - -/// Fail if the live statement just visited left `block`, the block it was -/// in when visited, because `block` was merged into another block. Only a -/// merge erases a block and keeps its statements, and splitting a block -/// keeps the original live, so a split passes. -pub(super) fn check_not_moved( - stage: &StageInfo, - visited: Statement, - block: Option, -) -> Result<(), WalkError> { - match block { - Some(block) if block.get_info(stage).is_none() => Err(WalkError::VisitedMoved(visited)), - _ => Ok(()), - } -} - -/// The block `next` is in, for [`check_not_moved`] at the next step. -pub(super) fn visited_block( - stage: &StageInfo, - next: Option, -) -> Option { - match next?.get_info(stage)?.parent? { - StatementParent::Block(block) => Some(block), - StatementParent::DiGraph(_) | StatementParent::UnGraph(_) => None, - } -} - /// The first statement after `stmt`'s subtree within `root`'s: the next /// sibling of `stmt` or of its nearest ancestor that has one, stopping at /// `root`. diff --git a/crates/kirin-ir/src/rewrite/walk/order.rs b/crates/kirin-ir/src/rewrite/walk/order.rs index 5bbf428ac..743356066 100644 --- a/crates/kirin-ir/src/rewrite/walk/order.rs +++ b/crates/kirin-ir/src/rewrite/walk/order.rs @@ -16,7 +16,7 @@ use crate::{ #[expect( private_bounds, reason = "`CursorOrder` is crate-private so other crates cannot step a cursor by hand: \ - that would skip `walk_mut`'s check that the visited statement was not moved" + that would skip `walk_mut`'s checks of the walk contract" )] pub trait WalkOrder: CursorOrder { /// Whether an edit may erase or move the statement being visited (see @@ -30,7 +30,7 @@ pub trait WalkOrder: CursorOrder { fn schedule(stage: &StageInfo, root: Statement) -> Vec { let mut cursor = Self::start(stage, root); let mut ordered_nodes = Vec::new(); - while let Ok(Some(stmt)) = cursor.next(stage) { + while let Some(stmt) = cursor.next(stage) { ordered_nodes.push(stmt); } ordered_nodes @@ -45,7 +45,7 @@ pub trait WalkOrder: CursorOrder { mut visit: impl FnMut(&StageInfo, Statement) -> WalkResult, ) -> WalkStatus { let mut cursor = Self::start(stage, root); - while let Ok(Some(stmt)) = cursor.next(stage) { + while let Some(stmt) = cursor.next(stage) { if let Some(status) = visit(stage, stmt).steer(&mut cursor) { return status; } @@ -64,19 +64,18 @@ pub trait WalkOrder: CursorOrder { mut visit: impl FnMut(&mut Rewriter<'r, L>, Statement) -> WalkResult, ) -> Result { let mut cursor = Self::start(rewriter.stage(), root); - while let Some(stmt) = cursor.next(rewriter.stage())? { + while let Some(stmt) = cursor.next(rewriter.stage()) { let seen = rewriter.events().len(); let place = VisitedPlace::of(rewriter.stage(), stmt); let result = visit(rewriter, stmt); + // Check the contract before stepping on, so the cursor never steps + // from a statement erased or moved where the order forbids it. + cursor.check_erasure(rewriter.stage())?; if !Self::ALLOWS_ERASING_OR_MOVING_VISITED { check_not_moved_since(rewriter, stmt, place, seen)?; } // `None` on `Advance` and `Skip`, the final status on `Interrupt`. if let Some(status) = result.steer(&mut cursor) { - // The cursor checks erases and merges of the visited - // statement only when it steps; the walk ends here, so check - // them now. - cursor.check_visited(rewriter.stage())?; return Ok(status); } } @@ -84,34 +83,40 @@ pub trait WalkOrder: CursorOrder { } } -/// Fail with [`WalkError::VisitedMoved`] if an event recorded since the first -/// `seen` events of `rewriter` moved `visited`: a move of `visited` itself, or a splice of the block or CFG it was in at `place` +/// Fail with [`WalkError::VisitedMoved`] if the edits since the first `seen` +/// events of `rewriter` moved `visited` from `place`: a merge of its block +/// into another, a move of `visited` itself, or a splice of its block or CFG /// that took it along. A splice that left it behind, such as the terminator -/// of a spliced block, or one that moved nothing, is not a move. -/// Panics if `rewriter` has recorded fewer than `seen` events. +/// of a spliced block, or one that moved nothing, is not a move; nor is a +/// split, which keeps the original block live. An erased `visited` passes: +/// erasure is [`WalkCursor::check_erasure`]'s to judge. Panics if `rewriter` +/// has recorded fewer than `seen` events. fn check_not_moved_since( rewriter: &Rewriter<'_, L>, visited: Statement, place: VisitedPlace, seen: usize, ) -> Result<(), WalkError> { - // Where `visited` is now. An erased one has no place; the cursor - // reports its erasure. - let now = || { - visited - .get_info(rewriter.stage()) - .map(|_| VisitedPlace::of(rewriter.stage(), visited)) - }; - let moved = rewriter.events()[seen..].iter().any(|event| match *event { - MutationEvent::MovedStatement { stmt } => stmt == visited, - MutationEvent::SplicedBlock { source, .. } => { - place.block == Some(source) && now().is_some_and(|now| now.block != place.block) - } - MutationEvent::SplicedCfg { source, .. } => { - place.cfg == Some(source) && now().is_some_and(|now| now.cfg != place.cfg) - } - _ => false, - }); + let stage = rewriter.stage(); + if visited.get_info(stage).is_none() { + return Ok(()); + } + // Only a merge erases a block and keeps its statements. + let merged = place + .block + .is_some_and(|block| block.get_info(stage).is_none()); + let now = VisitedPlace::of(stage, visited); + let moved = merged + || rewriter.events()[seen..].iter().any(|event| match *event { + MutationEvent::MovedStatement { stmt } => stmt == visited, + MutationEvent::SplicedBlock { source, .. } => { + place.block == Some(source) && now.block != place.block + } + MutationEvent::SplicedCfg { source, .. } => { + place.cfg == Some(source) && now.cfg != place.cfg + } + _ => false, + }); if moved { Err(WalkError::VisitedMoved(visited)) } else { @@ -120,9 +125,10 @@ fn check_not_moved_since( } /// Where a visited statement was when its visit started: its block, and that -/// block's CFG. [`check_not_moved_since`] needs it for splices: it names the -/// splices that could have moved the statement, and comparing it with where -/// the statement is afterwards tells whether one did. +/// block's CFG. [`check_not_moved_since`] needs it for merges, which erase +/// the block, and for splices: it names the splices that could have moved the +/// statement, and comparing it with where the statement is afterwards tells +/// whether one did. #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct VisitedPlace { block: Option, @@ -188,7 +194,7 @@ pub enum WalkStatus { /// Crate-private, so that outside this crate a walk is only /// [`WalkOrder::walk`], [`WalkOrder::walk_mut`] or [`WalkOrder::schedule`]: /// an editing walk then always goes through `walk_mut`, which checks the -/// whole contract, while a cursor stepped by hand would miss moves. +/// whole contract, while a cursor stepped by hand checks none of it. pub(crate) trait CursorOrder { /// The cursor walking this order. type Cursor: WalkCursor; @@ -204,16 +210,15 @@ pub(crate) trait WalkCursor { /// The next statement to visit, or `None` once the walk is done. /// /// The IR may be edited between calls within the walk's contract (see - /// the [module docs](super)). Fails if the statement returned by the - /// previous call has since been erased where the order does not allow - /// it, or moved by merging its block into another. - fn next(&mut self, stage: &StageInfo) -> Result, WalkError>; + /// the [module docs](super)). The cursor does not check that contract: + /// [`WalkOrder::walk_mut`], the only walk that edits, checks it after + /// each visit, before stepping on. + fn next(&mut self, stage: &StageInfo) -> Option; /// Fail if the statement returned by the previous [`next`](Self::next) - /// has since been erased where the order does not allow it, or moved by - /// merging its block into another. Only reads; the cursor does not move. - /// [`next`](Self::next) runs this check first. - fn check_visited(&self, stage: &StageInfo) -> Result<(), WalkError>; + /// has since been erased where the order does not allow it. Only reads; + /// the cursor does not move. + fn check_erasure(&self, stage: &StageInfo) -> Result<(), WalkError>; /// Do not walk the statements nested in the statement the previous /// [`next`](Self::next) returned, including any added while visiting it. From a14d0dfa1342a2960dc0019b1acebdeac6f48707 Mon Sep 17 00:00:00 2001 From: Alex Gatea Date: Wed, 11 Nov 2026 16:08:09 -0500 Subject: [PATCH 14/14] docs(plans): remove remaining IR operations plan Co-Authored-By: Claude Opus 5.5 --- .../2026-11-04-remaining-ir-operations.md | 259 ------------------ 1 file changed, 259 deletions(-) delete mode 100644 docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md diff --git a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md b/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md deleted file mode 100644 index 981b05017..000000000 --- a/docs/plans/kirin-ir/2026-11-04-remaining-ir-operations.md +++ /dev/null @@ -1,259 +0,0 @@ -# Remaining IR operations in the `kirin-ir` `Rewriter` - -Branch: `alex/remaining-ir-operations`, off `rust`. - -Adds the `Rewriter` operations needed for code motion and inlining: moving -statements and blocks, splicing whole bodies, copying bodies, and (as an -attempt) inlining a call. - -## Background - -### What exists - -| Body | Operations | -|---|---| -| block body | `erase_statement`, `insert_before` / `insert_after` (new statements only), `replace_statement`, `insert_terminator` | -| CFG | `append_block`, `insert_block_before` / `_after` (new empty blocks only), `erase_block`, `split_block`, `merge_blocks`, `retarget_successor`, `insert_block_arg`, `erase_block_arg` | -| values | `replace_operand`, `replace_all_uses`, `replace_results`, `new_result` | -| `DiGraph` | `insert_node`, `replace_yield` (removal through `erase_statement`) | -| `UnGraph` | `insert_ungraph_node`, `insert_ungraph_edge` (removal through `erase_statement`) | - -Nothing moves an existing statement or block, and nothing copies IR. The one -operation that moves IR is `merge_blocks`, which appends one block's -statements to the end of another block of the same CFG. - -### Cost model - -Each statement stores its parent block (`StatementParent::Block`), and each -block stores its parent CFG (`BlockParent::CFG`). Values store their defining -statement or block, and uses store the using statement. So: - -- moving a **statement** costs O(1): relink it, change its parent; -- moving a **block** costs O(1), and its statements don't change; -- moving the **statements of a block** into another block costs O(number - moved), one parent write each (this is why `split_block` costs O(statements - after the split point)); -- moving a **whole CFG's blocks** costs O(number of blocks). - -Ids never change when IR moves, so uses stay valid for free. A copy creates new -ids and must map every operand, which costs O(size of what is copied). - -### Comparison with MLIR - -| MLIR (`RewriterBase`, `Region`, `Operation`) | Kirin today | This plan | -|---|---|---| -| `moveOpBefore` / `moveOpAfter` | none | step 2 | -| `Block::moveBefore` | none | not in this plan (step 3, dropped) | -| `inlineBlockBefore(source, dest, before, args)` | `merge_blocks` (end of block, same CFG only) | step 4 | -| `inlineRegionBefore(region, parent, before)` | none | step 5 | -| `Operation::clone`, `Region::cloneInto` with `IRMapping` | none | step 6 | -| `inlineCall` + `DialectInlinerInterface` | none | step 7 (attempt) | -| `replaceUsesWithIf`, `replaceAllUsesExcept` | `replace_all_uses` only | step 1 | -| `Operation::isBeforeInBlock` | none | not in this plan (dominance comes later) | -| `Value::setType` | none | not in this plan | -| `replaceOpWithNewOp` | `replace_statement` (same result count) | not needed | -| `eraseOpResults` | `remove_result` deferred | not in this plan | - -## Steps (pause for review after each) - -Each step adds the operation to `Rewriter`, its rejections as `RewriteError` -variants where needed, its `MutationEvent`, tests in `kirin-ir`, and its row in -the operation table of `docs/design/rewrite-engine.md`. - -### 1. Replace only some uses - -```rust -rewriter.replace_uses_if(old, new, |stage: &StageInfo, site: Use| ...)?; -rewriter.replace_all_uses_except(old, new, stmt)?; -``` - -MLIR's `replaceUsesWithIf` / `replaceAllUsesExcept`. Like `replace_all_uses`, -but only the use sites the predicate accepts, or all but those of `stmt`. The -predicate gets the stage, so it can look up the using statement's parent or -neighbors; it runs on every use before anything is changed. CSE -and code motion need this when only some uses may change, e.g. only the uses -after a point. Same events as `replace_all_uses`, counting only the changed -slots. - -Cost: O(uses of `old`). - -### 2. Move a statement - -```rust -rewriter.move_before(stmt, anchor)?; -rewriter.move_after(stmt, anchor)?; -``` - -Same anchor rules as `insert_before` / `insert_after`: `move_before` a -terminator appends to the end of that block's statements; nothing can go after -a terminator. The anchor may be in any block, of any CFG or statement-owned -body. - -Example, hoisting out of an `scf.if` arm: - -``` -^entry: ^entry: - %r = if %c then { %t = neg %x // moved - %t = neg %x ==> %r = if %c then { - yield %t yield %t - } ... } ... -``` - -Rejects: -- `stmt` or `anchor` unknown; -- `stmt` is a terminator (terminators are not in the statement chain); -- `stmt` or `anchor` is a graph member; -- the anchor is inside `stmt`'s own nested bodies (moving an `scf.if` into its - own arm would make it contain itself). - -Not checked, like `insert_before`: that every use of `stmt`'s results still -comes after it. That is the caller's job. - -Records a new `MutationEvent::MovedStatement { stmt }`. `ErasedStatement` + -`InsertedStatement` would wrongly say the statement was tombstoned. - -**Walk check.** Moving the visited statement breaks pre-order, post-order and -BFS, which find the next statement through the visited one. Today the walk -only notices the `merge_blocks` case (the visited statement's old block was -erased). This step adds a check for moves, read from the `Rewriter`'s event -log, so nothing new is stored in the IR: - -- Every operation that moves IR records an event naming what moved: - `MovedStatement { stmt }` here, and the events of steps 3–5. -- `WalkCursor` gets a variant of `next` that takes the `Rewriter` instead of - the stage. At each visit the cursor remembers the visited statement, its - block, and the length of the event log. On its next step it scans only the - events recorded since, and reports `WalkError::VisitedMoved` if one moved the - visited statement or its block. When nothing was recorded, that is one - length comparison. -- `walk_mut` uses the new variant. The read-only `walk` and `schedule` keep - `next(stage)`, since nothing moves during them. -- The reverse BFS orders allow moves (their order is built at the start), so - they skip the check. -- `split_block` records no move event, so the split the contract allows still - passes. Moving an ancestor of the visited statement is not checked. - -Cost: O(1), plus O(nesting depth) for the cycle check. - -### 3. Move a block within its CFG (dropped) - -Not added. A CFG's block order has one semantic effect, choosing the entry -block (the first in the list); otherwise it only changes printing and -iteration order. A new entry block can already be made with -`insert_block_before` plus a branch, and steps 5 and 7 place blocks while -splicing, so nothing needs a separate move. Add it when block layout starts -to matter (e.g. a backend or an analysis seeded in reverse postorder). - -### 4. Splice a block's statements into another block - -```rust -rewriter.splice_block_before(source, anchor, &arg_replacements)?; -rewriter.splice_block_after(source, anchor, &arg_replacements)?; -``` - -MLIR's `inlineBlockBefore`. Moves `source`'s statements (not its terminator) -next to `anchor`, and redirects uses of `source`'s arguments to -`arg_replacements`. The anchor rules are those of `move_before` / -`move_after`. `splice_block_after` is for step 7, where the callee's entry is -spliced after the call, the last statement of an unterminated block. `source` is left empty, keeping its arguments and -terminator, and is not erased: the caller decides what to do with it. -Generalizes `merge_blocks`: any position, in any block of any CFG or -statement-owned body. Records a new `SplicedBlock { source, dest }`, which the -walk check reads as "every statement of `source` moved". Afterwards a spliced -statement's parent is `dest`, so the check also takes the visited statement's -parent from before the visit, and compares it with `source`. - -Rejects an anchor in `source` or nested in one of its statements -(`SplicedIntoItself`). - -Example, folding `scf.if` with a known condition: splice the `then` arm in -front of the `if`, replace the `if`'s results with the arm's `yield` values, -erase the `if`. - -Cost: O(statements moved). - -### 5. Splice a whole CFG into another - -```rust -rewriter.splice_cfg_after(source, anchor_block)?; -``` - -MLIR's `inlineRegionBefore`. Moves every block of `source` into `anchor_block`'s -CFG, right after `anchor_block`. Branches between the moved blocks stay valid, -since they all move together. `source` is left empty and not erased; its -owner (e.g. the callee function) still names it, and the caller decides -whether to erase the owner. Records `ChangedBlockList` for both CFGs and a new -`SplicedCfg { source, dest }`, which the walk check reads as "every block of -`source` moved". Afterwards a moved block's parent is `dest`, so, as in step 4, -the check compares against where the visited statement was before the visit: -a `VisitedPlace` holding its block and that block's CFG. -`ChangedBlockList` alone could not drive the check, since `split_block`, which -a walk allows, records it too. - -Rejects an anchor block in `source` or nested in one of its blocks -(`CfgSplicedIntoItself`). There is no `_before` variant and no way to splice -into an empty CFG, which has no anchor; inlining needs neither. - -Cost: O(blocks moved). This is the move-instead-of-copy path for inlining a -callee's last use. - -### 6. Copy a statement or a CFG - -```rust -let mut map = IrMap::default(); // MLIR's IRMapping -let copy = rewriter.clone_before(stmt, anchor, &mut map)?; -let blocks = rewriter.clone_cfg_after(source, anchor_block, &mut map)?; -``` - -A deep copy with fresh ids: statements, results, blocks, block arguments and -nested bodies, including `DiGraph` and `UnGraph` bodies. A copied graph gets -new ports, copied members and yields, and its topology is rebuilt from the -copied operands with the helpers `insert_node` / `insert_ungraph_node` -already use. Every operand and branch target inside the copy is mapped -through `map`. A value from outside the copied IR keeps pointing at the -original. The caller can pre-seed `map`, e.g. to map the callee's entry -arguments to the call's operands. `map` afterwards holds old → new for every -value and block copied, so the caller can find the copied `ret`s. - -Cost: O(size of what is copied). - -### 7. Inline a call (attempt) - -An inliner in `kirin-function`, which knows its own call and `ret` ops, built -from steps 4–6 and the existing `kirin-ir` operations. It takes a closure that -builds the branch op, since `kirin-function` doesn't know which branch op a -language uses. It inlines `CallSpecialized` calls within the same stage; the -other call kinds first need resolving to a specialization. For a call `c` in block `B`: - -1. `split_block` after `c`. The continuation block `C` gets the rest of `B`. -2. Copy the callee's CFG between `B` and `C`, mapping its entry arguments to - `c`'s operands (step 6). For the callee's last use, splice instead (step 5). -3. Splice the copied entry block into `B` (step 4), or branch to it. -4. Replace each copied `ret v` with a branch to `C` passing `v`. -5. Give `C` one argument per result of `c`, redirect `c`'s results to them, and - erase `c`. - -Cost per call: O(statements after `c` in `B`) for the split, plus O(callee -size) for the copy. Inlining a block's calls from last to first makes the -splits cost O(block size) in total, since each statement moves at most once. - -### 8. Docs - -`Rewriter` module docs, the design doc's operation table, and the -`MutationEvent` count in Gap 3. The `kirin_ir::walk` contract gets a line on -the new moves. - -## Decisions - -1. Moves record new events naming what moved: `MovedStatement`, - `SplicedBlock`, `SplicedCfg`. -2. The walk checks for a moved visited statement or block by reading the - `Rewriter`'s event log (step 2). Ancestors are not checked. -3. Moving a block within its CFG is not added (step 3). -4. A block's statements can be spliced into any block. -5. Splicing leaves the source block or CFG empty and does not erase it. -6. Copying includes graph bodies. -7. Inlining lives in `kirin-function`, with a closure building the branch. -8. Only `CallSpecialized` within the same stage is inlined. -9. Of the extra MLIR operations, only `replace_uses_if` / - `replace_all_uses_except` (step 1). Their predicate gets the stage.