Rewrites work on any HugrMut (#282)

acl-cqc · web-flow · commit cc3bd6a29bfd · 2023-09-04T17:50:13.000Z
Rewrite::verify takes &amp;impl HugrView (not &amp;Hugr); Rewrite::apply takes
&amp;mut impl HugrMut (not &amp;mut Hugr).

TransactionalRewrite backs up the portion inside the HugrMut and then
restores - this feels a little awkward but may be ok, still not tested
though.
diff --git a/src/hugr/rewrite.rs b/src/hugr/rewrite.rs
@@ -3,11 +3,12 @@
 pub mod insert_identity;
 pub mod outline_cfg;
 pub mod simple_replace;
-use std::mem;
 
-use crate::Hugr;
+use crate::{Hugr, HugrView};
 pub use simple_replace::{SimpleReplacement, SimpleReplacementError};
 
+use super::HugrMut;
+
 /// An operation that can be applied to mutate a Hugr
 pub trait Rewrite {
     /// The type of Error with which this Rewrite may fail
@@ -21,7 +22,7 @@ pub trait Rewrite {
     /// Checks whether the rewrite would succeed on the specified Hugr.
     /// If this call succeeds, [self.apply] should also succeed on the same `h`
     /// If this calls fails, [self.apply] would fail with the same error.
-    fn verify(&self, h: &Hugr) -> Result<(), Self::Error>;
+    fn verify(&self, h: &impl HugrView) -> Result<(), Self::Error>;
 
     /// Mutate the specified Hugr, or fail with an error.
     /// Returns [`Self::ApplyResult`] if successful.
@@ -31,7 +32,7 @@ pub trait Rewrite {
     /// May panic if-and-only-if `h` would have failed [Hugr::validate]; that is,
     /// implementations may begin with `assert!(h.validate())`, with `debug_assert!(h.validate())`
     /// being preferred.
-    fn apply(self, h: &mut Hugr) -> Result<Self::ApplyResult, Self::Error>;
+    fn apply(self, h: &mut impl HugrMut) -> Result<Self::ApplyResult, Self::Error>;
 }
 
 /// Wraps any rewrite into a transaction (i.e. that has no effect upon failure)
@@ -46,19 +47,28 @@ impl<R: Rewrite> Rewrite for Transactional<R> {
     type ApplyResult = R::ApplyResult;
     const UNCHANGED_ON_FAILURE: bool = true;
 
-    fn verify(&self, h: &Hugr) -> Result<(), Self::Error> {
+    fn verify(&self, h: &impl HugrView) -> Result<(), Self::Error> {
         self.underlying.verify(h)
     }
 
-    fn apply(self, h: &mut Hugr) -> Result<Self::ApplyResult, Self::Error> {
+    fn apply(self, h: &mut impl HugrMut) -> Result<Self::ApplyResult, Self::Error> {
         if R::UNCHANGED_ON_FAILURE {
             return self.underlying.apply(h);
         }
-        let backup = h.clone();
+        // Try to backup just the contents of this HugrMut.
+        let mut backup = Hugr::new(h.root_type().clone());
+        backup.insert_from_view(backup.root(), h).unwrap();
         let r = self.underlying.apply(h);
+        fn first_child(h: &impl HugrView) -> Option<crate::Node> {
+            h.children(h.root()).next()
+        }
         if r.is_err() {
-            // drop the old h, it was undefined
-            let _ = mem::replace(h, backup);
+            // Try to restore backup.
+            h.replace_op(h.root(), backup.root_type().clone());
+            while let Some(child) = first_child(h) {
+                h.remove_node(child).unwrap();
+            }
+            h.insert_from_view(h.root(), &backup).unwrap();
         }
         r
     }
diff --git a/src/hugr/rewrite/insert_identity.rs b/src/hugr/rewrite/insert_identity.rs
@@ -3,7 +3,7 @@
 use crate::hugr::{HugrMut, Node};
 use crate::ops::{LeafOp, OpTag, OpTrait};
 use crate::types::EdgeKind;
-use crate::{Direction, Hugr, HugrView, Port};
+use crate::{Direction, HugrView, Port};
 
 use super::Rewrite;
 
@@ -52,7 +52,7 @@ impl Rewrite for IdentityInsertion {
     /// The inserted node.
     type ApplyResult = Node;
     const UNCHANGED_ON_FAILURE: bool = true;
-    fn verify(&self, _h: &Hugr) -> Result<(), IdentityInsertionError> {
+    fn verify(&self, _h: &impl HugrView) -> Result<(), IdentityInsertionError> {
         /*
         Assumptions:
         1. Value kind inputs can only have one connection.
@@ -65,7 +65,7 @@ impl Rewrite for IdentityInsertion {
 
         unimplemented!()
     }
-    fn apply(self, h: &mut Hugr) -> Result<Self::ApplyResult, IdentityInsertionError> {
+    fn apply(self, h: &mut impl HugrMut) -> Result<Self::ApplyResult, IdentityInsertionError> {
         if self.post_port.direction() != Direction::Incoming {
             return Err(IdentityInsertionError::PortIsOutput);
         }
@@ -78,6 +78,7 @@ impl Rewrite for IdentityInsertion {
         let (pre_node, pre_port) = h
             .linked_ports(self.post_node, self.post_port)
             .exactly_one()
+            .ok()
             .expect("Value kind input can only have one connection.");
 
         h.disconnect(self.post_node, self.post_port).unwrap();
diff --git a/src/hugr/rewrite/outline_cfg.rs b/src/hugr/rewrite/outline_cfg.rs
@@ -6,12 +6,11 @@ use thiserror::Error;
 
 use crate::builder::{BlockBuilder, Container, Dataflow, SubContainer};
 use crate::extension::PRELUDE_REGISTRY;
-use crate::hugr::hugrmut::sealed::HugrMutInternals;
 use crate::hugr::rewrite::Rewrite;
 use crate::hugr::{HugrMut, HugrView};
 use crate::ops;
 use crate::ops::{BasicBlock, OpTag, OpTrait, OpType};
-use crate::{type_row, Hugr, Node};
+use crate::{type_row, Node};
 
 /// Moves part of a Control-flow Sibling Graph into a new CFG-node
 /// that is the only child of a new Basic Block in the original CSG.
@@ -27,7 +26,10 @@ impl OutlineCfg {
         }
     }
 
-    fn compute_entry_exit_outside(&self, h: &Hugr) -> Result<(Node, Node, Node), OutlineCfgError> {
+    fn compute_entry_exit_outside(
+        &self,
+        h: &impl HugrView,
+    ) -> Result<(Node, Node, Node), OutlineCfgError> {
         let cfg_n = match self
             .blocks
             .iter()
@@ -86,11 +88,11 @@ impl Rewrite for OutlineCfg {
     type ApplyResult = ();
 
     const UNCHANGED_ON_FAILURE: bool = true;
-    fn verify(&self, h: &Hugr) -> Result<(), OutlineCfgError> {
+    fn verify(&self, h: &impl HugrView) -> Result<(), OutlineCfgError> {
         self.compute_entry_exit_outside(h)?;
         Ok(())
     }
-    fn apply(self, h: &mut Hugr) -> Result<(), OutlineCfgError> {
+    fn apply(self, h: &mut impl HugrMut) -> Result<(), OutlineCfgError> {
         let (entry, exit, outside) = self.compute_entry_exit_outside(h)?;
         // 1. Compute signature
         // These panic()s only happen if the Hugr would not have passed validate()
@@ -128,12 +130,13 @@ impl Rewrite for OutlineCfg {
             .children(new_block)
             .filter(|n| h.get_optype(*n).tag() == OpTag::Cfg)
             .exactly_one()
+            .ok() // HugrMut::Children is not Debug
             .unwrap();
-        let inner_exit = h.children(cfg_node).exactly_one().unwrap();
+        let inner_exit = h.children(cfg_node).exactly_one().ok().unwrap();
 
         // 4. Entry edges. Change any edges into entry_block from outside, to target new_block
         let preds: Vec<_> = h
-            .linked_ports(entry, h.node_inputs(entry).exactly_one().unwrap())
+            .linked_ports(entry, h.node_inputs(entry).exactly_one().ok().unwrap())
             .collect();
         for (pred, br) in preds {
             if !self.blocks.contains(&pred) {
@@ -168,6 +171,7 @@ impl Rewrite for OutlineCfg {
                 t == outside
             })
             .exactly_one()
+            .ok() // NodePorts does not implement Debug
             .unwrap();
         h.disconnect(exit, exit_port).unwrap();
         h.connect(exit, exit_port.index(), inner_exit, 0).unwrap();
diff --git a/src/hugr/rewrite/simple_replace.rs b/src/hugr/rewrite/simple_replace.rs
@@ -55,11 +55,11 @@ impl Rewrite for SimpleReplacement {
 
     const UNCHANGED_ON_FAILURE: bool = true;
 
-    fn verify(&self, _h: &Hugr) -> Result<(), SimpleReplacementError> {
+    fn verify(&self, _h: &impl HugrView) -> Result<(), SimpleReplacementError> {
         unimplemented!()
     }
 
-    fn apply(self, h: &mut Hugr) -> Result<(), SimpleReplacementError> {
+    fn apply(self, h: &mut impl HugrMut) -> Result<(), SimpleReplacementError> {
         // 1. Check the parent node exists and is a DataflowParent.
         if !OpTag::DataflowParent.is_superset(h.get_optype(self.parent).tag()) {
             return Err(SimpleReplacementError::InvalidParentNode());
@@ -114,6 +114,7 @@ impl Rewrite for SimpleReplacement {
                 let (rem_inp_pred_node, rem_inp_pred_port) = h
                     .linked_ports(*rem_inp_node, *rem_inp_port)
                     .exactly_one()
+                    .ok() // PortLinks does not implement Debug
                     .unwrap();
                 h.disconnect(*rem_inp_node, *rem_inp_port).unwrap();
                 let new_inp_node = index_map.get(rep_inp_node).unwrap();
@@ -155,6 +156,7 @@ impl Rewrite for SimpleReplacement {
                 let (rem_inp_pred_node, rem_inp_pred_port) = h
                     .linked_ports(*rem_inp_node, *rem_inp_port)
                     .exactly_one()
+                    .ok() // PortLinks does not implement Debug
                     .unwrap();
                 h.disconnect(*rem_inp_node, *rem_inp_port).unwrap();
                 h.disconnect(*rem_out_node, *rem_out_port).unwrap();
