dirstate-tree: Add tree traversal/iteration
authorSimon Sapin <simon.sapin@octobus.net>
Tue, 06 Apr 2021 21:07:12 +0200
changeset 47114 caa3031c9ed5
parent 47113 3da19db33cbc
child 47115 5d62243c7732
dirstate-tree: Add tree traversal/iteration Like Python’s, Rust’s iterators are "external" in that they are driven by a caller who calls a `next` method. This is as opposed to "internal" iterators who drive themselves and call a callback for each item. Writing an internal iterator traversing a tree is easy with recursion, but internal iterators cannot rely on the call stack in that way, they must save in an explicit object all state that they need to be preserved across two `next` calls. This algorithm uses a `Vec` as a stack that contains what would be local variables on the call stack if we could use recursion. Differential Revision: https://phab.mercurial-scm.org/D10370
rust/hg-core/src/dirstate_tree/dirstate_map.rs
--- a/rust/hg-core/src/dirstate_tree/dirstate_map.rs	Tue Apr 06 14:35:39 2021 +0200
+++ b/rust/hg-core/src/dirstate_tree/dirstate_map.rs	Tue Apr 06 21:07:12 2021 +0200
@@ -43,6 +43,13 @@
     children: ChildNodes,
 }
 
+/// `(full_path, entry, copy_source)`
+type NodeDataMut<'a> = (
+    &'a WithBasename<HgPathBuf>,
+    &'a mut Option<DirstateEntry>,
+    &'a mut Option<HgPathBuf>,
+);
+
 impl DirstateMap {
     pub fn new() -> Self {
         Self {
@@ -95,6 +102,87 @@
             }
         }
     }
+
+    fn iter_nodes<'a>(
+        &'a self,
+    ) -> impl Iterator<Item = (&'a WithBasename<HgPathBuf>, &'a Node)> + 'a
+    {
+        // Depth first tree traversal.
+        //
+        // If we could afford internal iteration and recursion,
+        // this would look like:
+        //
+        // ```
+        // fn traverse_children(
+        //     children: &ChildNodes,
+        //     each: &mut impl FnMut(&Node),
+        // ) {
+        //     for child in children.values() {
+        //         traverse_children(&child.children, each);
+        //         each(child);
+        //     }
+        // }
+        // ```
+        //
+        // However we want an external iterator and therefore can’t use the
+        // call stack. Use an explicit stack instead:
+        let mut stack = Vec::new();
+        let mut iter = self.root.iter();
+        std::iter::from_fn(move || {
+            while let Some((key, child_node)) = iter.next() {
+                // Pseudo-recursion
+                let new_iter = child_node.children.iter();
+                let old_iter = std::mem::replace(&mut iter, new_iter);
+                stack.push((key, child_node, old_iter));
+            }
+            // Found the end of a `children.iter()` iterator.
+            if let Some((key, child_node, next_iter)) = stack.pop() {
+                // "Return" from pseudo-recursion by restoring state from the
+                // explicit stack
+                iter = next_iter;
+
+                Some((key, child_node))
+            } else {
+                // Reached the bottom of the stack, we’re done
+                None
+            }
+        })
+    }
+
+    /// Mutable iterator for the `(entry, copy source)` of each node.
+    ///
+    /// It would not be safe to yield mutable references to nodes themeselves
+    /// with `-> impl Iterator<Item = &mut Node>` since child nodes are
+    /// reachable from their ancestor nodes, potentially creating multiple
+    /// `&mut` references to a given node.
+    fn iter_node_data_mut<'a>(
+        &'a mut self,
+    ) -> impl Iterator<Item = NodeDataMut<'a>> + 'a {
+        // Explict stack for pseudo-recursion, see `iter_nodes` above.
+        let mut stack = Vec::new();
+        let mut iter = self.root.iter_mut();
+        std::iter::from_fn(move || {
+            while let Some((key, child_node)) = iter.next() {
+                // Pseudo-recursion
+                let data =
+                    (key, &mut child_node.entry, &mut child_node.copy_source);
+                let new_iter = child_node.children.iter_mut();
+                let old_iter = std::mem::replace(&mut iter, new_iter);
+                stack.push((data, old_iter));
+            }
+            // Found the end of a `children.values_mut()` iterator.
+            if let Some((data, next_iter)) = stack.pop() {
+                // "Return" from pseudo-recursion by restoring state from the
+                // explicit stack
+                iter = next_iter;
+
+                Some(data)
+            } else {
+                // Reached the bottom of the stack, we’re done
+                None
+            }
+        })
+    }
 }
 
 impl super::dispatch::DirstateMapMethods for DirstateMap {
@@ -242,6 +330,7 @@
         _parents: DirstateParents,
         _now: Duration,
     ) -> Result<Vec<u8>, DirstateError> {
+        let _ = self.iter_node_data_mut();
         todo!()
     }
 
@@ -278,7 +367,11 @@
     }
 
     fn copy_map_iter(&self) -> CopyMapIter<'_> {
-        todo!()
+        Box::new(self.iter_nodes().filter_map(|(path, node)| {
+            node.copy_source
+                .as_ref()
+                .map(|copy_source| (path.full_path(), copy_source))
+        }))
     }
 
     fn copy_map_contains_key(&self, key: &HgPath) -> bool {
@@ -318,6 +411,8 @@
     }
 
     fn iter(&self) -> StateMapIter<'_> {
-        todo!()
+        Box::new(self.iter_nodes().filter_map(|(path, node)| {
+            node.entry.as_ref().map(|entry| (path.full_path(), entry))
+        }))
     }
 }