feat(persistent): More efficient HugrView iterators for PersistentHugr

lmondada · lmondada · commit d02aa6ca2b8c · 2025-09-30T11:44:45.000+01:00
diff --git a/hugr-persistent/src/persistent_hugr.rs b/hugr-persistent/src/persistent_hugr.rs
@@ -1,4 +1,5 @@
 use std::{
+    cell::RefCell,
     collections::{BTreeSet, HashMap, VecDeque},
     vec,
 };
@@ -14,6 +15,9 @@ use crate::{
     Commit, CommitData, CommitId, CommitStateSpace, InvalidCommit, PatchNode, PersistentReplacement,
 };
 
+mod cache;
+use cache::PersistentHugrCache;
+
 pub mod serial;
 
 /// A HUGR-like object that tracks its mutation history.
@@ -56,7 +60,7 @@ pub mod serial;
 /// subgraphs within dataflow regions are supported.
 ///
 /// [`SimpleReplacement`]: hugr_core::SimpleReplacement
-#[derive(Clone, Debug)]
+#[derive(Clone, derive_more::Debug)]
 pub struct PersistentHugr {
     /// The state space of all commits.
     ///
@@ -72,6 +76,9 @@ pub struct PersistentHugr {
     /// other commits are [`CommitData::Replacement`]s, and are descendants
     /// of this.
     base_commit_id: CommitId,
+    /// Cache some properties for performance.
+    #[debug(skip)]
+    cache: RefCell<PersistentHugrCache>,
 }
 
 impl PersistentHugr {
@@ -144,6 +151,7 @@ impl PersistentHugr {
         Ok(Self {
             graph,
             base_commit_id: base_commit,
+            cache: RefCell::new(PersistentHugrCache::default()),
         })
     }
 
@@ -247,6 +255,11 @@ impl PersistentHugr {
             }
 
             self.graph.insert_node(new_commit.clone().into());
+
+            // Invalidate cache
+            for parent in new_commit.parents() {
+                self.cache.borrow_mut().invalidate_children(parent.id());
+            }
         }
 
         Ok(())
@@ -331,7 +344,11 @@ impl PersistentHugr {
     }
 
     pub fn children_commits(&self, commit_id: CommitId) -> impl Iterator<Item = CommitId> + '_ {
-        self.graph.children(commit_id)
+        self.cache
+            .borrow_mut()
+            .children_or_insert(commit_id, || self.graph.children(commit_id).collect())
+            .clone()
+            .into_iter()
     }
 
     pub fn parent_commits(&self, commit_id: CommitId) -> impl Iterator<Item = CommitId> + '_ {
@@ -461,6 +478,21 @@ fn get_ancestors_while<'a>(
     all_commits
 }
 
+/// A node in a commit of a [`PersistentHugr`] is either a valid node of the
+/// HUGR, a node deleted by a child commit in that [`PersistentHugr`], or an
+/// input or output node in a replacement graph.
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum NodeStatus {
+    /// A node deleted by a child commit in that [`PersistentHugr`].
+    ///
+    /// The ID of the child commit is stored in the variant.
+    Deleted(CommitId),
+    /// An input or output node in the replacement graph of a Commit
+    ReplacementIO,
+    /// A valid node in the [`PersistentHugr`]
+    Valid,
+}
+
 // non-public methods
 impl PersistentHugr {
     /// Convert a node ID specific to a commit HUGR into a patch node in the
@@ -480,6 +512,26 @@ impl PersistentHugr {
             child.deleted_parent_nodes().contains(&node)
         })
     }
+
+    /// Whether a node is valid in `self`, is deleted or is an IO node in a
+    /// replacement graph.
+    pub(crate) fn node_status(
+        &self,
+        per_node @ PatchNode(commit_id, node): PatchNode,
+    ) -> NodeStatus {
+        debug_assert!(self.contains_id(commit_id), "unknown commit");
+        if self
+            .get_commit(commit_id)
+            .replacement()
+            .is_some_and(|repl| repl.get_replacement_io().contains(&node))
+        {
+            NodeStatus::ReplacementIO
+        } else if let Some(commit_id) = self.find_deleting_commit(per_node) {
+            NodeStatus::Deleted(commit_id)
+        } else {
+            NodeStatus::Valid
+        }
+    }
 }
 
 impl<'a> IntoIterator for &'a PersistentHugr {
diff --git a/hugr-persistent/src/persistent_hugr/cache.rs b/hugr-persistent/src/persistent_hugr/cache.rs
@@ -0,0 +1,23 @@
+use std::collections::HashMap;
+
+use crate::CommitId;
+
+/// A cache for storing computed properties of a `PersistentHugr`.
+#[derive(Debug, Default, Clone)]
+pub(super) struct PersistentHugrCache {
+    children_cache: HashMap<CommitId, Vec<CommitId>>,
+}
+
+impl PersistentHugrCache {
+    pub fn invalidate_children(&mut self, commit: CommitId) {
+        self.children_cache.remove(&commit);
+    }
+
+    pub fn children_or_insert(
+        &mut self,
+        commit: CommitId,
+        children: impl FnOnce() -> Vec<CommitId>,
+    ) -> &Vec<CommitId> {
+        self.children_cache.entry(commit).or_insert_with(children)
+    }
+}
diff --git a/hugr-persistent/src/trait_impls.rs b/hugr-persistent/src/trait_impls.rs
@@ -1,4 +1,7 @@
-use std::collections::HashMap;
+use std::{
+    collections::{BTreeSet, HashMap, VecDeque},
+    iter::FusedIterator,
+};
 
 use itertools::{Either, Itertools};
 
@@ -16,7 +19,7 @@ use hugr_core::{
     ops::{OpTag, OpTrait, OpType},
 };
 
-use crate::CommitId;
+use crate::{CommitId, persistent_hugr::NodeStatus};
 
 use super::{
     InvalidCommit, PatchNode, PersistentHugr, PersistentReplacement, state_space::CommitData,
@@ -100,14 +103,18 @@ impl HugrView for PersistentHugr {
     }
 
     fn get_parent(&self, node: Self::Node) -> Option<Self::Node> {
-        assert!(self.contains_node(node), "invalid node");
-        let (hugr, node_map) = self.apply_all();
-        let parent = hugr.get_parent(node_map[&node])?;
-        let parent_inv = node_map
-            .iter()
-            .find_map(|(&k, &v)| (v == parent).then_some(k))
-            .expect("parent not found in node map");
-        Some(parent_inv)
+        debug_assert!(self.contains_node(node), "invalid node");
+
+        if node.owner() == self.base() {
+            self.base_hugr()
+                .get_parent(node.1)
+                .map(|n| PatchNode(self.base(), n))
+        } else {
+            // all nodes in children commits are applied on the sibling DFG of the
+            // entrypoint
+            // TODO: generalise this for the case that commits introduce nested DFGs.
+            Some(self.entrypoint())
+        }
     }
 
     fn get_optype(&self, PatchNode(commit_id, node): Self::Node) -> &OpType {
@@ -216,24 +223,16 @@ impl HugrView for PersistentHugr {
 
     fn children(&self, node: Self::Node) -> impl DoubleEndedIterator<Item = Self::Node> + Clone {
         // Only children of dataflow parents may change
+        let cm = self.get_commit(node.owner());
+        let commit_hugr = cm.commit_hugr();
+        let children = commit_hugr.children(node.1).map(|n| cm.to_patch_node(n));
         if OpTag::DataflowParent.is_superset(self.get_optype(node).tag()) {
-            // TODO: make this more efficient by traversing from inputs to outputs
-            let (hugr, node_map) = self.apply_all();
-            let children = hugr.children(node_map[&node]).collect_vec();
-            let inv_node_map: HashMap<_, _> = node_map.into_iter().map(|(k, v)| (v, k)).collect();
-            Either::Right(children.into_iter().map(move |child| {
-                *inv_node_map
-                    .get(&child)
-                    .expect("node not found in node map")
-            }))
+            // we must filter out children nodes that are invalidated by later commits, and
+            // on the other hand add nodes in those commits
+            Either::Left(IterValidNodes::new(self, children.fuse()))
         } else {
-            // children are children of the commit hugr
-            let cm = self.get_commit(node.owner());
-            Either::Left(
-                cm.commit_hugr()
-                    .children(node.1)
-                    .map(|n| cm.to_patch_node(n)),
-            )
+            // children are precisely children of the commit hugr
+            Either::Right(children)
         }
     }
 
@@ -341,6 +340,79 @@ impl HugrView for PersistentHugr {
     }
 }
 
+/// An iterator over nodes in a `PersistentHugr` that filters out invalid nodes.
+///
+/// For any invalid node encountered, it will traverse and return the nodes in
+/// the commit deleting the node instead.
+#[derive(Debug, Clone)]
+pub struct IterValidNodes<'a, I> {
+    /// The original iterator over nodes.
+    nodes_iter: I,
+    /// Nodes discovered in commits deleting nodes in the original iterator.
+    discovered_nodes: VecDeque<PatchNode>,
+    /// Commits discovered that delete nodes in the original iterator.
+    discovered_commits: VecDeque<CommitId>,
+    /// Commits discovered across all time, to make sure we only process each
+    /// commit once.
+    processed_commits: BTreeSet<CommitId>,
+    /// The persistent hugr that the nodes belong to.
+    hugr: &'a PersistentHugr,
+}
+
+impl<'a, I> IterValidNodes<'a, I> {
+    fn new(hugr: &'a PersistentHugr, nodes_iter: impl IntoIterator<IntoIter = I>) -> Self {
+        Self {
+            nodes_iter: nodes_iter.into_iter(),
+            discovered_nodes: VecDeque::new(),
+            discovered_commits: VecDeque::new(),
+            processed_commits: BTreeSet::new(),
+            hugr,
+        }
+    }
+}
+
+impl<I: FusedIterator<Item = PatchNode>> Iterator for IterValidNodes<'_, I> {
+    type Item = PatchNode;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let Some(node) = self
+                .nodes_iter
+                .next()
+                .or_else(|| self.discovered_nodes.pop_front())
+            else {
+                break;
+            };
+            match self.hugr.node_status(node) {
+                NodeStatus::Deleted(commit_id) => {
+                    if self.processed_commits.insert(commit_id) {
+                        self.discovered_commits.push_back(commit_id);
+                    }
+                }
+                NodeStatus::ReplacementIO | NodeStatus::Valid => return Some(node),
+            }
+        }
+
+        // Add nodes in next commit to queue
+        let next_commit_id = self.discovered_commits.pop_front()?;
+        let next_commit = self.hugr.get_commit(next_commit_id);
+
+        self.discovered_nodes.extend(
+            next_commit
+                .inserted_nodes()
+                .map(|n| next_commit.to_patch_node(n)),
+        );
+
+        self.next()
+    }
+}
+
+impl<I: FusedIterator<Item = PatchNode>> DoubleEndedIterator for IterValidNodes<'_, I> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        unimplemented!("cannot go backwards")
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashSet;
diff --git a/hugr-persistent/src/wire.rs b/hugr-persistent/src/wire.rs
@@ -6,7 +6,7 @@ use hugr_core::{
 };
 use itertools::Itertools;
 
-use crate::{CommitId, PatchNode, PersistentHugr, Walker};
+use crate::{CommitId, PatchNode, PersistentHugr, Walker, persistent_hugr::NodeStatus};
 
 /// A wire in a [`PersistentHugr`].
 ///
@@ -59,43 +59,11 @@ impl CommitWire {
     }
 }
 
-/// A node in a commit of a [`PersistentHugr`] is either a valid node of the
-/// HUGR, a node deleted by a child commit in that [`PersistentHugr`], or an
-/// input or output node in a replacement graph.
-#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-enum NodeStatus {
-    /// A node deleted by a child commit in that [`PersistentHugr`].
-    ///
-    /// The ID of the child commit is stored in the variant.
-    Deleted(CommitId),
-    /// An input or output node in the replacement graph of a Commit
-    ReplacementIO,
-    /// A valid node in the [`PersistentHugr`]
-    Valid,
-}
-
 impl PersistentHugr {
     pub fn get_wire(&self, node: PatchNode, port: impl Into<Port>) -> PersistentWire {
         PersistentWire::from_port(node, port, self)
     }
 
-    /// Whether a node is valid in `self`, is deleted or is an IO node in a
-    /// replacement graph.
-    fn node_status(&self, per_node @ PatchNode(commit_id, node): PatchNode) -> NodeStatus {
-        debug_assert!(self.contains_id(commit_id), "unknown commit");
-        if self
-            .get_commit(commit_id)
-            .replacement()
-            .is_some_and(|repl| repl.get_replacement_io().contains(&node))
-        {
-            NodeStatus::ReplacementIO
-        } else if let Some(commit_id) = self.find_deleting_commit(per_node) {
-            NodeStatus::Deleted(commit_id)
-        } else {
-            NodeStatus::Valid
-        }
-    }
-
     /// The unique outgoing port in `self` that `port` is attached to.
     ///
     /// # Panics
@@ -132,7 +100,7 @@ impl PersistentHugr {
 impl PersistentWire {
     /// Get the wire connected to a specified port of a pinned node in `hugr`.
     fn from_port(node: PatchNode, port: impl Into<Port>, per_hugr: &PersistentHugr) -> Self {
-        assert!(per_hugr.contains_node(node), "node not in hugr");
+        debug_assert!(per_hugr.contains_node(node), "node not in hugr");
 
         // Queue of wires within each commit HUGR, that combined will form the
         // persistent wire.
