switched to outgoing ports

Author: lmondada

hugr-core/src/hugr/patch/simple_replace.rs

@@ -1,6 +1,7 @@
 //! Implementation of the `SimpleReplace` operation.
 
 use std::collections::{BTreeSet, HashMap};
+use std::hash::Hash;
 
 use crate::core::HugrNode;
 use crate::hugr::hugrmut::InsertionResult;
@@ -353,41 +354,6 @@ impl<
             .map(Into::into)
     }
 
-    /// Map the host nodes in `self` according to `node_map`.
-    ///
-    /// `node_map` must map nodes in the current HUGR of the subgraph to
-    /// its equivalent nodes in some `new_hugr`.
-    ///
-    /// This converts a replacement that acts on nodes of type `HostNode` to
-    /// a replacement that acts on `new_hugr`, with nodes of type `N`.
-    ///
-    /// This does not check convexity. It is up to the caller to ensure that
-    /// the mapped replacement obtained from this applies on a convex subgraph
-    /// of the new HUGR.
-    pub(crate) fn map_host_nodes<N: HugrNode>(
-        &self,
-        node_map: impl Fn(HostNode) -> N,
-    ) -> SimpleReplacement<N> {
-        let Self {
-            subgraph,
-            replacement,
-            nu_inp,
-            nu_out,
-        } = self;
-        let nu_inp = nu_inp
-            .iter()
-            .map(|(&(node, port), &(host_node, host_port))| {
-                ((node, port), (node_map(host_node), host_port))
-            })
-            .collect();
-        let nu_out = nu_out
-            .iter()
-            .map(|(&(host_node, host_port), &port)| ((node_map(host_node), host_port), port))
-            .collect();
-        let subgraph = subgraph.map_nodes(node_map);
-        SimpleReplacement::new(subgraph, replacement.clone(), nu_inp, nu_out)
-    }
-
     /// Get all edges that the replacement would add between `host` and
     /// `self.replacement`.
     ///
@@ -421,6 +387,48 @@ impl<
     }
 }
 
+impl<HostNode: HugrNode, P: Copy + Eq + Hash>
+    SimpleReplacement<HostNode, DefaultInMap<HostNode>, OutputNodeBoundaryMap<HostNode, P>>
+{
+    /// Map the host nodes in `self` according to `node_map`.
+    ///
+    /// `node_map` must map nodes in the current HUGR of the subgraph to
+    /// its equivalent nodes in some `new_hugr`.
+    ///
+    /// This converts a replacement that acts on nodes of type `HostNode` to
+    /// a replacement that acts on `new_hugr`, with nodes of type `N`.
+    ///
+    /// This does not check convexity. It is up to the caller to ensure that
+    /// the mapped replacement obtained from this applies on a convex subgraph
+    /// of the new HUGR.
+    pub(crate) fn map_host_nodes<N: HugrNode>(
+        &self,
+        node_map: impl Fn(HostNode) -> N,
+    ) -> SimpleReplacement<N, DefaultInMap<N>, OutputNodeBoundaryMap<N, P>> {
+        let Self {
+            subgraph,
+            replacement,
+            nu_inp,
+            nu_out,
+        } = self;
+        let nu_inp = nu_inp
+            .iter()
+            .map(|(&repl_node_port, &(host_node, host_port))| {
+                (repl_node_port, (node_map(host_node), host_port))
+            })
+            .collect();
+        let nu_out: HashMap<_, _> = nu_out
+            .0
+            .iter()
+            .map(|(&(host_node, host_port), &repl_port)| {
+                ((node_map(host_node), host_port), repl_port)
+            })
+            .collect();
+        let subgraph = subgraph.map_nodes(node_map);
+        SimpleReplacement::new_unsafe(subgraph, replacement.clone(), nu_inp, nu_out.into())
+    }
+}
+
 fn check_valid_boundary<HostNode, InMap, OutMap>(
     subgraph: &SiblingSubgraph<HostNode>,
     host: &impl HugrView<Node = HostNode>,

hugr-core/src/hugr/persistent.rs

@@ -80,11 +80,17 @@ pub use resolver::PointerEqResolver;
 
 use crate::{
     hugr::patch::{simple_replace, Patch, PatchVerification},
-    Hugr, HugrView, Node, SimpleReplacement,
+    Hugr, HugrView, Node, OutgoingPort, SimpleReplacement,
 };
 
+use super::patch::{simple_replace::DefaultInMap, OutputNodeBoundaryMap};
+
 /// A replacement operation that can be applied to a [`PersistentHugr`].
-pub type PersistentReplacement = SimpleReplacement<PatchNode>;
+pub type PersistentReplacement = SimpleReplacement<
+    PatchNode,
+    DefaultInMap<PatchNode>,
+    OutputNodeBoundaryMap<PatchNode, OutgoingPort>,
+>;
 
 /// A patch that can be applied to a [`PersistentHugr`] or a
 /// [`CommitStateSpace`] as an atomic commit.
@@ -139,11 +145,7 @@ impl Commit {
         &self.0
     }
 
-    fn all_parents(&self) -> impl Iterator<Item = &Commit> + '_ {
-        self.0.all_parents().map_into()
-    }
-
-    fn replacement(&self) -> Option<&SimpleReplacement<PatchNode>> {
+    fn replacement(&self) -> Option<&PersistentReplacement> {
         match self.0.value() {
             CommitData::Base(_) => None,
             CommitData::Replacement(replacement) => Some(replacement),
@@ -278,7 +280,7 @@ impl PersistentHugr {
         let new_invalid_nodes = replacement
             .subgraph()
             .nodes()
-            .into_iter()
+            .iter()
             .map(|&PatchNode(id, node)| (id, node))
             .into_grouping_map()
             .collect::<BTreeSet<_>>();

hugr-core/src/hugr/persistent/tests.rs

@@ -5,7 +5,12 @@ use rstest::*;
 use crate::{
     builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr},
     extension::prelude::bool_t,
-    hugr::{patch::Patch, persistent::PatchNode, views::SiblingSubgraph, Hugr, HugrView},
+    hugr::{
+        patch::{simple_replace::DefaultInMap, OutputNodeBoundaryMap, Patch},
+        persistent::PatchNode,
+        views::SiblingSubgraph,
+        Hugr, HugrView,
+    },
     ops::handle::NodeHandle,
     std_extensions::logic::LogicOp,
     IncomingPort, Node, OutgoingPort, SimpleReplacement,
@@ -48,7 +53,10 @@ fn simple_hugr() -> (Hugr, [Node; 3]) {
 }
 
 /// Creates a replacement that replaces a node with a sequence of two NOT gates
-fn create_double_not_replacement(hugr: &Hugr, node_to_replace: Node) -> SimpleReplacement {
+fn create_double_not_replacement(
+    hugr: &Hugr,
+    node_to_replace: Node,
+) -> SimpleReplacement<Node, DefaultInMap<Node>, OutputNodeBoundaryMap<Node, OutgoingPort>> {
     // Create a simple hugr with two NOT gates in sequence
     let mut dfg_builder = DFGBuilder::new(inout_sig(vec![bool_t()], vec![bool_t()])).unwrap();
     let [input_wire] = dfg_builder.input_wires_arr();
@@ -67,11 +75,6 @@ fn create_double_not_replacement(hugr: &Hugr, node_to_replace: Node) -> SimpleRe
 
     let replacement_hugr = dfg_builder.finish_hugr_with_outputs([not2_out]).unwrap();
 
-    // Find the input and output of the node to replace
-    let host_output_target = hugr
-        .single_linked_input(node_to_replace, OutgoingPort::from(0))
-        .unwrap();
-
     // Create the mappings
     let mut nu_inp = HashMap::new();
     nu_inp.insert(
@@ -80,18 +83,30 @@ fn create_double_not_replacement(hugr: &Hugr, node_to_replace: Node) -> SimpleRe
     );
 
     let mut nu_out = HashMap::new();
-    nu_out.insert(host_output_target, IncomingPort::from(0));
+    nu_out.insert(
+        (node_to_replace, OutgoingPort::from(0)),
+        IncomingPort::from(0),
+    );
 
     // Create the subgraph with the single node
     let subgraph = SiblingSubgraph::try_from_nodes(vec![node_to_replace], hugr).unwrap();
 
     // Create the replacement
-    SimpleReplacement::new(subgraph, replacement_hugr, nu_inp, nu_out)
+    SimpleReplacement::try_new(
+        subgraph,
+        hugr,
+        replacement_hugr,
+        nu_inp,
+        OutputNodeBoundaryMap::from(nu_out),
+    )
+    .unwrap()
 }
 
 /// Creates a replacement that replaces the unique AND gate in `hugr` and its
 /// predecessor NOT gate on 1st input with an XOR gate
-fn create_not_and_to_xor_replacement(hugr: &Hugr) -> SimpleReplacement {
+fn create_not_and_to_xor_replacement(
+    hugr: &Hugr,
+) -> SimpleReplacement<Node, DefaultInMap<Node>, OutputNodeBoundaryMap<Node, OutgoingPort>> {
     // Create second replacement that replaces the second NOT gate from the first
     // replacement
     // Find the AND gate in the hugr
@@ -133,13 +148,19 @@ fn create_not_and_to_xor_replacement(hugr: &Hugr) -> SimpleReplacement {
 
     // Output mapping - AND gate's output to XOR's output
     let mut nu_out = HashMap::new();
-    let and_output_port = hugr.single_linked_input(and_gate, 0).unwrap();
-    nu_out.insert(and_output_port, IncomingPort::from(0));
+    nu_out.insert((and_gate, OutgoingPort::from(0)), IncomingPort::from(0));
 
     // Create subgraph with both the AND gate and NOT0 node
     let subgraph = SiblingSubgraph::try_from_nodes(vec![not_node, and_gate], &hugr).unwrap();
 
-    SimpleReplacement::new(subgraph, replacement_hugr, nu_inp, nu_out)
+    SimpleReplacement::try_new(
+        subgraph,
+        hugr,
+        replacement_hugr,
+        nu_inp,
+        OutputNodeBoundaryMap::from(nu_out),
+    )
+    .unwrap()
 }
 
 #[fixture]