tree: rewrite recursive diff iterator to not use machine stack

While mutable borrow in the next() method is a bit tricky, I think it's
easier to follow than invoking iterators recursively.

lib/src/tree.rs

@@ -387,76 +387,59 @@ pub fn recursive_tree_diff(root1: Tree, root2: Tree, matcher: &dyn Matcher) -> T
     TreeDiffIterator::new(RepoPath::root(), root1, root2, matcher)
 }
 
-enum MaybeTreeEntryDiffIterator<'trees> {
-    WalkEntries(TreeEntryDiffIterator<'trees>),
-    Nothing,
-}
-
-impl<'trees> Iterator for MaybeTreeEntryDiffIterator<'trees> {
-    type Item = (
-        &'trees RepoPathComponent,
-        Option<&'trees TreeValue>,
-        Option<&'trees TreeValue>,
-    );
-
-    fn next(&mut self) -> Option<Self::Item> {
-        match self {
-            MaybeTreeEntryDiffIterator::WalkEntries(it) => it.next(),
-            MaybeTreeEntryDiffIterator::Nothing => None,
-        }
-    }
-}
-
 pub struct TreeDiffIterator<'matcher> {
-    dir: RepoPath,
+    stack: Vec<TreeDiffItem>,
     matcher: &'matcher dyn Matcher,
+}
+
+struct TreeDiffDirItem {
+    path: RepoPath,
     // Iterator over the diffs between tree1 and tree2
-    entry_iterator: MaybeTreeEntryDiffIterator<'static>,
+    entry_iterator: TreeEntryDiffIterator<'static>,
     // On drop, tree1 and tree2 must outlive entry_iterator
     tree1: Pin<Box<Tree>>,
     tree2: Pin<Box<Tree>>,
+}
+
+enum TreeDiffItem {
+    Dir(TreeDiffDirItem),
     // This is used for making sure that when a directory gets replaced by a file, we
     // yield the value for the addition of the file after we yield the values
     // for removing files in the directory.
-    added_file: Option<(RepoPath, TreeValue)>,
+    File(RepoPath, Diff<TreeValue>),
-    // Iterator over the diffs of a subdirectory, if we're currently visiting one.
-    subdir_iterator: Option<Box<TreeDiffIterator<'matcher>>>,
 }
 
 impl<'matcher> TreeDiffIterator<'matcher> {
-    fn new(
+    fn new(dir: RepoPath, tree1: Tree, tree2: Tree, matcher: &'matcher dyn Matcher) -> Self {
-        dir: RepoPath,
+        let mut stack = Vec::new();
-        tree1: Tree,
+        if !matcher.visit(&dir).is_nothing() {
-        tree2: Tree,
+            stack.push(TreeDiffItem::Dir(TreeDiffDirItem::new(dir, tree1, tree2)));
-        matcher: &'matcher dyn Matcher,
+        };
-    ) -> TreeDiffIterator {
+        Self { stack, matcher }
+    }
+}
+
+impl TreeDiffDirItem {
+    fn new(path: RepoPath, tree1: Tree, tree2: Tree) -> Self {
         let tree1 = Box::pin(tree1);
         let tree2 = Box::pin(tree2);
-        let root_entry_iterator = if matcher.visit(&dir).is_nothing() {
+        let iter: TreeEntryDiffIterator = diff_entries(&tree1, &tree2);
-            MaybeTreeEntryDiffIterator::Nothing
+        let iter: TreeEntryDiffIterator<'static> = unsafe { std::mem::transmute(iter) };
-        } else {
-            let iter: TreeEntryDiffIterator = diff_entries(&tree1, &tree2);
-            let iter: TreeEntryDiffIterator<'static> = unsafe { std::mem::transmute(iter) };
-            MaybeTreeEntryDiffIterator::WalkEntries(iter)
-        };
         Self {
-            dir,
+            path,
-            matcher,
+            entry_iterator: iter,
-            entry_iterator: root_entry_iterator,
             tree1,
             tree2,
-            added_file: None,
-            subdir_iterator: None,
         }
     }
 
-    fn visit_subdir(
+    fn subdir(
-        &mut self,
+        &self,
         name: &RepoPathComponent,
         before: Option<&TreeValue>,
         after: Option<&TreeValue>,
-    ) {
+    ) -> Self {
-        let subdir_path = self.dir.join(name);
+        let subdir_path = self.path.join(name);
         let before_tree = match before {
             Some(TreeValue::Tree(id_before)) => self.tree1.known_sub_tree(name, id_before),
             _ => Tree::null(self.tree1.store().clone(), subdir_path.clone()),
@@ -465,12 +448,7 @@ impl<'matcher> TreeDiffIterator<'matcher> {
             Some(TreeValue::Tree(id_after)) => self.tree2.known_sub_tree(name, id_after),
             _ => Tree::null(self.tree2.store().clone(), subdir_path.clone()),
         };
-        self.subdir_iterator = Some(Box::new(TreeDiffIterator::new(
+        Self::new(subdir_path, before_tree, after_tree)
-            subdir_path,
-            before_tree,
-            after_tree,
-            self.matcher,
-        )));
     }
 }
 
@@ -478,27 +456,37 @@ impl Iterator for TreeDiffIterator<'_> {
     type Item = (RepoPath, Diff<TreeValue>);
 
     fn next(&mut self) -> Option<Self::Item> {
-        loop {
+        while let Some(top) = self.stack.last_mut() {
-            // First return results from any subdirectory we're currently visiting.
+            let (dir, (name, before, after)) = match top {
-            if let Some(subdir_iterator) = &mut self.subdir_iterator {
+                TreeDiffItem::Dir(dir) => {
-                if let Some(element) = subdir_iterator.next() {
+                    if let Some(entry) = dir.entry_iterator.next() {
-                    return Some(element);
+                        (dir, entry)
+                    } else {
+                        self.stack.pop().unwrap();
+                        continue;
+                    }
                 }
-                self.subdir_iterator = None;
+                TreeDiffItem::File(..) => {
-            }
+                    if let TreeDiffItem::File(name, diff) = self.stack.pop().unwrap() {
-
+                        return Some((name, diff));
-            if let Some((name, value)) = self.added_file.take() {
+                    } else {
-                return Some((name, Diff::Added(value)));
+                        unreachable!();
-            }
+                    }
+                }
+            };
 
             // Note: whenever we say "file" below, it may also be a symlink or a conflict.
-            let (name, before, after) = self.entry_iterator.next()?;
+            let file_path = dir.path.join(name);
             let tree_before = matches!(before, Some(TreeValue::Tree(_)));
             let tree_after = matches!(after, Some(TreeValue::Tree(_)));
-            if tree_before || tree_after {
+            let post_subdir =
-                self.visit_subdir(name, before, after);
+                if (tree_before || tree_after) && !self.matcher.visit(&file_path).is_nothing() {
-            }
+                    let subdir = dir.subdir(name, before, after);
-            let file_path = self.dir.join(name);
+                    self.stack.push(TreeDiffItem::Dir(subdir));
+                    self.stack.len() - 1
+                } else {
+                    self.stack.len()
+                };
             if self.matcher.matches(&file_path) {
                 if !tree_before && tree_after {
                     if let Some(file_before) = before {
@@ -506,7 +494,10 @@ impl Iterator for TreeDiffIterator<'_> {
                     }
                 } else if tree_before && !tree_after {
                     if let Some(file_after) = after {
-                        self.added_file = Some((file_path, file_after.clone()));
+                        self.stack.insert(
+                            post_subdir,
+                            TreeDiffItem::File(file_path, Diff::Added(file_after.clone())),
+                        );
                     }
                 } else if !tree_before && !tree_after {
                     match (before, after) {
@@ -529,6 +520,7 @@ impl Iterator for TreeDiffIterator<'_> {
                 }
             }
         }
+        None
     }
 }