From ea8e818f63d32e650531ff95ea97b40e4fabcb69 Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Wed, 2 Oct 2024 10:57:39 +0200
Subject: [PATCH] feat!: `HugrMut::remove_node` and `SimpleReplacement` return
 removed weights (#1516)

Closes #476

BREAKING CHANGE: `remove_node` now returns an OpType, and the
`ApplyResult` of `SimpleReplacement` holds replaced nodes and weights.
---
 hugr-core/src/hugr/hugrmut.rs                |  8 ++++----
 hugr-core/src/hugr/rewrite/replace.rs        |  4 +++-
 hugr-core/src/hugr/rewrite/simple_replace.rs | 16 +++++++++-------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/hugr-core/src/hugr/hugrmut.rs b/hugr-core/src/hugr/hugrmut.rs
index 16f8749b2..b608d632f 100644
--- a/hugr-core/src/hugr/hugrmut.rs
+++ b/hugr-core/src/hugr/hugrmut.rs
@@ -105,13 +105,13 @@ pub trait HugrMut: HugrMutInternals {
         self.hugr_mut().add_node_after(sibling, op)
     }
 
-    /// Remove a node from the graph.
+    /// Remove a node from the graph and return the node weight.
     ///
     /// # Panics
     ///
     /// If the node is not in the graph, or if the node is the root node.
     #[inline]
-    fn remove_node(&mut self, node: Node) {
+    fn remove_node(&mut self, node: Node) -> OpType {
         panic_invalid_non_root(self, node);
         self.hugr_mut().remove_node(node)
     }
@@ -264,11 +264,11 @@ impl<T: RootTagged<RootHandle = Node> + AsMut<Hugr>> HugrMut for T {
         node
     }
 
-    fn remove_node(&mut self, node: Node) {
+    fn remove_node(&mut self, node: Node) -> OpType {
         panic_invalid_non_root(self, node);
         self.as_mut().hierarchy.remove(node.pg_index());
         self.as_mut().graph.remove_node(node.pg_index());
-        self.as_mut().op_types.remove(node.pg_index());
+        self.as_mut().op_types.take(node.pg_index())
     }
 
     fn connect(
diff --git a/hugr-core/src/hugr/rewrite/replace.rs b/hugr-core/src/hugr/rewrite/replace.rs
index 979b0d928..dcb171efe 100644
--- a/hugr-core/src/hugr/rewrite/replace.rs
+++ b/hugr-core/src/hugr/rewrite/replace.rs
@@ -324,7 +324,9 @@ impl Rewrite for Replacement {
         }
 
         // 7. Remove remaining nodes
-        to_remove.into_iter().for_each(|n| h.remove_node(n));
+        to_remove.into_iter().for_each(|n| {
+            h.remove_node(n);
+        });
         Ok(node_map)
     }
 
diff --git a/hugr-core/src/hugr/rewrite/simple_replace.rs b/hugr-core/src/hugr/rewrite/simple_replace.rs
index a9872be0a..3018adce0 100644
--- a/hugr-core/src/hugr/rewrite/simple_replace.rs
+++ b/hugr-core/src/hugr/rewrite/simple_replace.rs
@@ -55,14 +55,14 @@ impl SimpleReplacement {
 
 impl Rewrite for SimpleReplacement {
     type Error = SimpleReplacementError;
-    type ApplyResult = ();
+    type ApplyResult = Vec<(Node, OpType)>;
     const UNCHANGED_ON_FAILURE: bool = true;
 
     fn verify(&self, _h: &impl HugrView) -> Result<(), SimpleReplacementError> {
         unimplemented!()
     }
 
-    fn apply(mut self, h: &mut impl HugrMut) -> Result<(), SimpleReplacementError> {
+    fn apply(mut self, h: &mut impl HugrMut) -> Result<Self::ApplyResult, Self::Error> {
         let parent = self.subgraph.get_parent(h);
         // 1. Check the parent node exists and is a DataflowParent.
         if !OpTag::DataflowParent.is_superset(h.get_optype(parent).tag()) {
@@ -184,10 +184,12 @@ impl Rewrite for SimpleReplacement {
             });
 
         // 3.5. Remove all nodes in self.removal and edges between them.
-        for &node in self.subgraph.nodes() {
-            h.remove_node(node);
-        }
-        Ok(())
+        Ok(self
+            .subgraph
+            .nodes()
+            .iter()
+            .map(|&node| (node, h.remove_node(node)))
+            .collect())
     }
 
     #[inline]
@@ -831,7 +833,7 @@ pub(in crate::hugr::rewrite) mod test {
     }
 
     fn apply_simple(h: &mut Hugr, rw: SimpleReplacement) {
-        h.apply_rewrite(rw).unwrap()
+        h.apply_rewrite(rw).unwrap();
     }
 
     fn apply_replace(h: &mut Hugr, rw: SimpleReplacement) {