diff --git a/lib/src/backend.rs b/lib/src/backend.rs index f5b801c98..2b912a54d 100644 --- a/lib/src/backend.rs +++ b/lib/src/backend.rs @@ -276,11 +276,11 @@ impl<'a> TreeEntry<'a> { } } -pub struct TreeEntriesNonRecursiveIter<'a> { +pub struct TreeEntriesNonRecursiveIterator<'a> { iter: std::collections::btree_map::Iter<'a, RepoPathComponent, TreeValue>, } -impl<'a> Iterator for TreeEntriesNonRecursiveIter<'a> { +impl<'a> Iterator for TreeEntriesNonRecursiveIterator<'a> { type Item = TreeEntry<'a>; fn next(&mut self) -> Option { @@ -300,8 +300,8 @@ impl Tree { self.entries.is_empty() } - pub fn entries(&self) -> TreeEntriesNonRecursiveIter { - TreeEntriesNonRecursiveIter { + pub fn entries(&self) -> TreeEntriesNonRecursiveIterator { + TreeEntriesNonRecursiveIterator { iter: self.entries.iter(), } } diff --git a/lib/src/tree.rs b/lib/src/tree.rs index 3d69963e4..29106488a 100644 --- a/lib/src/tree.rs +++ b/lib/src/tree.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use itertools::Itertools; use crate::backend::{ - BackendError, Conflict, ConflictId, ConflictPart, TreeEntriesNonRecursiveIter, TreeEntry, + BackendError, Conflict, ConflictId, ConflictPart, TreeEntriesNonRecursiveIterator, TreeEntry, TreeId, TreeValue, }; use crate::files::MergeResult; @@ -96,12 +96,12 @@ impl Tree { &self.data } - pub fn entries_non_recursive(&self) -> TreeEntriesNonRecursiveIter { + pub fn entries_non_recursive(&self) -> TreeEntriesNonRecursiveIterator { self.data.entries() } - pub fn entries(&self) -> TreeEntriesIter { - TreeEntriesIter::new(self.clone()) + pub fn entries(&self) -> TreeEntriesIterator { + TreeEntriesIterator::new(self.clone()) } pub fn entry(&self, basename: &RepoPathComponent) -> Option { @@ -206,51 +206,53 @@ impl Tree { } } -pub struct TreeEntriesIter { - stack: Vec<(Pin>, TreeEntriesNonRecursiveIter<'static>)>, +pub struct TreeEntriesIterator { + tree: Pin>, + entry_iterator: TreeEntriesNonRecursiveIterator<'static>, + subdir_iterator: Option>, } -impl TreeEntriesIter { +impl TreeEntriesIterator { fn new(tree: Tree) -> Self { let tree = Box::pin(tree); - let iter = tree.entries_non_recursive(); - let iter: TreeEntriesNonRecursiveIter<'static> = unsafe { std::mem::transmute(iter) }; + let entry_iterator = tree.entries_non_recursive(); + let entry_iterator: TreeEntriesNonRecursiveIterator<'static> = + unsafe { std::mem::transmute(entry_iterator) }; Self { - stack: vec![(tree, iter)], + tree, + entry_iterator, + subdir_iterator: None, } } } -impl Iterator for TreeEntriesIter { +impl Iterator for TreeEntriesIterator { type Item = (RepoPath, TreeValue); fn next(&mut self) -> Option { - 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(entry.name(), id); - let subtree = Box::pin(subtree); - let iter = subtree.entries_non_recursive(); - let subtree_iter: TreeEntriesNonRecursiveIter<'static> = - unsafe { std::mem::transmute(iter) }; - self.stack.push((subtree, subtree_iter)); - } - other => { - let path = tree.dir().join(entry.name()); - return Some((path, other.clone())); - } - }; + loop { + // First return results from any subdirectory we're currently visiting. + if let Some(subdir_iter) = &mut self.subdir_iterator { + if let Some(item) = subdir_iter.next() { + return Some(item); } + self.subdir_iterator = None; + } + if let Some(entry) = self.entry_iterator.next() { + match entry.value() { + TreeValue::Tree(id) => { + let subtree = self.tree.known_sub_tree(entry.name(), id); + self.subdir_iterator = Some(Box::new(TreeEntriesIterator::new(subtree))); + } + other => { + let path = self.tree.dir().join(entry.name()); + return Some((path, other.clone())); + } + }; + } else { + return None; } } - None } } @@ -280,8 +282,8 @@ impl Diff { } struct TreeEntryDiffIterator<'trees, 'matcher> { - it1: Peekable>, - it2: Peekable>, + it1: Peekable>, + it2: Peekable>, // TODO: Restrict walk according to Matcher::visit() _matcher: &'matcher dyn Matcher, } @@ -420,6 +422,7 @@ impl Iterator for TreeDiffIterator<'_> { if let Some(element) = subdir_iterator.next() { return Some(element); } + self.subdir_iterator = None; } if let Some((name, value)) = self.added_file.take() {