feat(ir): move and splice statements, blocks and CFGs in Rewriter

crates/kirin-ir/src/node/linked_list.rs

@@ -48,6 +48,15 @@ impl<Ptr: Copy + PartialEq> LinkedList<Ptr> {
         }
     }
 
+    /// 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()
     }

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 });

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<Statement>) {
+        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<Statement>) -> Vec<Statement> {
         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;
-        }
-    }
 }

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);

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;

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(())
+    }
+}

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<Block, RewriteError> {
-        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.