Import and export order edges via metadata.

zrho · zrho · commit 054f8d100ef3 · 2025-03-26T17:34:57.000Z
diff --git a/hugr-core/src/export.rs b/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,
     },
@@ -12,10 +12,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};
@@ -467,10 +467,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,
@@ -577,6 +574,7 @@ 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);
@@ -591,9 +589,25 @@ impl<'a> Context<'a> {
                     targets = self.make_ports(child, Direction::Incoming, output.types.len());
                     output_types = Some(&output.types);
                 }
-                _ => {
+                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| !self.hugr.get_optype(*successor).is_output());
+
+                        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]));
+                        }
                     }
                 }
             }
@@ -623,7 +637,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,
         };
@@ -1002,11 +1016,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()
diff --git a/hugr-core/src/import.rs b/hugr-core/src/import.rs
@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use crate::{
     extension::{ExtensionId, ExtensionRegistry, ExtensionSet, SignatureError},
-    hugr::HugrMut,
+    hugr::{HugrMut, IdentList, NodeMetadata},
     ops::{
         constant::{CustomConst, CustomSerialized, OpaqueValue},
         AliasDecl, AliasDefn, Call, CallIndirect, Case, Conditional, Const, DataflowBlock,
@@ -180,11 +180,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)
@@ -600,11 +616,84 @@ 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;
+                };
+
+                // TODO: Error on duplicate key
+                order_keys.insert(*key, *child_id);
+            }
+        }
+
+        // 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;
+            };
+
+            // TODO: Proper error on non-existing key
+            let a = order_keys.get(a).expect("unknown order key");
+            let b = order_keys.get(b).expect("unknown order key");
+
+            // 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];
+
+            // Find the order ports
+            // TODO: Proper error on non-existing order port
+            let a_port = self
+                .hugr
+                .get_optype(a_node)
+                .other_output_port()
+                .expect("order hint on node without order port");
+
+            let b_port = self
+                .hugr
+                .get_optype(b_node)
+                .other_input_port()
+                .expect("order hint on node without order port");
+
+            self.hugr.connect(a_node, a_port, b_node, b_port);
+        }
+
+        Ok(())
+    }
+
     fn import_adt_and_rest(
         &mut self,
         node_id: table::NodeId,
@@ -1341,28 +1430,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,
diff --git a/hugr-core/tests/model.rs b/hugr-core/tests/model.rs
@@ -77,3 +77,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"
+    )));
+}
diff --git a/hugr-core/tests/snapshots/model__roundtrip_order.snap b/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)
+
+(import order_hint.order)
+
+(import order_hint.key)
+
+(import core.fn)
+
+(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]
+    (ext))
+  (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]
+        (ext)))
+    (meta (order_hint.order 4 7))
+    (meta (order_hint.order 5 6))
+    (meta (order_hint.order 5 4))
+    (meta (order_hint.order 6 7))
+    (arithmetic.int.ineg [%0] [%4]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int] (ext)))
+      (meta (order_hint.key 4)))
+    (arithmetic.int.ineg [%1] [%5]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int] (ext)))
+      (meta (order_hint.key 5)))
+    (arithmetic.int.ineg [%2] [%6]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int] (ext)))
+      (meta (order_hint.key 6)))
+    (arithmetic.int.ineg [%3] [%7]
+      (signature
+        (core.fn [arithmetic.int.types.int] [arithmetic.int.types.int] (ext)))
+      (meta (order_hint.key 7)))))
diff --git a/hugr-model/src/v0/mod.rs b/hugr-model/src/v0/mod.rs
@@ -265,6 +265,33 @@ pub const COMPAT_META_JSON: &str = "compat.meta_json";
 /// - **Result:** `(core.const ?type)`
 pub const COMPAT_CONST_JSON: &str = "compat.const_json";
 
+/// Metadata constructor for order hint keys.
+///
+/// Nodes in a dataflow region can be annotated with a key. Each node may have
+/// at most one key and the key must be unique among all nodes in the same
+/// dataflow region. The parent dataflow graph can then use the
+/// `order_hint.order` metadata to imply a desired ordering relation, referring
+/// to the nodes by their key.
+///
+/// - **Parameter:** `?key : core.nat`
+/// - **Result:** `core.meta`
+pub const ORDER_HINT_KEY: &str = "order_hint.key";
+
+/// Metadata constructor for order hints.
+///
+/// When this metadata is attached to a dataflow region, it can indicate a
+/// preferred ordering relation between child nodes. Code generation must take
+/// this into account when deciding on an execution order. The child nodes are
+/// identified by a key, using the `order_hint.key` metadata.
+///
+/// The graph consisting of both value dependencies between nodes and order
+/// hints must be directed acyclic.
+///
+/// - **Parameter:** `?before : core.nat`
+/// - **Parameter:** `?after : core.nat`
+/// - **Result:** `core.meta`
+pub const ORDER_HINT_ORDER: &str = "order_hint.order";
+
 pub mod ast;
 pub mod binary;
 pub mod scope;
diff --git a/hugr-model/tests/fixtures/model-order.edn b/hugr-model/tests/fixtures/model-order.edn
diff --git a/hugr-py/src/hugr/model/export.py b/hugr-py/src/hugr/model/export.py
