Merge remote-tracking branch 'origin/main' into new/add-inferred-resources

Author: croyzor

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/extension/infer.rs

@@ -626,7 +626,7 @@ mod test {
     use super::*;
     use crate::builder::{BuildError, DFGBuilder, Dataflow, DataflowHugr};
     use crate::extension::ExtensionSet;
-    use crate::hugr::HugrInternalsMut;
+    use crate::hugr::HugrMut;
     use crate::hugr::{validate::ValidationError, Hugr, HugrView, NodeType};
     use crate::ops::{self, dataflow::IOTrait, handle::NodeHandle};
     use crate::type_row;

src/extension/op_def.rs

@@ -7,7 +7,7 @@ use super::{
     Extension, ExtensionBuildError, ExtensionId, ExtensionSet, SignatureError, TypeParametrised,
 };
 
-use crate::types::{SignatureDescription, TypeRow};
+use crate::types::SignatureDescription;
 
 use crate::types::FunctionType;
 
@@ -31,8 +31,8 @@ pub trait CustomSignatureFunc: Send + Sync {
         name: &SmolStr,
         arg_values: &[TypeArg],
         misc: &HashMap<String, serde_yaml::Value>,
-        // TODO: Make return type an FunctionType
-    ) -> Result<(TypeRow, TypeRow, ExtensionSet), SignatureError>;
+    ) -> Result<FunctionType, SignatureError>;
+
     /// Describe the signature of a node, given the operation name,
     /// values for the type parameters,
     /// and 'misc' data from the extension definition YAML.
@@ -48,14 +48,14 @@ pub trait CustomSignatureFunc: Send + Sync {
 
 impl<F> CustomSignatureFunc for F
 where
-    F: Fn(&[TypeArg]) -> Result<(TypeRow, TypeRow, ExtensionSet), SignatureError> + Send + Sync,
+    F: Fn(&[TypeArg]) -> Result<FunctionType, SignatureError> + Send + Sync,
 {
     fn compute_signature(
         &self,
         _name: &SmolStr,
         arg_values: &[TypeArg],
         _misc: &HashMap<String, serde_yaml::Value>,
-    ) -> Result<(TypeRow, TypeRow, ExtensionSet), SignatureError> {
+    ) -> Result<FunctionType, SignatureError> {
         self(arg_values)
     }
 }
@@ -217,7 +217,7 @@ impl OpDef {
     /// OpDef with statically-provided [TypeArg]s.
     pub fn compute_signature(&self, args: &[TypeArg]) -> Result<FunctionType, SignatureError> {
         self.check_args(args)?;
-        let (ins, outs, res) = match &self.signature_func {
+        let res = match &self.signature_func {
             SignatureFunc::FromDecl { .. } => {
                 // Sig should be computed solely from inputs + outputs + args.
                 todo!()
@@ -227,7 +227,7 @@ impl OpDef {
         // TODO bring this assert back once resource inference is done?
         // https://github.com/CQCL-DEV/hugr/issues/425
         // assert!(res.contains(self.extension()));
-        Ok(FunctionType::new(ins, outs).with_extension_delta(&res))
+        Ok(res)
     }
 
     /// Optional description of the ports in the signature.

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::{HashMap, 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")]
@@ -251,12 +251,31 @@ impl Hugr {
         }
     }
 
-    /// Produce a canonical ordering of the nodes.
+    /// Add a node to the graph, with the default conversion from OpType to NodeType
+    pub(crate) fn add_op(&mut self, op: impl Into<OpType>) -> Node {
+        // TODO: Default to `NodeType::open_extensions` once we can infer extensions
+        self.add_node(NodeType::pure(op))
+    }
+
+    /// Add a node to the graph.
+    pub(crate) fn add_node(&mut self, nodetype: NodeType) -> Node {
+        let node = self
+            .graph
+            .add_node(nodetype.input_count(), nodetype.output_count());
+        self.op_types[node] = nodetype;
+        node.into()
+    }
+
+    /// 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) {
@@ -265,6 +284,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 {
@@ -367,6 +426,9 @@ 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 {0:?}.")]
+    InvalidNode(Node),
 }
 
 #[cfg(feature = "pyo3")]