From 8ec100713d562badebf5f967852dfdb5f1811c9d Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Sun, 20 Dec 2020 00:07:51 -0800
Subject: [PATCH] tree: for walking tree, replace function with callback by
 iterator

Iterators are a lot easier to use.
---
 lib/src/tree.rs  | 69 +++++++++++++++++++++++++++++++++++++++++++-----
 lib/src/trees.rs | 22 +--------------
 src/commands.rs  |  8 +++---
 3 files changed, 68 insertions(+), 31 deletions(-)

diff --git a/lib/src/tree.rs b/lib/src/tree.rs
index 35e03acc4..f17d0b827 100644
--- a/lib/src/tree.rs
+++ b/lib/src/tree.rs
@@ -17,11 +17,14 @@ use std::fmt::{Debug, Error, Formatter};
 use std::sync::Arc;
 
 use crate::matchers::AlwaysMatcher;
-use crate::repo_path::{DirRepoPath, DirRepoPathComponent, FileRepoPath, RepoPath, RepoPathJoin};
+use crate::repo_path::{
+    DirRepoPath, DirRepoPathComponent, FileRepoPath, RepoPath, RepoPathComponent, RepoPathJoin,
+};
 use crate::store;
 use crate::store::{ConflictId, TreeEntriesIter, TreeEntry, TreeId, TreeValue};
 use crate::store_wrapper::StoreWrapper;
-use crate::trees::{recursive_tree_diff, walk_entries, TreeValueDiff};
+use crate::trees::{recursive_tree_diff, TreeValueDiff};
+use std::pin::Pin;
 
 #[derive(Clone)]
 pub struct Tree {
@@ -91,6 +94,10 @@ impl Tree {
         self.data.entries()
     }
 
+    pub fn entries_recursive(&self) -> TreeWalk {
+        TreeWalk::new(self.clone())
+    }
+
     pub fn entry<N>(&self, basename: &N) -> Option<TreeEntry>
     where
         N: Borrow<str> + ?Sized,
@@ -187,13 +194,63 @@ impl Tree {
 
     pub fn conflicts(&self) -> Vec<(RepoPath, ConflictId)> {
         let mut conflicts = vec![];
-        walk_entries(&self, &mut |name, value| -> Result<(), ()> {
+        for (name, value) in self.entries_recursive() {
             if let TreeValue::Conflict(id) = value {
                 conflicts.push((name.clone(), id.clone()));
             }
-            Ok(())
-        })
-        .unwrap();
+        }
         conflicts
     }
 }
+
+pub struct TreeWalk {
+    stack: Vec<(Pin<Box<Tree>>, TreeEntriesIter<'static>)>,
+}
+
+impl TreeWalk {
+    fn new(tree: Tree) -> Self {
+        let tree = Box::pin(tree);
+        let iter = tree.entries();
+        let iter: TreeEntriesIter<'static> = unsafe { std::mem::transmute(iter) };
+        Self {
+            stack: vec![(tree, iter)],
+        }
+    }
+}
+
+impl Iterator for TreeWalk {
+    type Item = (RepoPath, TreeValue);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while !self.stack.is_empty() {
+            let (tree, iter) = self.stack.last_mut().unwrap();
+            match iter.next() {
+                None => {
+                    // No more entries in this directory
+                    self.stack.pop().unwrap();
+                }
+                Some(entry) => {
+                    match entry.value() {
+                        TreeValue::Tree(id) => {
+                            let subtree =
+                                tree.known_sub_tree(&DirRepoPathComponent::from(entry.name()), id);
+                            let subtree = Box::pin(subtree);
+                            let iter = subtree.entries();
+                            let subtree_iter: TreeEntriesIter<'static> =
+                                unsafe { std::mem::transmute(iter) };
+                            self.stack.push((subtree, subtree_iter));
+                        }
+                        other => {
+                            let path = RepoPath::new(
+                                tree.dir().clone(),
+                                RepoPathComponent::from(entry.name()),
+                            );
+                            return Some((path, other.clone()));
+                        }
+                    };
+                }
+            }
+        }
+        None
+    }
+}
diff --git a/lib/src/trees.rs b/lib/src/trees.rs
index 54ac1b688..9267eab5d 100644
--- a/lib/src/trees.rs
+++ b/lib/src/trees.rs
@@ -16,33 +16,13 @@ use crate::files;
 use crate::files::MergeResult;
 use crate::matchers::Matcher;
 use crate::repo_path::{
-    DirRepoPath, DirRepoPathComponent, FileRepoPath, FileRepoPathComponent, RepoPath,
-    RepoPathComponent, RepoPathJoin,
+    DirRepoPath, DirRepoPathComponent, FileRepoPath, FileRepoPathComponent, RepoPathJoin,
 };
 use crate::store::{Conflict, ConflictPart, StoreError, TreeId, TreeValue};
 use crate::store_wrapper::StoreWrapper;
 use crate::tree::Tree;
 use std::cmp::Ordering;
 
-pub fn walk_entries<E>(
-    tree: &Tree,
-    callback: &mut impl FnMut(&RepoPath, &TreeValue) -> Result<(), E>,
-) -> Result<(), E> {
-    for entry in tree.entries() {
-        let path = RepoPath::new(tree.dir().clone(), RepoPathComponent::from(entry.name()));
-        match entry.value() {
-            TreeValue::Tree(id) => {
-                let subtree = tree.known_sub_tree(&DirRepoPathComponent::from(entry.name()), id);
-                walk_entries(&subtree, callback)?;
-            }
-            other => {
-                callback(&path, other)?;
-            }
-        };
-    }
-    Ok(())
-}
-
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Diff<T> {
     Modified(T, T),
diff --git a/src/commands.rs b/src/commands.rs
index 4315f8bb5..fc360c68c 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -47,7 +47,7 @@ use jj_lib::settings::UserSettings;
 use jj_lib::store::{CommitId, Timestamp};
 use jj_lib::store::{StoreError, TreeValue};
 use jj_lib::tree::Tree;
-use jj_lib::trees::{walk_entries, TreeValueDiff};
+use jj_lib::trees::TreeValueDiff;
 use jj_lib::working_copy::{CheckoutStats, WorkingCopy};
 
 use self::chrono::{FixedOffset, TimeZone, Utc};
@@ -565,10 +565,10 @@ fn cmd_files(
     let mut repo = get_repo(ui, &matches)?;
     let mut_repo = Arc::get_mut(&mut repo).unwrap();
     let commit = resolve_revision_arg(ui, mut_repo, sub_matches)?;
-    walk_entries(&commit.tree(), &mut |name, _value| {
+    for (name, _value) in commit.tree().entries_recursive() {
         writeln!(ui, "{}", name.to_internal_string());
-        Ok(())
-    })
+    }
+    Ok(())
 }
 
 fn print_diff(left: &[u8], right: &[u8], styler: &mut dyn Styler) {