Unify all worktree traversal in a Traversal cursor/iterator

This uses our new ability to have a SeekTarget that's distinct from Dimension.

zed/src/fuzzy.rs

@@ -278,21 +278,19 @@ pub async fn match_paths(
 
                             let start = max(tree_start, segment_start) - tree_start;
                             let end = min(tree_end, segment_end) - tree_start;
-                            let entries = if include_ignored {
-                                snapshot.files(start).take(end - start)
-                            } else {
-                                snapshot.visible_files(start).take(end - start)
-                            };
-                            let paths = entries.map(|entry| {
-                                if let EntryKind::File(char_bag) = entry.kind {
-                                    PathMatchCandidate {
-                                        path: &entry.path,
-                                        char_bag,
+                            let paths = snapshot
+                                .files(include_ignored, start)
+                                .take(end - start)
+                                .map(|entry| {
+                                    if let EntryKind::File(char_bag) = entry.kind {
+                                        PathMatchCandidate {
+                                            path: &entry.path,
+                                            char_bag,
+                                        }
+                                    } else {
+                                        unreachable!()
                                     }
-                                } else {
-                                    unreachable!()
-                                }
-                            });
+                                });
 
                             matcher.match_paths(
                                 snapshot.id(),

zed/src/workspace.rs

@@ -1191,7 +1191,7 @@ impl WorkspaceHandle for ViewHandle<Workspace> {
             .flat_map(|tree| {
                 let tree_id = tree.id();
                 tree.read(cx)
-                    .files(0)
+                    .files(true, 0)
                     .map(move |f| (tree_id, f.path.clone()))
             })
             .collect::<Vec<_>>()

zed/src/worktree.rs

@@ -17,7 +17,7 @@ use futures::{Stream, StreamExt};
 pub use fuzzy::{match_paths, PathMatch};
 use gpui::{
     executor,
-    sum_tree::{self, Cursor, Edit, SumTree},
+    sum_tree::{self, Edit, SeekTarget, SumTree},
     AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, MutableAppContext, Task,
     UpgradeModelHandle, WeakModelHandle,
 };
@@ -1395,26 +1395,28 @@ impl Snapshot {
             .peekable();
         loop {
             match (self_entries.peek(), other_entries.peek()) {
-                (Some(self_entry), Some(other_entry)) => match self_entry.id.cmp(&other_entry.id) {
-                    Ordering::Less => {
-                        let entry = self.entry_for_id(self_entry.id).unwrap().into();
-                        updated_entries.push(entry);
-                        self_entries.next();
-                    }
-                    Ordering::Equal => {
-                        if self_entry.scan_id != other_entry.scan_id {
+                (Some(self_entry), Some(other_entry)) => {
+                    match Ord::cmp(&self_entry.id, &other_entry.id) {
+                        Ordering::Less => {
                             let entry = self.entry_for_id(self_entry.id).unwrap().into();
                             updated_entries.push(entry);
+                            self_entries.next();
                         }
+                        Ordering::Equal => {
+                            if self_entry.scan_id != other_entry.scan_id {
+                                let entry = self.entry_for_id(self_entry.id).unwrap().into();
+                                updated_entries.push(entry);
+                            }
 
-                        self_entries.next();
-                        other_entries.next();
+                            self_entries.next();
+                            other_entries.next();
+                        }
+                        Ordering::Greater => {
+                            removed_entries.push(other_entry.id as u64);
+                            other_entries.next();
+                        }
                     }
-                    Ordering::Greater => {
-                        removed_entries.push(other_entry.id as u64);
-                        other_entries.next();
-                    }
-                },
+                }
                 (Some(self_entry), None) => {
                     let entry = self.entry_for_id(self_entry.id).unwrap().into();
                     updated_entries.push(entry);
@@ -1478,8 +1480,50 @@ impl Snapshot {
         self.entries_by_path.summary().visible_file_count
     }
 
-    pub fn files(&self, start: usize) -> FileIter {
-        FileIter::all(self, start)
+    fn traverse_from_offset(
+        &self,
+        include_dirs: bool,
+        include_ignored: bool,
+        start_offset: usize,
+    ) -> Traversal {
+        let mut cursor = self.entries_by_path.cursor();
+        cursor.seek(
+            &TraversalTarget::Count {
+                count: start_offset,
+                include_dirs,
+                include_ignored,
+            },
+            Bias::Right,
+            &(),
+        );
+        Traversal {
+            cursor,
+            include_dirs,
+            include_ignored,
+        }
+    }
+
+    fn traverse_from_path(
+        &self,
+        include_dirs: bool,
+        include_ignored: bool,
+        path: &Path,
+    ) -> Traversal {
+        let mut cursor = self.entries_by_path.cursor();
+        cursor.seek(&TraversalTarget::Path { path }, Bias::Right, &());
+        Traversal {
+            cursor,
+            include_dirs,
+            include_ignored,
+        }
+    }
+
+    pub fn files(&self, include_ignored: bool, start: usize) -> Traversal {
+        self.traverse_from_offset(false, include_ignored, start)
+    }
+
+    pub fn entries(&self, include_ignored: bool) -> Traversal {
+        self.traverse_from_offset(true, include_ignored, 0)
     }
 
     pub fn paths(&self) -> impl Iterator<Item = &Arc<Path>> {
@@ -1490,12 +1534,22 @@ impl Snapshot {
             .map(|entry| &entry.path)
     }
 
-    pub fn visible_files(&self, start: usize) -> FileIter {
-        FileIter::visible(self, start)
-    }
-
     fn child_entries<'a>(&'a self, path: &'a Path) -> ChildEntriesIter<'a> {
-        ChildEntriesIter::new(path, self)
+        let mut cursor = self.entries_by_path.cursor();
+        cursor.seek(&TraversalTarget::Path { path }, Bias::Right, &());
+        let mut traversal = Traversal {
+            cursor,
+            include_dirs: true,
+            include_ignored: true,
+        };
+        traversal.advance();
+        let mut done = false;
+        if let Some(item) = traversal.cursor.item() {
+            if !item.path.starts_with(path) {
+                done = true;
+            }
+        }
+        ChildEntriesIter { traversal, done }
     }
 
     pub fn root_entry(&self) -> Option<&Entry> {
@@ -1507,12 +1561,7 @@ impl Snapshot {
     }
 
     fn entry_for_path(&self, path: impl AsRef<Path>) -> Option<&Entry> {
-        let mut cursor = self.entries_by_path.cursor::<PathSearch>();
-        if cursor.seek(&PathSearch::Exact(path.as_ref()), Bias::Left, &()) {
-            cursor.item()
-        } else {
-            None
-        }
+        self.traverse_from_path(true, true, path.as_ref()).entry()
     }
 
     fn entry_for_id(&self, id: usize) -> Option<&Entry> {
@@ -1600,9 +1649,19 @@ impl Snapshot {
         let mut new_entries;
         let removed_entry_ids;
         {
-            let mut cursor = self.entries_by_path.cursor::<PathSearch>();
-            new_entries = cursor.slice(&PathSearch::Exact(path), Bias::Left, &());
-            removed_entry_ids = cursor.slice(&PathSearch::Successor(path), Bias::Left, &());
+            let mut cursor = self.entries_by_path.cursor::<TraversalProgress>();
+            new_entries = cursor.slice(&TraversalTarget::Path { path }, Bias::Left, &());
+            let count = cursor.start().count;
+            removed_entry_ids = cursor.slice(
+                &TraversalTarget::PathSuccessor {
+                    path,
+                    count,
+                    include_ignored: true,
+                    include_dirs: true,
+                },
+                Bias::Left,
+                &(),
+            );
             new_entries.push_tree(cursor.suffix(&()), &());
         }
         self.entries_by_path = new_entries;
@@ -1890,15 +1949,12 @@ impl sum_tree::Item for Entry {
     type Summary = EntrySummary;
 
     fn summary(&self) -> Self::Summary {
+        let visible_count = if self.is_ignored { 0 } else { 1 };
         let file_count;
         let visible_file_count;
         if self.is_file() {
             file_count = 1;
-            if self.is_ignored {
-                visible_file_count = 0;
-            } else {
-                visible_file_count = 1;
-            }
+            visible_file_count = visible_count;
         } else {
             file_count = 0;
             visible_file_count = 0;
@@ -1906,6 +1962,8 @@ impl sum_tree::Item for Entry {
 
         EntrySummary {
             max_path: self.path.clone(),
+            count: 1,
+            visible_count,
             file_count,
             visible_file_count,
         }
@@ -1923,6 +1981,8 @@ impl sum_tree::KeyedItem for Entry {
 #[derive(Clone, Debug)]
 pub struct EntrySummary {
     max_path: Arc<Path>,
+    count: usize,
+    visible_count: usize,
     file_count: usize,
     visible_file_count: usize,
 }
@@ -1931,6 +1991,8 @@ impl Default for EntrySummary {
     fn default() -> Self {
         Self {
             max_path: Arc::from(Path::new("")),
+            count: 0,
+            visible_count: 0,
             file_count: 0,
             visible_file_count: 0,
         }
@@ -2005,64 +2067,6 @@ impl<'a> sum_tree::Dimension<'a, EntrySummary> for PathKey {
     }
 }
 
-#[derive(Copy, Clone, Debug, PartialEq, Eq)]
-enum PathSearch<'a> {
-    Exact(&'a Path),
-    Successor(&'a Path),
-}
-
-impl<'a> Ord for PathSearch<'a> {
-    fn cmp(&self, other: &Self) -> cmp::Ordering {
-        match (self, other) {
-            (Self::Exact(a), Self::Exact(b)) => a.cmp(b),
-            (Self::Successor(a), Self::Exact(b)) => {
-                if b.starts_with(a) {
-                    cmp::Ordering::Greater
-                } else {
-                    a.cmp(b)
-                }
-            }
-            _ => unreachable!("not sure we need the other two cases"),
-        }
-    }
-}
-
-impl<'a> PartialOrd for PathSearch<'a> {
-    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl<'a> Default for PathSearch<'a> {
-    fn default() -> Self {
-        Self::Exact(Path::new("").into())
-    }
-}
-
-impl<'a: 'b, 'b> sum_tree::Dimension<'a, EntrySummary> for PathSearch<'b> {
-    fn add_summary(&mut self, summary: &'a EntrySummary, _: &()) {
-        *self = Self::Exact(summary.max_path.as_ref());
-    }
-}
-
-#[derive(Copy, Clone, Default, Debug, Eq, PartialEq, Ord, PartialOrd)]
-pub struct FileCount(usize);
-
-impl<'a> sum_tree::Dimension<'a, EntrySummary> for FileCount {
-    fn add_summary(&mut self, summary: &'a EntrySummary, _: &()) {
-        self.0 += summary.file_count;
-    }
-}
-
-#[derive(Copy, Clone, Default, Debug, Eq, PartialEq, Ord, PartialOrd)]
-pub struct VisibleFileCount(usize);
-
-impl<'a> sum_tree::Dimension<'a, EntrySummary> for VisibleFileCount {
-    fn add_summary(&mut self, summary: &'a EntrySummary, _: &()) {
-        self.0 += summary.visible_file_count;
-    }
-}
-
 struct BackgroundScanner {
     fs: Arc<dyn Fs>,
     snapshot: Arc<Mutex<Snapshot>>,
@@ -2555,86 +2559,175 @@ impl WorktreeHandle for ModelHandle<Worktree> {
     }
 }
 
-pub enum FileIter<'a> {
-    All(Cursor<'a, Entry, FileCount>),
-    Visible(Cursor<'a, Entry, VisibleFileCount>),
+#[derive(Clone, Debug)]
+struct TraversalProgress<'a> {
+    max_path: &'a Path,
+    count: usize,
+    visible_count: usize,
+    file_count: usize,
+    visible_file_count: usize,
 }
 
-impl<'a> FileIter<'a> {
-    fn all(snapshot: &'a Snapshot, start: usize) -> Self {
-        let mut cursor = snapshot.entries_by_path.cursor();
-        cursor.seek(&FileCount(start), Bias::Right, &());
-        Self::All(cursor)
-    }
-
-    fn visible(snapshot: &'a Snapshot, start: usize) -> Self {
-        let mut cursor = snapshot.entries_by_path.cursor();
-        cursor.seek(&VisibleFileCount(start), Bias::Right, &());
-        Self::Visible(cursor)
-    }
-
-    fn next_internal(&mut self) {
-        match self {
-            Self::All(cursor) => {
-                let ix = cursor.start().0;
-                cursor.seek_forward(&FileCount(ix + 1), Bias::Right, &());
-            }
-            Self::Visible(cursor) => {
-                let ix = cursor.start().0;
-                cursor.seek_forward(&VisibleFileCount(ix + 1), Bias::Right, &());
-            }
-        }
-    }
-
-    fn item(&self) -> Option<&'a Entry> {
-        match self {
-            Self::All(cursor) => cursor.item(),
-            Self::Visible(cursor) => cursor.item(),
+impl<'a> TraversalProgress<'a> {
+    fn count(&self, include_dirs: bool, include_ignored: bool) -> usize {
+        match (include_ignored, include_dirs) {
+            (true, true) => self.count,
+            (true, false) => self.file_count,
+            (false, true) => self.visible_count,
+            (false, false) => self.visible_file_count,
         }
     }
 }
 
-impl<'a> Iterator for FileIter<'a> {
+impl<'a> sum_tree::Dimension<'a, EntrySummary> for TraversalProgress<'a> {
+    fn add_summary(&mut self, summary: &'a EntrySummary, _: &()) {
+        self.max_path = summary.max_path.as_ref();
+        self.count += summary.count;
+        self.visible_count += summary.visible_count;
+        self.file_count += summary.file_count;
+        self.visible_file_count += summary.visible_file_count;
+    }
+}
+
+impl<'a> Default for TraversalProgress<'a> {
+    fn default() -> Self {
+        Self {
+            max_path: Path::new(""),
+            count: 0,
+            visible_count: 0,
+            file_count: 0,
+            visible_file_count: 0,
+        }
+    }
+}
+
+pub struct Traversal<'a> {
+    cursor: sum_tree::Cursor<'a, Entry, TraversalProgress<'a>>,
+    include_ignored: bool,
+    include_dirs: bool,
+}
+
+impl<'a> Traversal<'a> {
+    pub fn advance(&mut self) -> bool {
+        self.cursor.seek_forward(
+            &TraversalTarget::Count {
+                count: self
+                    .cursor
+                    .start()
+                    .count(self.include_dirs, self.include_ignored)
+                    + 1,
+                include_dirs: self.include_dirs,
+                include_ignored: self.include_ignored,
+            },
+            Bias::Right,
+            &(),
+        )
+    }
+
+    pub fn advance_sibling(&mut self) -> bool {
+        let count = self
+            .cursor
+            .start()
+            .count(self.include_dirs, self.include_ignored);
+        let prev_path = self.cursor.start().max_path.as_ref();
+        self.cursor.seek_forward(
+            &TraversalTarget::PathSuccessor {
+                count,
+                path: prev_path,
+                include_dirs: self.include_dirs,
+                include_ignored: self.include_ignored,
+            },
+            Bias::Right,
+            &(),
+        );
+        if let Some(item) = self.entry() {
+            item.path.components().count() == prev_path.components().count()
+        } else {
+            false
+        }
+    }
+
+    pub fn entry(&self) -> Option<&'a Entry> {
+        self.cursor.item()
+    }
+}
+
+impl<'a> Iterator for Traversal<'a> {
     type Item = &'a Entry;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if let Some(entry) = self.item() {
-            self.next_internal();
-            Some(entry)
+        if let Some(item) = self.entry() {
+            self.advance();
+            Some(item)
         } else {
             None
         }
     }
 }
 
-struct ChildEntriesIter<'a> {
-    parent_path: &'a Path,
-    cursor: Cursor<'a, Entry, PathSearch<'a>>,
+#[derive(Debug)]
+enum TraversalTarget<'a> {
+    Path {
+        path: &'a Path,
+    },
+    PathSuccessor {
+        path: &'a Path,
+        count: usize,
+        include_ignored: bool,
+        include_dirs: bool,
+    },
+    Count {
+        count: usize,
+        include_ignored: bool,
+        include_dirs: bool,
+    },
 }
 
-impl<'a> ChildEntriesIter<'a> {
-    fn new(parent_path: &'a Path, snapshot: &'a Snapshot) -> Self {
-        let mut cursor = snapshot.entries_by_path.cursor();
-        cursor.seek(&PathSearch::Exact(parent_path), Bias::Right, &());
-        Self {
-            parent_path,
-            cursor,
+impl<'a, 'b> SeekTarget<'a, EntrySummary, TraversalProgress<'a>> for TraversalTarget<'b> {
+    fn cmp(&self, cursor_location: &TraversalProgress<'a>, _: &()) -> Ordering {
+        match self {
+            TraversalTarget::Path { path } => path.cmp(&cursor_location.max_path),
+            TraversalTarget::PathSuccessor {
+                path,
+                count,
+                include_dirs,
+                include_ignored,
+            } => {
+                if !cursor_location.max_path.starts_with(path)
+                    && cursor_location.count(*include_dirs, *include_ignored) > *count
+                {
+                    Ordering::Equal
+                } else {
+                    Ordering::Greater
+                }
+            }
+            TraversalTarget::Count {
+                count,
+                include_dirs,
+                include_ignored,
+            } => Ord::cmp(
+                count,
+                &cursor_location.count(*include_dirs, *include_ignored),
+            ),
         }
     }
 }
 
+struct ChildEntriesIter<'a> {
+    traversal: Traversal<'a>,
+    done: bool,
+}
+
 impl<'a> Iterator for ChildEntriesIter<'a> {
     type Item = &'a Entry;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if let Some(item) = self.cursor.item() {
-            if item.path.starts_with(self.parent_path) {
-                self.cursor
-                    .seek_forward(&PathSearch::Successor(&item.path), Bias::Left, &());
-                Some(item)
-            } else {
-                None
-            }
+        if self.done {
+            return None;
+        }
+        if let Some(item) = self.traversal.entry() {
+            self.done = !self.traversal.advance_sibling();
+            Some(item)
         } else {
             None
         }
@@ -3417,8 +3510,8 @@ mod tests {
 
     impl Snapshot {
         fn check_invariants(&self) {
-            let mut files = self.files(0);
-            let mut visible_files = self.visible_files(0);
+            let mut files = self.files(true, 0);
+            let mut visible_files = self.files(false, 0);
             for entry in self.entries_by_path.cursor::<()>() {
                 if entry.is_file() {
                     assert_eq!(files.next().unwrap().inode, entry.inode);