From 935399cba02c2246dafaf776af86371ca7f923d7 Mon Sep 17 00:00:00 2001
From: Daniel Ploch <dploch@google.com>
Date: Fri, 20 Jan 2023 18:05:52 -0500
Subject: [PATCH] repo-path-tree: make public and add an iterator API

The internal 'set-of-dirs' structure is useful to extensions that need to do
their own tree-walking routines. It may also be useful internally in other
places.
---
 lib/src/matchers.rs | 128 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 122 insertions(+), 6 deletions(-)

diff --git a/lib/src/matchers.rs b/lib/src/matchers.rs
index 6498ddb186..7d5b8ab968 100644
--- a/lib/src/matchers.rs
+++ b/lib/src/matchers.rs
@@ -14,7 +14,7 @@
 
 #![allow(dead_code, missing_docs)]
 
-use std::collections::{HashMap, HashSet};
+use std::collections::{hash_map, HashMap, HashSet, VecDeque};
 use std::iter;
 
 use tracing::instrument;
@@ -271,7 +271,7 @@ struct RepoPathTree {
 }
 
 impl RepoPathTree {
-    fn new() -> Self {
+    pub fn new() -> Self {
         RepoPathTree {
             entries: HashMap::new(),
             is_dir: false,
@@ -290,15 +290,15 @@ impl RepoPathTree {
         })
     }
 
-    fn add_dir(&mut self, dir: &RepoPath) {
+    pub fn add_dir(&mut self, dir: &RepoPath) {
         self.add(dir).is_dir = true;
     }
 
-    fn add_file(&mut self, file: &RepoPath) {
+    pub fn add_file(&mut self, file: &RepoPath) {
         self.add(file).is_file = true;
     }
 
-    fn get(&self, dir: &RepoPath) -> Option<&RepoPathTree> {
+    pub fn get(&self, dir: &RepoPath) -> Option<&RepoPathTree> {
         dir.components()
             .iter()
             .try_fold(self, |sub, name| sub.entries.get(name))
@@ -310,7 +310,7 @@ impl RepoPathTree {
             .unwrap_or(Visit::Nothing)
     }
 
-    fn walk_to<'a>(
+    pub fn walk_to<'a>(
         &'a self,
         dir: &'a RepoPath,
     ) -> impl Iterator<Item = (&RepoPathTree, &[RepoPathComponent])> + 'a {
@@ -338,6 +338,87 @@ impl RepoPathTree {
     }
 }
 
+pub struct RepoPathTreeIterator<'a> {
+    tree: &'a RepoPathTree,
+    map_iter: Option<hash_map::Iter<'a, RepoPathComponent, RepoPathTree>>,
+    map_component: Option<&'a RepoPathComponent>,
+    tree_iter: Option<Box<RepoPathTreeIterator<'a>>>,
+}
+
+impl<'a> RepoPathTreeIterator<'a> {
+    fn new(tree: &'a RepoPathTree) -> Self {
+        Self {
+            tree,
+            map_iter: None,
+            map_component: None,
+            tree_iter: None,
+        }
+    }
+}
+
+pub struct RepoPathTreeElement<'a> {
+    path: VecDeque<&'a RepoPathComponent>,
+    is_dir: bool,
+    is_file: bool,
+}
+
+impl<'a> RepoPathTreeElement<'a> {
+    pub fn repo_path(&self) -> RepoPath {
+        RepoPath::from_components(self.path.iter().cloned().cloned().collect())
+    }
+
+    pub fn is_dir(&self) -> bool {
+        self.is_dir
+    }
+
+    pub fn is_file(&self) -> bool {
+        self.is_file
+    }
+}
+
+impl<'a> Iterator for RepoPathTreeIterator<'a> {
+    type Item = RepoPathTreeElement<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(map_iter) = self.map_iter.as_mut() {
+            if let Some(tree_iter) = self.tree_iter.as_mut() {
+                if let Some(mut child) = tree_iter.next() {
+                    let map_component = *self.map_component.as_ref().unwrap();
+                    child.path.push_front(map_component);
+                    return Some(child);
+                }
+            }
+
+            if let Some((component, tree)) = map_iter.next() {
+                self.map_component = Some(component);
+                let mut iter = tree.into_iter();
+                let mut first = iter.next().unwrap();
+                first.path.push_front(component);
+                self.tree_iter = Some(Box::new(iter));
+                Some(first)
+            } else {
+                None
+            }
+        } else {
+            self.map_iter = Some(self.tree.entries.iter());
+            Some(Self::Item {
+                path: Default::default(),
+                is_dir: self.tree.is_dir,
+                is_file: self.tree.is_file,
+            })
+        }
+    }
+}
+
+impl<'a> IntoIterator for &'a RepoPathTree {
+    type Item = RepoPathTreeElement<'a>;
+    type IntoIter = RepoPathTreeIterator<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        Self::IntoIter::new(self)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use maplit::hashset;
@@ -387,6 +468,41 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_repo_path_tree_iterator() {
+        let mut tree = RepoPathTree::new();
+        tree.add_file(&RepoPath::from_internal_string("dir1/dir2/file"));
+        tree.add_file(&RepoPath::from_internal_string("dir1/dir3/file"));
+
+        let elements: Vec<_> = tree.into_iter().collect();
+        assert_eq!(elements.len(), 6);
+        assert_eq!(elements[0].repo_path(), RepoPath::root());
+        assert!(elements[0].is_dir);
+        assert!(!elements[0].is_file);
+        assert_eq!(
+            elements[1].repo_path(),
+            RepoPath::from_internal_string("dir1")
+        );
+        assert!(elements[1].is_dir);
+        assert!(!elements[1].is_file);
+        assert_eq!(
+            hashset! {elements[2].repo_path(), elements[4].repo_path()},
+            hashset! {RepoPath::from_internal_string("dir1/dir2"), RepoPath::from_internal_string("dir1/dir3")}
+        );
+        assert!(elements[2].is_dir);
+        assert!(!elements[2].is_file);
+        assert!(elements[4].is_dir);
+        assert!(!elements[4].is_file);
+        assert_eq!(
+            hashset! {elements[3].repo_path(), elements[5].repo_path()},
+            hashset! {RepoPath::from_internal_string("dir1/dir2/file"), RepoPath::from_internal_string("dir1/dir3/file")}
+        );
+        assert!(!elements[3].is_dir);
+        assert!(elements[3].is_file);
+        assert!(!elements[5].is_dir);
+        assert!(elements[5].is_file);
+    }
+
     #[test]
     fn test_nothingmatcher() {
         let m = NothingMatcher;