Re-spec and Implement OutlineCfg rewrite

specification/hugr.md

@@ -1226,11 +1226,26 @@ nodes as children.
 
 ###### `OutlineCFG`
 
-Replace a CFG-convex subgraph (of sibling BasicBlock nodes) having a
-single entry node with a single BasicBlock node having a CFG node child
-which has as its children the original BasicBlock nodes and an exit node
-that has inputs coming from all edges of the original CFG that don’t
-terminate in it.
+Replace a set of CFG sibling nodes with a single BasicBlock node having a
+CFG node child which has as its children the set of BasicBlock nodes
+originally specified. The set of basic blocks must satisfy constraints:
+* There must be at most one node in the set with incoming (controlflow) edges
+ from nodes outside the set. Specifically,
+   * *either* the set includes the CFG's entry node, and any edges from outside
+     the set (there may be none or more) target said entry node;
+   * *or* the set does not include the CFG's entry node, but contains exactly one
+     node which is the target of at least one edge(s) from outside the set.
+* The set may not include the Exit block.
+* There must be exactly one edge from a node in the set to a node outside it.
+
+Situations in which multiple nodes have edges leaving the set, or where the Exit block
+would be in the set, can be converted to this form by a combination of InsertIdentity
+operations and one Replace. For example, rather than moving the Exit block into the nested CFG:
+1. An Identity node with a single successor can be added onto each edge into the Exit
+2. If there is more than one edge into the Exit, these Identity nodes can then all be combined
+   by a Replace operation changing them all for a single Identity (keeping the same number
+   of in-edges, but reducing to one out-edge to the Exit).
+3. The single edge to the Exit node can then be used as the exiting edge.
 
 ##### Inlining methods
 

src/algorithm/nest_cfgs.rs

@@ -471,22 +471,15 @@ pub(crate) mod test {
         //      \-> right -/     \-<--<-/
         // Here we would like two consecutive regions, but there is no *edge* between
         // the conditional and the loop to indicate the boundary, so we cannot separate them.
-        let mut cfg_builder = CFGBuilder::new(type_row![NAT], type_row![NAT])?;
-        let mut entry = cfg_builder.simple_entry_builder(type_row![NAT], 2)?;
-        let pred_const = entry.add_constant(ConstValue::simple_predicate(0, 2))?; // Nothing here cares which
-        let const_unit = entry.add_constant(ConstValue::simple_unary_predicate())?;
-
-        let entry = n_identity(entry, &pred_const)?;
-        let merge = build_then_else_merge_from_if(&mut cfg_builder, &const_unit, entry)?;
-        let tail = build_loop_from_header(&mut cfg_builder, &pred_const, merge)?;
-        cfg_builder.branch(&merge, 0, &tail)?; // trivial "loop body"
-        let exit = cfg_builder.exit_block();
-        cfg_builder.branch(&tail, 0, &exit)?;
-
-        let h = cfg_builder.finish_hugr()?;
-
-        let (entry, exit) = (entry.node(), exit.node());
+        let (h, merge, tail) = build_cond_then_loop_cfg(false)?;
         let (merge, tail) = (merge.node(), tail.node());
+        let [entry, exit]: [Node; 2] = h
+            .children(h.root())
+            .take(2)
+            .collect_vec()
+            .try_into()
+            .unwrap();
+
         let edge_classes = EdgeClassifier::get_edge_classes(&SimpleCfgView::new(&h));
         let [&left,&right] = edge_classes.keys().filter(|(s,_)| *s == entry).map(|(_,t)|t).collect::<Vec<_>>()[..] else {panic!("Entry node should have two successors");};
 
@@ -513,23 +506,14 @@ pub(crate) mod test {
         //  entry -> head -> split            > merge -> tail -> exit
         //             |          \-> right -/             |
         //             \---<---<---<---<---<---<---<---<---/
+        // split is unique successor of head
+        let split = h.output_neighbours(head).exactly_one().unwrap();
+        // merge is unique predecessor of tail
+        let merge = h.input_neighbours(tail).exactly_one().unwrap();
+
         let v = SimpleCfgView::new(&h);
         let edge_classes = EdgeClassifier::get_edge_classes(&v);
         let SimpleCfgView { h: _, entry, exit } = v;
-        // split is unique successor of head
-        let split = *edge_classes
-            .keys()
-            .filter(|(s, _)| *s == head)
-            .map(|(_, t)| t)
-            .exactly_one()
-            .unwrap();
-        // merge is unique predecessor of tail
-        let merge = *edge_classes
-            .keys()
-            .filter(|(_, t)| *t == tail)
-            .map(|(s, _)| s)
-            .exactly_one()
-            .unwrap();
         let [&left,&right] = edge_classes.keys().filter(|(s,_)| *s == split).map(|(_,t)|t).collect::<Vec<_>>()[..] else {panic!("Split node should have two successors");};
         let classes = group_by(edge_classes);
         assert_eq!(
@@ -655,6 +639,36 @@ pub(crate) mod test {
         Ok((header, tail))
     }
 
+    // Result is merge and tail; loop header is (merge, if separate==true; unique successor of merge, if separate==false)
+    pub fn build_cond_then_loop_cfg(
+        separate: bool,
+    ) -> Result<(Hugr, BasicBlockID, BasicBlockID), BuildError> {
+        let mut cfg_builder = CFGBuilder::new(type_row![NAT], type_row![NAT])?;
+        let mut entry = cfg_builder.simple_entry_builder(type_row![NAT], 2)?;
+        let pred_const = entry.add_constant(ConstValue::simple_predicate(0, 2))?; // Nothing here cares which
+        let const_unit = entry.add_constant(ConstValue::simple_unary_predicate())?;
+
+        let entry = n_identity(entry, &pred_const)?;
+        let merge = build_then_else_merge_from_if(&mut cfg_builder, &const_unit, entry)?;
+        let head = if separate {
+            let h = n_identity(
+                cfg_builder.simple_block_builder(type_row![NAT], type_row![NAT], 1)?,
+                &const_unit,
+            )?;
+            cfg_builder.branch(&merge, 0, &h)?;
+            h
+        } else {
+            merge
+        };
+        let tail = build_loop_from_header(&mut cfg_builder, &pred_const, head)?;
+        cfg_builder.branch(&head, 0, &tail)?; // trivial "loop body"
+        let exit = cfg_builder.exit_block();
+        cfg_builder.branch(&tail, 0, &exit)?;
+
+        let h = cfg_builder.finish_hugr()?;
+        Ok((h, merge, tail))
+    }
+
     // Build a CFG, returning the Hugr
     pub fn build_conditional_in_loop_cfg(
         separate_headers: bool,

src/builder/cfg.rs

@@ -1,3 +1,5 @@
+use itertools::Itertools;
+
 use super::{
     build_traits::SubContainer,
     dataflow::{DFGBuilder, DFGWrapper},
@@ -95,6 +97,22 @@ impl<B: AsMut<Hugr> + AsRef<Hugr>> CFGBuilder<B> {
             inputs: Some(input),
         })
     }
+
+    /// Create a CFGBuilder for an existing CFG node (that already has entry + exit nodes)
+    pub(crate) fn from_existing(base: B, cfg_node: Node) -> Result<Self, BuildError> {
+        let OpType::CFG(crate::ops::controlflow::CFG {outputs, ..}) = base.get_optype(cfg_node)
+            else {return Err(BuildError::UnexpectedType{node: cfg_node, op_desc: "Any CFG"});};
+        let n_out_wires = outputs.len();
+        let (_, exit_node) = base.children(cfg_node).take(2).collect_tuple().unwrap();
+        Ok(Self {
+            base,
+            cfg_node,
+            inputs: None, // This will prevent creating an entry node
+            exit_node,
+            n_out_wires,
+        })
+    }
+
     /// Return a builder for a non-entry [`BasicBlock::DFB`] child graph with `inputs`
     /// and `outputs` and the variants of the branching predicate Sum value
     /// specified by `predicate_variants`.
@@ -278,6 +296,8 @@ impl BlockBuilder<Hugr> {
 
 #[cfg(test)]
 mod test {
+    use std::collections::HashSet;
+
     use cool_asserts::assert_matches;
 
     use crate::builder::build_traits::HugrBuilder;
@@ -319,6 +339,36 @@ mod test {
 
         Ok(())
     }
+    #[test]
+    fn from_existing() -> Result<(), BuildError> {
+        let mut cfg_builder = CFGBuilder::new(type_row![NAT], type_row![NAT])?;
+        build_basic_cfg(&mut cfg_builder)?;
+        let h = cfg_builder.finish_hugr()?;
+
+        let mut new_builder = CFGBuilder::from_existing(h.clone(), h.root())?;
+        assert_matches!(new_builder.simple_entry_builder(type_row![NAT], 1), Err(_));
+        let h2 = new_builder.finish_hugr()?;
+        assert_eq!(h, h2); // No new nodes added
+
+        let mut new_builder = CFGBuilder::from_existing(h.clone(), h.root())?;
+        let block_builder = new_builder.simple_block_builder(
+            vec![SimpleType::new_simple_predicate(1), NAT].into(),
+            type_row![NAT],
+            1,
+        )?;
+        let new_bb = block_builder.container_node();
+        let [pred, nat]: [Wire; 2] = block_builder.input_wires_arr();
+        block_builder.finish_with_outputs(pred, [nat])?;
+        let h2 = new_builder.finish_hugr()?;
+        let expected_nodes = h
+            .children(h.root())
+            .chain([new_bb])
+            .collect::<HashSet<Node>>();
+        assert_eq!(expected_nodes, HashSet::from_iter(h2.children(h2.root())));
+
+        Ok(())
+    }
+
     fn build_basic_cfg<T: AsMut<Hugr> + AsRef<Hugr>>(
         cfg_builder: &mut CFGBuilder<T>,
     ) -> Result<(), BuildError> {

src/hugr/rewrite.rs

@@ -1,5 +1,6 @@
 //! Rewrite operations on the HUGR - replacement, outlining, etc.
 
+pub mod outline_cfg;
 pub mod simple_replace;
 use std::mem;