feat: Split methods between HugrMut and HugrMutInternals

src/builder/build_traits.rs

@@ -29,7 +29,7 @@ use super::{
 
 use crate::Hugr;
 
-use crate::hugr::HugrInternalsMut;
+use crate::hugr::HugrMut;
 
 /// Trait for HUGR container builders.
 /// Containers are nodes that are parents of sibling graphs.

src/builder/cfg.rs

@@ -12,7 +12,7 @@ use crate::{ops::handle::NodeHandle, types::Type};
 
 use crate::Node;
 use crate::{
-    hugr::{HugrInternalsMut, NodeType},
+    hugr::{HugrMut, NodeType},
     type_row, Hugr,
 };
 

src/builder/conditional.rs

@@ -15,7 +15,7 @@ use super::{
 
 use crate::Node;
 use crate::{
-    hugr::{HugrInternalsMut, NodeType},
+    hugr::{HugrMut, NodeType},
     Hugr,
 };
 

src/builder/dataflow.rs

@@ -11,7 +11,7 @@ use crate::types::{FunctionType, Signature};
 
 use crate::extension::ExtensionSet;
 use crate::Node;
-use crate::{hugr::HugrInternalsMut, Hugr};
+use crate::{hugr::HugrMut, Hugr};
 
 /// Builder for a [`ops::DFG`] node.
 #[derive(Debug, Clone, PartialEq)]

src/builder/module.rs

@@ -4,6 +4,7 @@ use super::{
     BuildError, Container,
 };
 
+use crate::hugr::hugrmut::sealed::HugrMutInternals;
 use crate::{
     hugr::{views::HugrView, ValidationError},
     ops,
@@ -18,10 +19,7 @@ use crate::types::Signature;
 use crate::Node;
 use smol_str::SmolStr;
 
-use crate::{
-    hugr::{HugrInternalsMut, NodeType},
-    Hugr,
-};
+use crate::{hugr::NodeType, Hugr};
 
 /// Builder for a HUGR module.
 #[derive(Debug, Clone, PartialEq)]

src/hugr.rs

@@ -1,6 +1,6 @@
 //! The Hugr data structure, and its basic component handles.
 
-mod hugrmut;
+pub mod hugrmut;
 
 pub mod rewrite;
 pub mod serialize;
@@ -10,14 +10,14 @@ pub mod views;
 use std::collections::VecDeque;
 use std::iter;
 
-pub(crate) use self::hugrmut::HugrInternalsMut;
+pub(crate) use self::hugrmut::HugrMut;
 pub use self::validate::ValidationError;
 
 use derive_more::From;
 pub use rewrite::{Rewrite, SimpleReplacement, SimpleReplacementError};
 
 use portgraph::multiportgraph::MultiPortGraph;
-use portgraph::{Hierarchy, PortMut, UnmanagedDenseMap};
+use portgraph::{Hierarchy, NodeIndex, PortMut, UnmanagedDenseMap};
 use thiserror::Error;
 
 #[cfg(feature = "pyo3")]
@@ -224,12 +224,16 @@ impl Hugr {
         }
     }
 
-    /// Produce a canonical ordering of the nodes.
+    /// Produce a canonical ordering of the descendant nodes of a root,
+    /// following the graph hierarchy.
+    ///
+    /// This starts with the root, and then proceeds in BFS order through the
+    /// contained regions.
     ///
     /// Used by [`HugrMut::canonicalize_nodes`] and the serialization code.
-    fn canonical_order(&self) -> impl Iterator<Item = Node> + '_ {
+    fn canonical_order(&self, root: Node) -> impl Iterator<Item = Node> + '_ {
         // Generate a BFS-ordered list of nodes based on the hierarchy
-        let mut queue = VecDeque::from([self.root.into()]);
+        let mut queue = VecDeque::from([root]);
         iter::from_fn(move || {
             let node = queue.pop_front()?;
             for child in self.children(node) {
@@ -238,6 +242,46 @@ impl Hugr {
             Some(node)
         })
     }
+
+    /// Compact the nodes indices of the hugr to be contiguous, and order them as a breadth-first
+    /// traversal of the hierarchy.
+    ///
+    /// The rekey function is called for each moved node with the old and new indices.
+    ///
+    /// After this operation, a serialization and deserialization of the Hugr is guaranteed to
+    /// preserve the indices.
+    pub fn canonicalize_nodes(&mut self, mut rekey: impl FnMut(Node, Node)) {
+        // Generate the ordered list of nodes
+        let mut ordered = Vec::with_capacity(self.node_count());
+        let root = self.root();
+        ordered.extend(self.as_mut().canonical_order(root));
+
+        // Permute the nodes in the graph to match the order.
+        //
+        // Invariant: All the elements before `position` are in the correct place.
+        for position in 0..ordered.len() {
+            // Find the element's location. If it originally came from a previous position
+            // then it has been swapped somewhere else, so we follow the permutation chain.
+            let mut source: Node = ordered[position];
+            while position > source.index.index() {
+                source = ordered[source.index.index()];
+            }
+
+            let target: Node = NodeIndex::new(position).into();
+            if target != source {
+                self.graph.swap_nodes(target.index, source.index);
+                self.op_types.swap(target.index, source.index);
+                self.hierarchy.swap_nodes(target.index, source.index);
+                rekey(source, target);
+            }
+        }
+        self.root = NodeIndex::new(0);
+
+        // Finish by compacting the copy nodes.
+        // The operation nodes will be left in place.
+        // This step is not strictly necessary.
+        self.graph.compact_nodes(|_, _| {});
+    }
 }
 
 impl Port {
@@ -340,6 +384,12 @@ pub enum HugrError {
     /// An error occurred while manipulating the hierarchy.
     #[error("An error occurred while manipulating the hierarchy.")]
     HierarchyError(#[from] portgraph::hierarchy::AttachError),
+    /// The node doesn't exist.
+    #[error("Invalid node {node:?}.")]
+    InvalidNode {
+        /// The missing node
+        node: Node,
+    },
 }
 
 #[cfg(feature = "pyo3")]