Replace linear scan of entries with a custom FileIter

zed/src/workspace/workspace.rs

@@ -224,7 +224,7 @@ impl WorkspaceHandle for ModelHandle<Workspace> {
             .iter()
             .flat_map(|tree| {
                 let tree_id = tree.id();
-                tree.read(app).files().map(move |inode| (tree_id, inode))
+                tree.read(app).files(0).map(move |f| (tree_id, f.inode()))
             })
             .collect::<Vec<_>>()
     }
@@ -253,7 +253,7 @@ mod tests {
 
             // Get the first file entry.
             let tree = app.read(|ctx| workspace.read(ctx).worktrees.iter().next().unwrap().clone());
-            let file_inode = app.read(|ctx| tree.read(ctx).files().next().unwrap());
+            let file_inode = app.read(|ctx| tree.read(ctx).files(0).next().unwrap().inode());
             let entry = (tree.id(), file_inode);
 
             // Open the same entry twice before it finishes loading.

zed/src/worktree.rs

@@ -3,7 +3,7 @@ mod fuzzy;
 
 use crate::{
     editor::{History, Snapshot as BufferSnapshot},
-    sum_tree::{self, Edit, SumTree},
+    sum_tree::{self, Cursor, Edit, SeekBias, SumTree},
 };
 use anyhow::{anyhow, Context, Result};
 use fuzzy::PathEntry;
@@ -153,10 +153,6 @@ impl Worktree {
         self.snapshot.entries.get(&inode).is_some()
     }
 
-    pub fn file_count(&self) -> usize {
-        self.snapshot.entries.summary().file_count
-    }
-
     pub fn abs_path_for_inode(&self, ino: u64) -> Result<PathBuf> {
         let mut result = self.snapshot.path.to_path_buf();
         result.push(self.path_for_inode(ino, false)?);
@@ -195,20 +191,6 @@ impl Worktree {
             Ok(())
         })
     }
-
-    #[cfg(test)]
-    pub fn files<'a>(&'a self) -> impl Iterator<Item = u64> + 'a {
-        self.snapshot
-            .entries
-            .cursor::<(), ()>()
-            .filter_map(|entry| {
-                if let Entry::File { inode, .. } = entry {
-                    Some(*inode)
-                } else {
-                    None
-                }
-            })
-    }
 }
 
 impl Entity for Worktree {
@@ -248,6 +230,14 @@ impl Snapshot {
         self.entries.summary().visible_file_count
     }
 
+    pub fn files(&self, start: usize) -> FileIter {
+        FileIter::all(self, start)
+    }
+
+    pub fn visible_files(&self, start: usize) -> FileIter {
+        FileIter::visible(self, start)
+    }
+
     pub fn root_entry(&self) -> Option<&Entry> {
         self.root_inode.and_then(|inode| self.entries.get(&inode))
     }
@@ -614,7 +604,7 @@ impl Entry {
         }
     }
 
-    fn inode(&self) -> u64 {
+    pub fn inode(&self) -> u64 {
         match self {
             Entry::Dir { inode, .. } => *inode,
             Entry::File { inode, .. } => *inode,
@@ -692,7 +682,7 @@ impl<'a> sum_tree::Dimension<'a, EntrySummary> for u64 {
 }
 
 #[derive(Copy, Clone, Default, Debug, Eq, PartialEq, Ord, PartialOrd)]
-struct FileCount(usize);
+pub struct FileCount(usize);
 
 impl<'a> sum_tree::Dimension<'a, EntrySummary> for FileCount {
     fn add_summary(&mut self, summary: &'a EntrySummary) {
@@ -701,7 +691,7 @@ impl<'a> sum_tree::Dimension<'a, EntrySummary> for FileCount {
 }
 
 #[derive(Copy, Clone, Default, Debug, Eq, PartialEq, Ord, PartialOrd)]
-struct VisibleFileCount(usize);
+pub struct VisibleFileCount(usize);
 
 impl<'a> sum_tree::Dimension<'a, EntrySummary> for VisibleFileCount {
     fn add_summary(&mut self, summary: &'a EntrySummary) {
@@ -1167,6 +1157,58 @@ impl WorktreeHandle for ModelHandle<Worktree> {
     }
 }
 
+pub enum FileIter<'a> {
+    All(Cursor<'a, Entry, FileCount, FileCount>),
+    Visible(Cursor<'a, Entry, VisibleFileCount, VisibleFileCount>),
+}
+
+impl<'a> FileIter<'a> {
+    fn all(snapshot: &'a Snapshot, start: usize) -> Self {
+        let mut cursor = snapshot.entries.cursor();
+        cursor.seek(&FileCount(start), SeekBias::Right);
+        Self::All(cursor)
+    }
+
+    fn visible(snapshot: &'a Snapshot, start: usize) -> Self {
+        let mut cursor = snapshot.entries.cursor();
+        cursor.seek(&VisibleFileCount(start), SeekBias::Right);
+        Self::Visible(cursor)
+    }
+
+    fn next_internal(&mut self) {
+        match self {
+            Self::All(cursor) => {
+                let ix = *cursor.start();
+                cursor.seek_forward(&FileCount(ix.0 + 1), SeekBias::Right);
+            }
+            Self::Visible(cursor) => {
+                let ix = *cursor.start();
+                cursor.seek_forward(&VisibleFileCount(ix.0 + 1), SeekBias::Right);
+            }
+        }
+    }
+
+    fn item(&self) -> Option<&'a Entry> {
+        match self {
+            Self::All(cursor) => cursor.item(),
+            Self::Visible(cursor) => cursor.item(),
+        }
+    }
+}
+
+impl<'a> Iterator for FileIter<'a> {
+    type Item = &'a Entry;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(entry) = self.item() {
+            self.next_internal();
+            Some(entry)
+        } else {
+            None
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1248,7 +1290,7 @@ mod tests {
 
             let file_inode = app.read(|ctx| {
                 let tree = tree.read(ctx);
-                let inode = tree.files().next().unwrap();
+                let inode = tree.files(0).next().unwrap().inode();
                 assert_eq!(
                     tree.path_for_inode(inode, false)
                         .unwrap()
@@ -1532,10 +1574,16 @@ mod tests {
 
     impl Snapshot {
         fn check_invariants(&self) {
+            let mut path_entries = self.files(0);
             for entry in self.entries.items() {
                 let path = self.path_for_inode(entry.inode(), false).unwrap();
                 assert_eq!(self.inode_for_path(path).unwrap(), entry.inode());
+
+                if let Entry::File { inode, .. } = entry {
+                    assert_eq!(path_entries.next().unwrap().inode(), inode);
+                }
             }
+            assert!(path_entries.next().is_none());
         }
 
         fn to_vec(&self) -> Vec<(PathBuf, u64)> {

zed/src/worktree/fuzzy.rs

@@ -1,9 +1,5 @@
+use super::{char_bag::CharBag, Entry, Snapshot};
 use gpui::scoped_pool;
-
-use crate::sum_tree::SeekBias;
-
-use super::{char_bag::CharBag, Entry, FileCount, Snapshot, VisibleFileCount};
-
 use std::{
     cmp::{max, min, Ordering, Reverse},
     collections::BinaryHeap,
@@ -122,7 +118,19 @@ where
                     if tree_start < segment_end && segment_start < tree_end {
                         let start = max(tree_start, segment_start) - tree_start;
                         let end = min(tree_end, segment_end) - tree_start;
-                        let path_entries = path_entries_iter(snapshot, start, end, include_ignored);
+                        let entries = if include_ignored {
+                            snapshot.files(start).take(end - start)
+                        } else {
+                            snapshot.visible_files(start).take(end - start)
+                        };
+                        let path_entries = entries.map(|entry| {
+                            if let Entry::File { path, .. } = entry {
+                                path
+                            } else {
+                                unreachable!()
+                            }
+                        });
+
                         let skipped_prefix_len = if include_root_name {
                             0
                         } else if let Some(Entry::Dir { .. }) = snapshot.root_entry() {
@@ -171,44 +179,6 @@ where
     results
 }
 
-fn path_entries_iter<'a>(
-    snapshot: &'a Snapshot,
-    start: usize,
-    end: usize,
-    include_ignored: bool,
-) -> impl Iterator<Item = &'a PathEntry> {
-    let mut files_cursor = None;
-    let mut visible_files_cursor = None;
-    if include_ignored {
-        let mut cursor = snapshot.entries.cursor::<_, ()>();
-        cursor.seek(&FileCount(start), SeekBias::Right);
-        files_cursor = Some(cursor);
-    } else {
-        let mut cursor = snapshot.entries.cursor::<_, ()>();
-        cursor.seek(&VisibleFileCount(start), SeekBias::Right);
-        visible_files_cursor = Some(cursor);
-    }
-    files_cursor
-        .into_iter()
-        .flatten()
-        .chain(visible_files_cursor.into_iter().flatten())
-        .filter_map(move |e| {
-            if let Entry::File {
-                path, is_ignored, ..
-            } = e
-            {
-                if is_ignored.unwrap_or(false) && !include_ignored {
-                    None
-                } else {
-                    Some(path)
-                }
-            } else {
-                None
-            }
-        })
-        .take(end - start)
-}
-
 fn match_single_tree_paths<'a>(
     snapshot: &Snapshot,
     skipped_prefix_len: usize,