feat: Represent order edges in hugr-model as metadata. (#2027)

This PR introduces metadata to encode order edges in `hugr-model`. The
children of a dataflow region can be assigned a key via `order_hint.key`
metadata. Then `order_hint.order` metadata on the dataflow region
encodes order edges between keys. The PR includes the necessary import
and export code in `hugr-core`.

Author: zrho

hugr-core/src/export.rs

@@ -1,7 +1,7 @@
 //! Exporting HUGR graphs to their `hugr-model` representation.
 use crate::{
     extension::{ExtensionId, OpDef, SignatureFunc},
-    hugr::{IdentList, NodeMetadataMap},
+    hugr::IdentList,
     ops::{
         constant::CustomSerialized, DataflowBlock, DataflowOpTrait, OpName, OpTrait, OpType, Value,
     },
@@ -13,10 +13,10 @@ use crate::{
     types::{
         type_param::{TypeArgVariable, TypeParam},
         type_row::TypeRowBase,
-        CustomType, FuncTypeBase, MaybeRV, PolyFuncTypeBase, RowVariable, SumType, TypeArg,
-        TypeBase, TypeBound, TypeEnum, TypeRow,
+        CustomType, EdgeKind, FuncTypeBase, MaybeRV, PolyFuncTypeBase, RowVariable, SumType,
+        TypeArg, TypeBase, TypeBound, TypeEnum, TypeRow,
     },
-    Direction, Hugr, HugrView, IncomingPort, Node, Port,
+    Direction, Hugr, HugrView, IncomingPort, Node, NodeIndex as _, Port,
 };
 
 use fxhash::{FxBuildHasher, FxHashMap};
@@ -478,10 +478,7 @@ impl<'a> Context<'a> {
         let inputs = self.make_ports(node, Direction::Incoming, num_inputs);
         let outputs = self.make_ports(node, Direction::Outgoing, num_outputs);
 
-        let meta = match self.hugr.get_node_metadata(node) {
-            Some(metadata_map) => self.export_node_metadata(metadata_map),
-            None => &[],
-        };
+        let meta = self.export_node_metadata(node);
 
         self.module.nodes[node_id.index()] = table::Node {
             operation,
@@ -588,10 +585,13 @@ impl<'a> Context<'a> {
         let mut targets: &[_] = &[];
         let mut input_types = None;
         let mut output_types = None;
+        let mut meta = Vec::new();
 
         let children = self.hugr.children(node);
         let mut region_children = BumpVec::with_capacity_in(children.size_hint().0 - 2, self.bump);
 
+        let mut output_node = None;
+
         for child in children {
             match self.hugr.get_optype(child) {
                 OpType::Input(input) => {
@@ -601,10 +601,27 @@ impl<'a> Context<'a> {
                 OpType::Output(output) => {
                     targets = self.make_ports(child, Direction::Incoming, output.types.len());
                     output_types = Some(&output.types);
+                    output_node = Some(child);
                 }
-                _ => {
+                child_optype => {
                     if let Some(child_id) = self.export_node_shallow(child) {
                         region_children.push(child_id);
+
+                        // Record all order edges that originate from this node in metadata.
+                        let successors = child_optype
+                            .other_output_port()
+                            .into_iter()
+                            .flat_map(|port| self.hugr.linked_inputs(child, port))
+                            .map(|(successor, _)| successor)
+                            .filter(|successor| Some(*successor) != output_node);
+
+                        for successor in successors {
+                            let a =
+                                self.make_term(model::Literal::Nat(child.index() as u64).into());
+                            let b = self
+                                .make_term(model::Literal::Nat(successor.index() as u64).into());
+                            meta.push(self.make_term_apply(model::ORDER_HINT_ORDER, &[a, b]));
+                        }
                     }
                 }
             }
@@ -634,7 +651,7 @@ impl<'a> Context<'a> {
             sources,
             targets,
             children: region_children.into_bump_slice(),
-            meta: &[], // TODO: Export metadata
+            meta: self.bump.alloc_slice_copy(&meta),
             signature,
             scope,
         };
@@ -1013,11 +1030,37 @@ impl<'a> Context<'a> {
         }
     }
 
-    pub fn export_node_metadata(&mut self, metadata_map: &NodeMetadataMap) -> &'a [table::TermId] {
-        let mut meta = BumpVec::with_capacity_in(metadata_map.len(), self.bump);
+    pub fn export_node_metadata(&mut self, node: Node) -> &'a [table::TermId] {
+        let metadata_map = self.hugr.get_node_metadata(node);
+
+        let has_order_edges = {
+            fn is_relevant_node(hugr: &Hugr, node: Node) -> bool {
+                let optype = hugr.get_optype(node);
+                !optype.is_input() && !optype.is_output()
+            }
+
+            let optype = self.hugr.get_optype(node);
+
+            Direction::BOTH
+                .iter()
+                .filter(|dir| optype.other_port_kind(**dir) == Some(EdgeKind::StateOrder))
+                .filter_map(|dir| optype.other_port(*dir))
+                .flat_map(|port| self.hugr.linked_ports(node, port))
+                .any(|(other, _)| is_relevant_node(self.hugr, other))
+        };
+
+        let meta_capacity = metadata_map.map_or(0, |map| map.len()) + has_order_edges as usize;
+        let mut meta = BumpVec::with_capacity_in(meta_capacity, self.bump);
+
+        if let Some(metadata_map) = metadata_map {
+            for (name, value) in metadata_map {
+                meta.push(self.export_json_meta(name, value));
+            }
+        }
 
-        for (name, value) in metadata_map {
-            meta.push(self.export_json_meta(name, value));
+        if has_order_edges {
+            let key = self.make_term(model::Literal::Nat(node.index() as u64).into());
+            meta.push(self.make_term_apply(model::ORDER_HINT_KEY, &[key]));
         }
 
         meta.into_bump_slice()

hugr-core/src/import.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use crate::{
     extension::{ExtensionId, ExtensionRegistry, ExtensionSet, SignatureError},
-    hugr::HugrMut,
+    hugr::{HugrMut, NodeMetadata},
     ops::{
         constant::{CustomConst, CustomSerialized, OpaqueValue},
         AliasDecl, AliasDefn, Call, CallIndirect, Case, Conditional, Const, DataflowBlock,
@@ -68,6 +68,23 @@ pub enum ImportError {
     /// The model is not well-formed.
     #[error("validate error: {0}")]
     Model(#[from] table::ModelError),
+    /// Incorrect order hints.
+    #[error("incorrect order hint: {0}")]
+    OrderHint(#[from] OrderHintError),
+}
+
+/// Import error caused by incorrect order hints.
+#[derive(Debug, Clone, Error)]
+pub enum OrderHintError {
+    /// Duplicate order hint key in the same region.
+    #[error("duplicate order hint key {0}")]
+    DuplicateKey(table::NodeId, u64),
+    /// Order hint including a key not defined in the region.
+    #[error("order hint with unknown key {0}")]
+    UnknownKey(u64),
+    /// Order hint involving a node with no order port.
+    #[error("order hint on node with no order port: {0}")]
+    NoOrderPort(table::NodeId),
 }
 
 /// Helper macro to create an `ImportError::Unsupported` error with a formatted message.
@@ -197,11 +214,27 @@ impl<'a> Context<'a> {
         self.record_links(node, Direction::Incoming, node_data.inputs);
         self.record_links(node, Direction::Outgoing, node_data.outputs);
 
+        // Import the JSON metadata
         for meta_item in node_data.meta {
-            // TODO: For now we expect all metadata to be JSON since this is how
-            // it is handled in `hugr-core`.
-            let (name, value) = self.import_json_meta(*meta_item)?;
-            self.hugr.set_metadata(node, name, value);
+            let Some([name_arg, json_arg]) =
+                self.match_symbol(*meta_item, model::COMPAT_META_JSON)?
+            else {
+                continue;
+            };
+
+            let table::Term::Literal(model::Literal::Str(name)) = self.get_term(name_arg)? else {
+                return Err(table::ModelError::TypeError(*meta_item).into());
+            };
+
+            let table::Term::Literal(model::Literal::Str(json_str)) = self.get_term(json_arg)?
+            else {
+                return Err(table::ModelError::TypeError(*meta_item).into());
+            };
+
+            let json_value: NodeMetadata = serde_json::from_str(json_str)
+                .map_err(|_| table::ModelError::TypeError(*meta_item))?;
+
+            self.hugr.set_metadata(node, name, json_value);
         }
 
         Ok(node)
@@ -617,11 +650,82 @@ impl<'a> Context<'a> {
             self.import_node(*child, node)?;
         }
 
+        self.create_order_edges(region)?;
+
         self.region_scope = prev_region;
 
         Ok(())
     }
 
+    /// Create order edges between nodes of a dataflow region based on order hint metadata.
+    ///
+    /// This method assumes that the nodes for the children of the region have already been imported.
+    fn create_order_edges(&mut self, region_id: table::RegionId) -> Result<(), ImportError> {
+        let region_data = self.get_region(region_id)?;
+        debug_assert_eq!(region_data.kind, model::RegionKind::DataFlow);
+
+        // Collect order hint keys
+        // PERFORMANCE: It might be worthwhile to reuse the map to avoid allocations.
+        let mut order_keys = FxHashMap::<u64, table::NodeId>::default();
+
+        for child_id in region_data.children {
+            let child_data = self.get_node(*child_id)?;
+
+            for meta_id in child_data.meta {
+                let Some([key]) = self.match_symbol(*meta_id, model::ORDER_HINT_KEY)? else {
+                    continue;
+                };
+
+                let table::Term::Literal(model::Literal::Nat(key)) = self.get_term(key)? else {
+                    continue;
+                };
+
+                if order_keys.insert(*key, *child_id).is_some() {
+                    return Err(OrderHintError::DuplicateKey(*child_id, *key).into());
+                }
+            }
+        }
+
+        // Insert order edges
+        for meta_id in region_data.meta {
+            let Some([a, b]) = self.match_symbol(*meta_id, model::ORDER_HINT_ORDER)? else {
+                continue;
+            };
+
+            let table::Term::Literal(model::Literal::Nat(a)) = self.get_term(a)? else {
+                continue;
+            };
+
+            let table::Term::Literal(model::Literal::Nat(b)) = self.get_term(b)? else {
+                continue;
+            };
+
+            let a = order_keys.get(a).ok_or(OrderHintError::UnknownKey(*a))?;
+            let b = order_keys.get(b).ok_or(OrderHintError::UnknownKey(*b))?;
+
+            // NOTE: The lookups here are expected to succeed since we only
+            // process the order metadata after we have imported the nodes.
+            let a_node = self.nodes[a];
+            let b_node = self.nodes[b];
+
+            let a_port = self
+                .hugr
+                .get_optype(a_node)
+                .other_output_port()
+                .ok_or(OrderHintError::NoOrderPort(*a))?;
+
+            let b_port = self
+                .hugr
+                .get_optype(b_node)
+                .other_input_port()
+                .ok_or(OrderHintError::NoOrderPort(*b))?;
+
+            self.hugr.connect(a_node, a_port, b_node, b_port);
+        }
+
+        Ok(())
+    }
+
     fn import_adt_and_rest(
         &mut self,
         node_id: table::NodeId,
@@ -1358,28 +1462,6 @@ impl<'a> Context<'a> {
         }
     }
 
-    fn import_json_meta(
-        &mut self,
-        term_id: table::TermId,
-    ) -> Result<(&'a str, serde_json::Value), ImportError> {
-        let [name_arg, json_arg] = self
-            .match_symbol(term_id, model::COMPAT_META_JSON)?
-            .ok_or(table::ModelError::TypeError(term_id))?;
-
-        let table::Term::Literal(model::Literal::Str(name)) = self.get_term(name_arg)? else {
-            return Err(table::ModelError::TypeError(term_id).into());
-        };
-
-        let table::Term::Literal(model::Literal::Str(json_str)) = self.get_term(json_arg)? else {
-            return Err(table::ModelError::TypeError(term_id).into());
-        };
-
-        let json_value =
-            serde_json::from_str(json_str).map_err(|_| table::ModelError::TypeError(term_id))?;
-
-        Ok((name, json_value))
-    }
-
     fn import_value(
         &mut self,
         term_id: table::TermId,

hugr-core/tests/model.rs

@@ -78,3 +78,10 @@ pub fn test_roundtrip_const() {
         "../../hugr-model/tests/fixtures/model-const.edn"
     )));
 }
+
+#[test]
+pub fn test_roundtrip_order() {
+    insta::assert_snapshot!(roundtrip(include_str!(
+        "../../hugr-model/tests/fixtures/model-order.edn"
+    )));
+}

hugr-core/tests/snapshots/model__roundtrip_order.snap

@@ -0,0 +1,60 @@
+---
+source: hugr-core/tests/model.rs
+expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-order.edn\"))"
+---
+(hugr 0)
+
+(mod)
+
+(import core.order_hint.key)
+
+(import core.fn)
+
+(import core.order_hint.order)
+
+(import arithmetic.int.types.int)
+
+(import arithmetic.int.ineg)
+
+(define-func
+  main
+  (core.fn
+    [arithmetic.int.types.int
+     arithmetic.int.types.int
+     arithmetic.int.types.int
+     arithmetic.int.types.int]
+    [arithmetic.int.types.int
+     arithmetic.int.types.int
+     arithmetic.int.types.int
+     arithmetic.int.types.int])
+  (dfg [%0 %1 %2 %3] [%4 %5 %6 %7]
+    (signature
+      (core.fn
+        [arithmetic.int.types.int
+         arithmetic.int.types.int
+         arithmetic.int.types.int
+         arithmetic.int.types.int]
+        [arithmetic.int.types.int
+         arithmetic.int.types.int
+         arithmetic.int.types.int
+         arithmetic.int.types.int]))
+    (meta (core.order_hint.order 4 7))
+    (meta (core.order_hint.order 5 6))
+    (meta (core.order_hint.order 5 4))
+    (meta (core.order_hint.order 6 7))
+    (arithmetic.int.ineg [%0] [%4]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int]))
+      (meta (core.order_hint.key 4)))
+    (arithmetic.int.ineg [%1] [%5]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int]))
+      (meta (core.order_hint.key 5)))
+    (arithmetic.int.ineg [%2] [%6]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int]))
+      (meta (core.order_hint.key 6)))
+    (arithmetic.int.ineg [%3] [%7]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int]))
+      (meta (core.order_hint.key 7)))))