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..bd393846b 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, @@ -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(); @@ -203,11 +216,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 +230,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 da06fad96..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}; @@ -85,48 +85,17 @@ 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; - } - } - 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; - } + let pred_tail = pred.expect_info(self.stage).statements.tail; + self.link(succ_chain, pred, pred_tail, None); + self.set_terminator(pred, succ_terminator); // 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/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..4409d5c98 --- /dev/null +++ b/crates/kirin-ir/src/rewrite/cfg/splice.rs @@ -0,0 +1,61 @@ +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 + /// ([`RewriteError::UnknownBlock`]) 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(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(()) + } +} diff --git a/crates/kirin-ir/src/rewrite/cfg/split.rs b/crates/kirin-ir/src/rewrite/cfg/split.rs index feb667694..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}; @@ -21,12 +20,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,48 +34,21 @@ 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)); - } - } - 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) { - info.terminator = None; - } + self.link(moved, new_block, None, 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/list.rs b/crates/kirin-ir/src/rewrite/list.rs new file mode 100644 index 000000000..7f84143a0 --- /dev/null +++ b/crates/kirin-ir/src/rewrite/list.rs @@ -0,0 +1,315 @@ +//! 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 { + debug_assert!(false, "unlinking an item that is not live"); + 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); + debug_assert!(list.len > 0, "unlinking from an empty list"); + 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 df443ca50..ea8ab6c2e 100644 --- a/crates/kirin-ir/src/rewrite/mod.rs +++ b/crates/kirin-ir/src/rewrite/mod.rs @@ -23,6 +23,8 @@ mod alloc; mod bodies; pub(crate) mod cfg; pub(crate) mod digraph; +mod list; +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..951158b78 --- /dev/null +++ b/crates/kirin-ir/src/rewrite/moves.rs @@ -0,0 +1,186 @@ +use crate::arena::{GetInfo, GetInfoMut}; +use crate::node::linked_list::LinkedList; +use crate::rewrite::rewriter::{MutationEvent, RewriteError, Rewriter, Side}; +use crate::rewrite::walk::tree; +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 + /// 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 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. + /// + /// Rejects, leaving the stage untouched, if `stmt` or `anchor` is unknown + /// ([`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`]). + pub fn move_before(&mut self, stmt: Statement, anchor: Statement) -> Result<(), RewriteError> { + self.move_relative(stmt, anchor, Side::Before) + } + + /// 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, Side::After) + } + + fn move_relative( + &mut self, + stmt: Statement, + anchor: Statement, + side: Side, + ) -> 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. + 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) { + 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(stmt, block); + let (block, prev, next) = self + .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); + 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` 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 ([`RewriteError::UnknownValue`]). + /// + /// 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, Side::Before) + } + + /// 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, Side::After) + } + + fn splice_relative( + &mut self, + source: Block, + anchor: Statement, + arg_replacements: &[SSAValue], + side: Side, + ) -> Result<(), RewriteError> { + let args: Vec = self + .live_block(source)? + .arguments + .iter() + .map(|&arg| arg.into()) + .collect(); + 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. + 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, dest, prev, next); + self.events + .push(MutationEvent::SplicedBlock { source, dest }); + Ok(()) + } +} 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 b50a57c22..b59dbbc20 100644 --- a/crates/kirin-ir/src/rewrite/rewriter.rs +++ b/crates/kirin-ir/src/rewrite/rewriter.rs @@ -23,10 +23,13 @@ //! 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` +//! [`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 @@ -81,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; @@ -137,6 +141,16 @@ 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 }, + /// 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`]. @@ -160,11 +174,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 @@ -224,6 +239,18 @@ 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 }, + /// `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 }, + /// `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), @@ -299,7 +326,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) => { @@ -375,6 +402,24 @@ 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::SplicedIntoItself { source, anchor } => { + write!( + f, + "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}") } @@ -481,8 +526,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) } @@ -601,6 +648,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 +783,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 +884,25 @@ 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(); + // `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() { + 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 +933,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 }); } } @@ -832,7 +976,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, @@ -840,8 +984,6 @@ impl<'a, L: Dialect> Rewriter<'a, L> { }; ( block, - item.node.prev, - item.node.next, item.definition .arguments() .copied() @@ -871,37 +1013,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(stmt, block); } // Drop the def-use edges this statement contributed, erase the bodies @@ -1029,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`. @@ -1045,40 +1157,49 @@ 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 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`]). + /// Rejects an unknown `anchor` or one not in a block body. + pub(super) fn slot_at( + &self, + anchor: Statement, + 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)), + }; + 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())); + (tail, None) + } + (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) = { - 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)), - } - }; + let (block, prev, next) = self.slot_at(anchor, side)?; if definition.is_terminator() { return Err(RewriteError::UnexpectedTerminator); @@ -1096,34 +1217,7 @@ impl<'a, L: Dialect> Rewriter<'a, L> { &operands, ); - // Wire neighbours (or the block head/tail when inserting at an end). - match prev { - Some(p) => { - if let Some(info) = p.get_info_mut(self.stage) { - info.node.next = Some(new_stmt); - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.head = Some(new_stmt); - } - } - } - match next { - Some(n) => { - if let Some(info) = n.get_info_mut(self.stage) { - info.node.prev = Some(new_stmt); - } - } - None => { - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.tail = Some(new_stmt); - } - } - } - if let Some(block_info) = block.get_info_mut(self.stage) { - block_info.statements.len += 1; - } + self.link(LinkedList::single(new_stmt), block, prev, next); // Predecessor mirror does not need to be updated since we rejected terminators earlier. self.events @@ -1485,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/breadth_first.rs b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs index 7be45879a..b76025310 100644 --- a/crates/kirin-ir/src/rewrite/walk/breadth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/breadth_first.rs @@ -7,16 +7,34 @@ use std::collections::VecDeque; -use super::depth_first::{check_not_moved, live, visited_block}; -use super::order::Position; +use super::depth_first::live; +use super::order::{CursorOrder, Position, WalkCursor}; use super::tree::{append_children, first_child, next_sibling, owns_bodies}; -use super::{WalkCursor, WalkError, WalkOrder}; -use crate::{Block, Dialect, GetInfo, StageInfo, Statement}; +use super::{WalkError, WalkOrder}; +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. /// 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 +50,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 @@ -61,8 +60,6 @@ pub 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 @@ -73,35 +70,37 @@ 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_ERASING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { BFSCursor { root, position: Position::Start, parents: VecDeque::new(), - visited_block: None, } } } impl WalkCursor for BFSCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + 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_erasure` allowed. if visited.get_info(stage).is_none() { - return self.after_erased(stage, visited); + return 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 +127,47 @@ 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); - Ok(next) + 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> { + /// 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 Ok(()); + } 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) + next } /// The first child of the front parent or a later one, dropping parents @@ -182,8 +189,11 @@ impl BFSCursor { } impl WalkOrder for ReverseBFS { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = true; +} + +impl CursorOrder for ReverseBFS { type Cursor = BuiltOrderCursor; - const ALLOWS_ERASING_VISITED: bool = true; fn start(stage: &StageInfo, root: Statement) -> Self::Cursor { let mut ordered_nodes = breadth_first(stage, root, |_| {}); @@ -193,8 +203,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_ERASING_VISITED: bool = true; fn start(stage: &StageInfo, root: Statement) -> Self::Cursor { let mut level_starts = Vec::new(); @@ -222,16 +235,22 @@ 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_erasure(&self, _stage: &StageInfo) -> Result<(), WalkError> { + Ok(()) } /// Nothing to skip: the order visits a statement after the statements diff --git a/crates/kirin-ir/src/rewrite/walk/depth_first.rs b/crates/kirin-ir/src/rewrite/walk/depth_first.rs index 8b10559ad..a1056c8ea 100644 --- a/crates/kirin-ir/src/rewrite/walk/depth_first.rs +++ b/crates/kirin-ir/src/rewrite/walk/depth_first.rs @@ -1,9 +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, Position, WalkCursor}; use super::tree::{first_child, next_sibling, parent}; -use super::{WalkCursor, WalkError, WalkOrder}; -use crate::{Block, Dialect, GetInfo, StageInfo, Statement, StatementParent}; +use super::{WalkError, WalkOrder}; +use crate::{Dialect, GetInfo, StageInfo, Statement}; /// Depth-first, each statement before the statements nested in it. #[derive(Clone, Copy, Debug, Default)] @@ -15,67 +15,71 @@ 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. - visited_block: Option, } /// 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. - visited_block: Option, } impl WalkOrder for PreOrder { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = false; +} + +impl CursorOrder for PreOrder { type Cursor = PreOrderCursor; - const ALLOWS_ERASING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { PreOrderCursor { root, position: Position::Start, - visited_block: None, } } } impl WalkOrder for PostOrder { + const ALLOWS_ERASING_OR_MOVING_VISITED: bool = false; +} + +impl CursorOrder for PostOrder { type Cursor = PostOrderCursor; - const ALLOWS_ERASING_VISITED: bool = false; fn start(_stage: &StageInfo, root: Statement) -> Self::Cursor { PostOrderCursor { root, position: Position::Start, - visited_block: None, } } } impl WalkCursor for PreOrderCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + fn next(&mut self, stage: &StageInfo) -> Option { 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); - Ok(next) + next + } + + 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) { @@ -84,12 +88,11 @@ impl WalkCursor for PreOrderCursor { } impl WalkCursor for PostOrderCursor { - fn next(&mut self, stage: &StageInfo) -> Result, WalkError> { + 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. 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,8 +104,16 @@ 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); - Ok(next) + next + } + + 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. @@ -114,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 just visited was erased, or moved by merging -/// `block`, the block it was in when visited, into another block. -fn check_in_place( - stage: &StageInfo, - visited: Statement, - block: Option, -) -> Result<(), WalkError> { - 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. `None` -/// for the root, whose place outside its subtree the walk never reads, so -/// moving it is harmless. -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? { - 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/mod.rs b/crates/kirin-ir/src/rewrite/walk/mod.rs index 2c3baca0d..317011966 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 -//! [`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. //! -//! 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 @@ -74,6 +71,18 @@ //! - 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`]; +//! - 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 (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; +//! - 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. @@ -87,14 +96,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; -mod tree; +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 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 f80fc3afc..743356066 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, Rewriter, StageInfo, Statement}; +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. /// @@ -9,24 +13,24 @@ use crate::{Dialect, Rewriter, StageInfo, Statement}; /// [`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; - - /// Whether an edit may erase the statement being visited, as the - /// [contract](super#contract) allows for this order. - const ALLOWS_ERASING_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; +#[expect( + private_bounds, + reason = "`CursorOrder` is crate-private so other crates cannot step a cursor by hand: \ + 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 + /// 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_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. 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 @@ -41,7 +45,7 @@ pub trait WalkOrder { 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; } @@ -51,16 +55,27 @@ 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. + /// `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, 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())? { - if let Some(status) = visit(rewriter, stmt).steer(&mut cursor) { + 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) { return Ok(status); } } @@ -68,14 +83,84 @@ pub trait WalkOrder { } } +/// 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; 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> { + 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 { + Ok(()) + } +} + +/// Where a visited statement was when its visit started: its block, and that +/// 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, + cfg: Option, +} + +impl VisitedPlace { + /// Where `stmt` is now in `stage`. + 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 { /// 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, @@ -104,16 +189,36 @@ 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 checks none of it. +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 - /// 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. 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. @@ -130,10 +235,17 @@ 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), or by splicing + /// its block with + /// [`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). VisitedMoved(Statement), } @@ -149,7 +261,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" ) } } @@ -169,6 +281,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/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/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 1b351c0cf..5777d9024 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. @@ -1003,10 +1004,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.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`. + 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) { @@ -1079,10 +1092,801 @@ 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. 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.events().to_vec() + }; + + 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.events().to_vec(); + 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.events().to_vec() + }; + 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.events().to_vec() + }; + 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(); +} + +/// ```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, + 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]); +} + +#[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()); +} + +/// ```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/rewriter_cfg.rs b/crates/kirin-ir/tests/rewriter_cfg.rs index b4d9b9ac4..1940c326a 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!( @@ -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(); @@ -161,7 +162,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 +280,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 +309,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 +319,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 +417,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 +560,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 +696,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 +977,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 +1099,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, @@ -1246,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/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 c98e9beb5..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!( @@ -814,3 +824,40 @@ 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 mut log = EventCursor::default(); + 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))])); + log.advance(&rw); + + // 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!( + log.read_and_advance(&rw), + 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/crates/kirin-ir/tests/walk.rs b/crates/kirin-ir/tests/walk.rs index bc1547e68..d6d506b07 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,15 +785,325 @@ fn reverse_bfs_orders_visit_statements_where_they_were_before_a_merge() { assert_eq!(result.unwrap(), without_br); } +/// 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_editing::(&mut t.stage, t.root, |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_editing::(&mut t.stage, t.root, |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)); + } +} + +/// 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_editing::(&mut t.stage, t.root, |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))); + } +} + +/// 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(); + 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_editing::(&mut t.stage, t.root, |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::()); +} + +/// ```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>) +{ + let mut t = tree(); + let a = t.a; + let schedule = O::schedule(&t.stage, t.root); + let result = walk_editing::(&mut t.stage, t.root, |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..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_ERASING_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 892445a34..c4fdf3c60 100644 --- a/docs/design/rewrite-engine.md +++ b/docs/design/rewrite-engine.md @@ -1114,6 +1114,9 @@ pub enum MutationEvent { ChangedDiGraphYields { graph: DiGraph }, ChangedUnGraphTopology { graph: UnGraph }, ChangedBodyNesting { owner: Statement }, + MovedStatement { stmt: Statement }, + SplicedBlock { source: Block, dest: Block }, + SplicedCfg { source: Cfg, dest: Cfg }, } ``` @@ -1749,8 +1752,24 @@ 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 splice_cfg_after(&mut self, source: CFG, anchor: Block) -> Result<(), RewriteError>; + pub fn merge_blocks( &mut self, pred: Block, @@ -1777,7 +1796,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; } ``` @@ -1792,9 +1811,13 @@ 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 | +| `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 | @@ -1805,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) | @@ -1968,9 +1992,9 @@ 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` | +| `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` | @@ -1982,7 +2006,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` | @@ -2004,12 +2028,17 @@ 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` | +| `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 `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 }`, +`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 @@ -2857,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_ERASING_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. @@ -3034,12 +3063,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