feat: add HierarchyTester for O(1) subtree-containment; rewrite nonlocal edge validation

hugr-core/src/hierarchy.rs

@@ -0,0 +1,150 @@
+//! Çontains [HierarchyTester], a tool for efficiently querying the hierarchy
+//! (constant-time in size and depth of Hugr)
+
+use std::collections::{HashMap, hash_map::Entry};
+
+use itertools::Itertools;
+
+use crate::HugrView;
+
+/// The entry and exit indices (inclusive, note every entry number is different);
+/// and the nearest strictly-enclosing FuncDefn.
+type NodeData<N> = (usize, usize, Option<N>);
+
+/// Caches enough information on the hierarchy of an immutably-held Hugr
+/// to allow efficient querying of [Self::is_ancestor_of] and [Self::nearest_enclosing_funcdefn].
+#[derive(Clone, Debug)]
+pub struct HierarchyTester<'a, H: HugrView> {
+    #[allow(unused)] // Just make sure the Hugr isn't changing behind our back
+    hugr: &'a H,
+    entry_exit: HashMap<H::Node, NodeData<H::Node>>, // for every node beneath entrypoint
+}
+
+impl<'a, H: HugrView> HierarchyTester<'a, H> {
+    /// Creates a new instance that knows about all descendants of the
+    /// specified Hugr's entrypoint
+    pub fn new(hugr: &'a H) -> Self {
+        let mut entry_exit = HashMap::new();
+        fn traverse<H: HugrView>(
+            hugr: &H,
+            n: H::Node,
+            mut fd: Option<H::Node>,
+            ee: &mut HashMap<H::Node, NodeData<H::Node>>,
+        ) {
+            let old = ee.insert(n, (ee.len(), usize::MAX, fd)); // second is placeholder for now
+            debug_assert!(old.is_none());
+            if hugr.get_optype(n).is_func_defn() {
+                fd = Some(n)
+            }
+            for ch in hugr.children(n) {
+                traverse(hugr, ch, fd, ee)
+            }
+            let end_idx = ee.len() - 1;
+            let Entry::Occupied(oe) = ee.entry(n) else {
+                panic!()
+            };
+            let (_, end, _) = oe.into_mut();
+            *end = end_idx;
+            debug_assert!(
+                // Could do this on every which_child_contains?!
+                hugr.children(n)
+                    .tuple_windows()
+                    .all(|(a, b)| ee.get(&a).unwrap().1 == ee.get(&b).unwrap().0 - 1)
+            );
+        }
+        traverse(hugr, hugr.entrypoint(), None, &mut entry_exit);
+        Self { hugr, entry_exit }
+    }
+
+    /// Returns true if `anc` is an ancestor of `desc`, including `anc == desc`.
+    /// (See also [Self::is_strict_ancestor_of].)
+    ///
+    /// # Panics
+    ///
+    /// if `n` is not an entry-descendant in the Hugr
+    pub fn is_ancestor_of(&self, anc: H::Node, desc: H::Node) -> bool {
+        let anc = self.entry_exit.get(&anc).unwrap();
+        let desc = self.entry_exit.get(&desc).unwrap();
+        anc.0 <= desc.0 && desc.0 <= anc.1
+    }
+
+    /// Returns true if `anc` is an ancestor of `desc`, excluding `anc == desc`.
+    /// (See also [Self::is_ancestor_of].)
+    /// Constant time regardless of size/depth of Hugr.
+    ///
+    /// # Panics
+    ///
+    /// if `n` is not an entry-descendant in the Hugr
+    pub fn is_strict_ancestor_of(&self, anc: H::Node, desc: H::Node) -> bool {
+        let anc = self.entry_exit.get(&anc).unwrap();
+        let desc = self.entry_exit.get(&desc).unwrap();
+        anc.0 < desc.0 && desc.1 <= anc.1
+    }
+
+    /// Returns the nearest strictly-enclosing [FuncDefn](crate::ops::FuncDefn) of a node
+    ///
+    /// # Panics
+    ///
+    /// if `n` is not an entry-descendant in the Hugr
+    pub fn nearest_enclosing_funcdefn(&self, n: H::Node) -> Option<H::Node> {
+        self.entry_exit.get(&n).unwrap().2
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::iter;
+
+    use proptest::prelude::{Just, Strategy};
+    use proptest::proptest;
+
+    use crate::builder::{DFGBuilder, Dataflow, HugrBuilder, SubContainer};
+    use crate::extension::prelude::usize_t;
+    use crate::types::Signature;
+    use crate::{Hugr, HugrView};
+
+    use super::HierarchyTester;
+
+    #[derive(Clone, Debug)]
+    struct Layout(Vec<Layout>);
+
+    fn make<H: AsMut<Hugr> + AsRef<Hugr>>(dfg: &mut DFGBuilder<H>, l: Layout) {
+        let [mut val] = dfg.input_wires_arr();
+        for c in l.0 {
+            let mut c_b = dfg
+                .dfg_builder(Signature::new_endo(usize_t()), [val])
+                .unwrap();
+            make(&mut c_b, c);
+            let [res] = c_b.finish_sub_container().unwrap().outputs_arr();
+            val = res;
+        }
+        dfg.set_outputs([val]).unwrap()
+    }
+
+    fn any_layout() -> impl Strategy<Value = Layout> {
+        Just(Layout(vec![])).prop_recursive(5, 10, 3, |elem| {
+            proptest::collection::vec(elem, 1..5).prop_map(Layout)
+        })
+    }
+
+    fn strict_ancestors<H: HugrView>(h: &H, n: H::Node) -> impl Iterator<Item = H::Node> {
+        iter::successors(h.get_parent(n), |n| h.get_parent(*n))
+    }
+
+    proptest! {
+        #[test]
+        fn hierarchy_test(ly in any_layout()) {
+            let mut h = DFGBuilder::new(Signature::new_endo(usize_t())).unwrap();
+            make(&mut h, ly);
+            let h = h.finish_hugr().unwrap();
+            let ht = HierarchyTester::new(&h);
+            for n1 in h.entry_descendants() {
+                for n2 in h.entry_descendants() {
+                    let is_strict_ancestor = strict_ancestors(&h, n1).any(|item| item==n2);
+                    assert_eq!(ht.is_strict_ancestor_of(n2, n1), is_strict_ancestor);
+                    assert_eq!(ht.is_ancestor_of(n2, n1), is_strict_ancestor || n1 == n2);
+                }
+            }
+        }
+    }
+}

hugr-core/src/hugr/validate.rs

@@ -1,7 +1,6 @@
 //! HUGR invariant checks.
 
 use std::collections::HashMap;
-use std::iter;
 
 use itertools::Itertools;
 use petgraph::algo::dominators::{self, Dominators};
@@ -11,6 +10,7 @@ use thiserror::Error;
 use crate::core::HugrNode;
 use crate::extension::SignatureError;
 
+use crate::hierarchy::HierarchyTester;
 use crate::ops::constant::ConstTypeError;
 use crate::ops::custom::{ExtensionOp, OpaqueOpError};
 use crate::ops::validate::{
@@ -32,6 +32,7 @@ use super::views::HugrView;
 /// Hugr to avoid recomputing it every time.
 pub(super) struct ValidationContext<'a, H: HugrView> {
     hugr: &'a H,
+    hierarchy_tester: Option<HierarchyTester<'a, H>>,
     /// Dominator tree for each CFG region, using the container node as index.
     dominators: HashMap<H::Node, (Dominators<portgraph::NodeIndex>, H::RegionPortgraphNodes)>,
 }
@@ -40,7 +41,11 @@ impl<'a, H: HugrView> ValidationContext<'a, H> {
     /// Create a new validation context.
     pub fn new(hugr: &'a H) -> Self {
         let dominators = HashMap::new();
-        Self { hugr, dominators }
+        Self {
+            hugr,
+            hierarchy_tester: None,
+            dominators,
+        }
     }
 
     /// Check the validity of the HUGR.
@@ -81,21 +86,6 @@ impl<'a, H: HugrView> ValidationContext<'a, H> {
         Ok(())
     }
 
-    /// Compute the dominator tree for a CFG region, identified by its container
-    /// node.
-    ///
-    /// The results of this computation should be cached in `self.dominators`.
-    /// We don't do it here to avoid mutable borrows.
-    fn compute_dominator(
-        &self,
-        parent: H::Node,
-    ) -> (Dominators<portgraph::NodeIndex>, H::RegionPortgraphNodes) {
-        let (region, node_map) = self.hugr.region_portgraph(parent);
-        let entry_node = self.hugr.children(parent).next().unwrap();
-        let doms = dominators::simple_fast(&region, node_map.to_portgraph(entry_node));
-        (doms, node_map)
-    }
-
     /// Check the constraints on a single node.
     ///
     /// This includes:
@@ -417,6 +407,17 @@ impl<'a, H: HugrView> ValidationContext<'a, H> {
         to: H::Node,
         to_offset: Port,
     ) -> Result<(), InterGraphEdgeError<H::Node>> {
+        fn containing_child<H: HugrView>(hugr: &H, parent: H::Node, desc: H::Node) -> H::Node {
+            let mut n = desc;
+            loop {
+                let p = hugr.get_parent(n).unwrap();
+                if p == parent {
+                    return n;
+                }
+                n = p;
+            }
+        }
+
         let from_parent = self
             .hugr
             .get_parent(from)
@@ -437,94 +438,91 @@ impl<'a, H: HugrView> ValidationContext<'a, H> {
             });
         }
 
-        // To detect either external or dominator edges, we traverse the ancestors
-        // of the target until we find either `from_parent` (in the external
-        // case), or the parent of `from_parent` (in the dominator case).
-        //
-        // This search could be sped-up with a pre-computed LCA structure, but
-        // for valid Hugrs this search should be very short.
-        //
-        // For Value edges only, we record any FuncDefn we went through; if there is
-        // any such, then that is an error, but we report that only if the dom/ext
-        // relation was otherwise ok (an error about an edge "entering" some ancestor
-        // node could be misleading if the source isn't where it's expected)
-        let mut err_entered_func = None;
+        let ht = &*self
+            .hierarchy_tester
+            .get_or_insert_with(|| HierarchyTester::new(self.hugr));
         let from_parent_parent = self.hugr.get_parent(from_parent);
-        for (ancestor, ancestor_parent) in
-            iter::successors(to_parent, |&p| self.hugr.get_parent(p)).tuple_windows()
+
+        if ht.is_strict_ancestor_of(from_parent, to) {
+            // External edge.
+            if is_static {
+                return Ok(());
+            }
+            //Check for Order edge: some Order-successor should contain the target
+            assert_eq!(
+                from_optype.other_port_kind(Direction::Outgoing),
+                Some(EdgeKind::StateOrder)
+            );
+
+            if !self
+                .hugr
+                .linked_inputs(from, from_optype.other_output_port().unwrap())
+                .any(|(sibling, _)| ht.is_strict_ancestor_of(sibling, to))
+            {
+                return Err(InterGraphEdgeError::MissingOrderEdge {
+                    from,
+                    from_offset,
+                    to,
+                    to_offset,
+                    to_ancestor: containing_child(self.hugr, from_parent, to),
+                })?;
+            }
+        } else if let Some(fpp) =
+            from_parent_parent.filter(|fpp| !is_static && ht.is_strict_ancestor_of(*fpp, to))
         {
-            if !is_static && self.hugr.get_optype(ancestor).is_func_defn() {
-                err_entered_func.get_or_insert(InterGraphEdgeError::ValueEdgeIntoFunc {
+            // Dominator edge
+            let from_grandparent_op = self.hugr.get_optype(fpp);
+            if from_grandparent_op.tag() != OpTag::Cfg {
+                return Err(InterGraphEdgeError::NonCFGAncestor {
+                    from,
+                    from_offset,
                     to,
                     to_offset,
+                    ancestor_parent_op: from_grandparent_op.clone(),
+                });
+            }
+            // Check domination
+            let (dominator_tree, node_map) = self.dominators.entry(fpp).or_insert_with(|| {
+                let (region, node_map) = self.hugr.region_portgraph(fpp);
+                let entry_node = self.hugr.children(fpp).next().unwrap();
+                let doms = dominators::simple_fast(&region, node_map.to_portgraph(entry_node));
+                (doms, node_map)
+            });
+            let ancestor = containing_child(self.hugr, fpp, to);
+            if !dominator_tree
+                .dominators(node_map.to_portgraph(ancestor))
+                .is_some_and(|mut ds| ds.any(|n| n == node_map.to_portgraph(from_parent)))
+            {
+                return Err(InterGraphEdgeError::NonDominatedAncestor {
                     from,
                     from_offset,
-                    func: ancestor,
+                    to,
+                    to_offset,
+                    from_parent,
+                    ancestor,
                 });
             }
-            if ancestor_parent == from_parent {
-                // External edge.
-                err_entered_func.map_or(Ok(()), Err)?;
-                if !is_static {
-                    // Must have an order edge.
-                    self.hugr
-                        .node_connections(from, ancestor)
-                        .find(|&[p, _]| from_optype.port_kind(p) == Some(EdgeKind::StateOrder))
-                        .ok_or(InterGraphEdgeError::MissingOrderEdge {
-                            from,
-                            from_offset,
-                            to,
-                            to_offset,
-                            to_ancestor: ancestor,
-                        })?;
-                }
-                return Ok(());
-            } else if Some(ancestor_parent) == from_parent_parent && !is_static {
-                // Dominator edge
-                let ancestor_parent_op = self.hugr.get_optype(ancestor_parent);
-                if ancestor_parent_op.tag() != OpTag::Cfg {
-                    return Err(InterGraphEdgeError::NonCFGAncestor {
-                        from,
-                        from_offset,
-                        to,
-                        to_offset,
-                        ancestor_parent_op: ancestor_parent_op.clone(),
-                    });
-                }
-                err_entered_func.map_or(Ok(()), Err)?;
-                // Check domination
-                let (dominator_tree, node_map) =
-                    if let Some(tree) = self.dominators.get(&ancestor_parent) {
-                        tree
-                    } else {
-                        let (tree, node_map) = self.compute_dominator(ancestor_parent);
-                        self.dominators.insert(ancestor_parent, (tree, node_map));
-                        self.dominators.get(&ancestor_parent).unwrap()
-                    };
-                if !dominator_tree
-                    .dominators(node_map.to_portgraph(ancestor))
-                    .is_some_and(|mut ds| ds.any(|n| n == node_map.to_portgraph(from_parent)))
-                {
-                    return Err(InterGraphEdgeError::NonDominatedAncestor {
-                        from,
-                        from_offset,
-                        to,
-                        to_offset,
-                        from_parent,
-                        ancestor,
-                    });
-                }
+        } else {
+            return Err(InterGraphEdgeError::NoRelation {
+                from,
+                from_offset,
+                to,
+                to_offset,
+            });
+        }
 
-                return Ok(());
+        if let Some(func) = ht.nearest_enclosing_funcdefn(to) {
+            if !ht.is_strict_ancestor_of(func, from) {
+                return Err(InterGraphEdgeError::ValueEdgeIntoFunc {
+                    to,
+                    to_offset,
+                    from,
+                    from_offset,
+                    func,
+                });
             }
         }
-
-        Err(InterGraphEdgeError::NoRelation {
-            from,
-            from_offset,
-            to,
-            to_offset,
-        })
+        Ok(())
     }
 
     /// Validates that `TypeArgs` are valid wrt the [`ExtensionRegistry`] and that nodes