Add FoldMap::folds_in_range to randomized test and fix issues it found

zed/src/editor/display_map/fold_map.rs

@@ -3,7 +3,7 @@ use super::{
     Anchor, Buffer, DisplayPoint, Edit, Point, TextSummary, ToOffset,
 };
 use crate::{
-    sum_tree::{self, Cursor, SeekBias, SumTree},
+    sum_tree::{self, Cursor, FilterCursor, SeekBias, SumTree},
     time,
 };
 use anyhow::{anyhow, Result};
@@ -80,15 +80,7 @@ impl FoldMap {
     where
         T: ToOffset,
     {
-        let buffer = self.buffer.read(ctx);
+        Ok(self.intersecting_folds(range, ctx)?.map(|f| &f.0))
-        let range = buffer.anchor_before(range.start)?..buffer.anchor_before(range.end)?;
-        Ok(self
-            .folds
-            .filter::<_, usize>(move |summary| {
-                range.start.cmp(&summary.max_end, buffer).unwrap() < Ordering::Equal
-                    && range.end.cmp(&summary.min_start, buffer).unwrap() >= Ordering::Equal
-            })
-            .map(|f| &f.0))
     }
 
     pub fn fold<T: ToOffset>(
@@ -150,25 +142,17 @@ impl FoldMap {
         let mut edits = Vec::new();
         let mut fold_ixs_to_delete = Vec::new();
         for range in ranges.into_iter() {
-            let start = buffer.anchor_before(range.start.to_offset(buffer)?)?;
-            let end = buffer.anchor_after(range.end.to_offset(buffer)?)?;
-            let range = start..end;
-
             // Remove intersecting folds and add their ranges to edits that are passed to apply_edits
-            let mut cursor = self.folds.filter::<_, usize>(|summary| {
+            let mut folds_cursor = self.intersecting_folds(range, ctx)?;
-                range.start.cmp(&summary.max_end, buffer).unwrap() < Ordering::Equal
+            while let Some(fold) = folds_cursor.item() {
-                    && range.end.cmp(&summary.min_start, buffer).unwrap() >= Ordering::Equal
-            });
-
-            while let Some(fold) = cursor.item() {
                 let offset_range =
                     fold.0.start.to_offset(buffer).unwrap()..fold.0.end.to_offset(buffer).unwrap();
                 edits.push(Edit {
                     old_range: offset_range.clone(),
                     new_range: offset_range,
                 });
-                fold_ixs_to_delete.push(*cursor.start());
+                fold_ixs_to_delete.push(*folds_cursor.start());
-                cursor.next();
+                folds_cursor.next();
             }
         }
         fold_ixs_to_delete.sort_unstable();
@@ -192,6 +176,23 @@ impl FoldMap {
         Ok(())
     }
 
+    fn intersecting_folds<'a, T>(
+        &self,
+        range: Range<T>,
+        ctx: &'a AppContext,
+    ) -> Result<FilterCursor<impl 'a + Fn(&FoldSummary) -> bool, Fold, usize>>
+    where
+        T: ToOffset,
+    {
+        let buffer = self.buffer.read(ctx);
+        let start = buffer.anchor_before(range.start.to_offset(buffer)?)?;
+        let end = buffer.anchor_after(range.end.to_offset(buffer)?)?;
+        Ok(self.folds.filter::<_, usize>(move |summary| {
+            start.cmp(&summary.max_end, buffer).unwrap() <= Ordering::Equal
+                && end.cmp(&summary.min_start, buffer).unwrap() >= Ordering::Equal
+        }))
+    }
+
     pub fn is_line_folded(&self, display_row: u32, ctx: &AppContext) -> bool {
         let transforms = self.sync(ctx);
         let mut cursor = transforms.cursor::<DisplayPoint, DisplayPoint>();
@@ -502,7 +503,7 @@ impl Default for FoldSummary {
         Self {
             start: Anchor::Start,
             end: Anchor::End,
-            min_start: Anchor::Start,
+            min_start: Anchor::End,
             max_end: Anchor::Start,
             count: 0,
         }
@@ -521,12 +522,12 @@ impl sum_tree::Summary for FoldSummary {
             self.max_end = other.max_end.clone();
         }
 
-        if cfg!(debug_assertions) {
+        #[cfg(debug_assertions)]
+        {
             let start_comparison = self.start.cmp(&other.start, buffer).unwrap();
-            let end_comparison = self.end.cmp(&other.end, buffer).unwrap();
             assert!(start_comparison <= Ordering::Equal);
             if start_comparison == Ordering::Equal {
-                assert!(end_comparison >= Ordering::Equal);
+                assert!(self.end.cmp(&other.end, buffer).unwrap() >= Ordering::Equal);
             }
         }
         self.start = other.start.clone();
@@ -821,6 +822,7 @@ mod tests {
                 fold_ranges,
                 vec![
                     Point::new(0, 2)..Point::new(2, 2),
+                    Point::new(0, 4)..Point::new(1, 0),
                     Point::new(1, 2)..Point::new(3, 2)
                 ]
             );
@@ -848,7 +850,7 @@ mod tests {
         };
 
         for seed in seed_range {
-            println!("{:?}", seed);
+            dbg!(seed);
             let mut rng = StdRng::seed_from_u64(seed);
 
             App::test((), |app| {
@@ -973,6 +975,30 @@ mod tests {
                         );
                         assert!(map.is_line_folded(display_point.row(), app.as_ref()));
                     }
+
+                    for _ in 0..5 {
+                        let end = rng.gen_range(0..=buffer.len());
+                        let start = rng.gen_range(0..=end);
+                        let expected_folds = map
+                            .folds
+                            .items()
+                            .into_iter()
+                            .filter(|fold| {
+                                let fold_start = fold.0.start.to_offset(buffer).unwrap();
+                                let fold_end = fold.0.end.to_offset(buffer).unwrap();
+                                start <= fold_end && end >= fold_start
+                            })
+                            .map(|fold| fold.0)
+                            .collect::<Vec<_>>();
+
+                        assert_eq!(
+                            map.folds_in_range(start..end, app.as_ref())
+                                .unwrap()
+                                .cloned()
+                                .collect::<Vec<_>>(),
+                            expected_folds
+                        );
+                    }
                 }
             });
         }

zed/src/sum_tree/cursor.rs

@@ -199,9 +199,6 @@ where
     }
 
     pub fn next(&mut self) {
-        if !self.did_seek {
-            self.descend_to_first_item(self.tree, |_| true)
-        }
         self.next_internal(|_| true)
     }
 
@@ -209,67 +206,88 @@ where
     where
         F: Fn(&T::Summary) -> bool,
     {
-        assert!(self.did_seek, "Must seek before calling this method");
+        let mut descend = false;
 
-        if self.stack.is_empty() {
+        if !self.did_seek {
-            if !self.at_end {
+            self.stack.push(StackEntry {
-                self.descend_to_first_item(self.tree, filter_node);
+                tree: self.tree,
-            }
+                index: 0,
-        } else {
+                seek_dimension: S::default(),
-            while self.stack.len() > 0 {
+                sum_dimension: U::default(),
-                let new_subtree = {
+            });
-                    let entry = self.stack.last_mut().unwrap();
+            descend = true;
-                    match entry.tree.0.as_ref() {
+            self.did_seek = true;
-                        Node::Internal {
+        }
-                            child_trees,
-                            child_summaries,
-                            ..
-                        } => {
-                            while entry.index < child_summaries.len() {
-                                entry
-                                    .seek_dimension
-                                    .add_summary(&child_summaries[entry.index]);
-                                entry
-                                    .sum_dimension
-                                    .add_summary(&child_summaries[entry.index]);
 
-                                entry.index += 1;
+        while self.stack.len() > 0 {
-                                if let Some(next_summary) = child_summaries.get(entry.index) {
+            let new_subtree = {
-                                    if filter_node(next_summary) {
+                let entry = self.stack.last_mut().unwrap();
-                                        break;
+                match entry.tree.0.as_ref() {
-                                    } else {
+                    Node::Internal {
-                                        self.seek_dimension.add_summary(next_summary);
+                        child_trees,
-                                        self.sum_dimension.add_summary(next_summary);
+                        child_summaries,
-                                    }
+                        ..
-                                }
+                    } => {
-                            }
+                        if !descend {
-
+                            let summary = &child_summaries[entry.index];
-                            child_trees.get(entry.index)
+                            entry.seek_dimension.add_summary(summary);
+                            entry.sum_dimension.add_summary(summary);
+                            entry.index += 1;
                         }
-                        Node::Leaf { item_summaries, .. } => loop {
+
+                        while entry.index < child_summaries.len() {
+                            let next_summary = &child_summaries[entry.index];
+                            if filter_node(next_summary) {
+                                break;
+                            } else {
+                                self.seek_dimension.add_summary(next_summary);
+                                self.sum_dimension.add_summary(next_summary);
+                            }
+                            entry.index += 1;
+                        }
+
+                        child_trees.get(entry.index)
+                    }
+                    Node::Leaf { item_summaries, .. } => {
+                        if !descend {
                             let item_summary = &item_summaries[entry.index];
                             self.seek_dimension.add_summary(item_summary);
                             entry.seek_dimension.add_summary(item_summary);
                             self.sum_dimension.add_summary(item_summary);
                             entry.sum_dimension.add_summary(item_summary);
                             entry.index += 1;
+                        }
+
+                        loop {
                             if let Some(next_item_summary) = item_summaries.get(entry.index) {
                                 if filter_node(next_item_summary) {
                                     return;
+                                } else {
+                                    self.seek_dimension.add_summary(next_item_summary);
+                                    entry.seek_dimension.add_summary(next_item_summary);
+                                    self.sum_dimension.add_summary(next_item_summary);
+                                    entry.sum_dimension.add_summary(next_item_summary);
+                                    entry.index += 1;
                                 }
                             } else {
                                 break None;
                             }
-                        },
+                        }
                     }
-                };
-
-                if let Some(subtree) = new_subtree {
-                    self.descend_to_first_item(subtree, filter_node);
-                    break;
-                } else {
-                    self.stack.pop();
                 }
+            };
+
+            if let Some(subtree) = new_subtree {
+                descend = true;
+                self.stack.push(StackEntry {
+                    tree: subtree,
+                    index: 0,
+                    seek_dimension: self.seek_dimension.clone(),
+                    sum_dimension: self.sum_dimension.clone(),
+                });
+            } else {
+                descend = false;
+                self.stack.pop();
             }
         }
 
@@ -277,67 +295,6 @@ where
         debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
     }
 
-    pub fn descend_to_first_item<F>(&mut self, mut subtree: &'a SumTree<T>, filter_node: F)
-    where
-        F: Fn(&T::Summary) -> bool,
-    {
-        self.did_seek = true;
-        loop {
-            subtree = match *subtree.0 {
-                Node::Internal {
-                    ref child_trees,
-                    ref child_summaries,
-                    ..
-                } => {
-                    let mut new_index = None;
-                    for (index, summary) in child_summaries.iter().enumerate() {
-                        if filter_node(summary) {
-                            new_index = Some(index);
-                            break;
-                        }
-                        self.seek_dimension.add_summary(summary);
-                        self.sum_dimension.add_summary(summary);
-                    }
-
-                    if let Some(new_index) = new_index {
-                        self.stack.push(StackEntry {
-                            tree: subtree,
-                            index: new_index,
-                            seek_dimension: self.seek_dimension.clone(),
-                            sum_dimension: self.sum_dimension.clone(),
-                        });
-                        &child_trees[new_index]
-                    } else {
-                        break;
-                    }
-                }
-                Node::Leaf {
-                    ref item_summaries, ..
-                } => {
-                    let mut new_index = None;
-                    for (index, item_summary) in item_summaries.iter().enumerate() {
-                        if filter_node(item_summary) {
-                            new_index = Some(index);
-                            break;
-                        }
-                        self.seek_dimension.add_summary(item_summary);
-                        self.sum_dimension.add_summary(item_summary);
-                    }
-
-                    if let Some(new_index) = new_index {
-                        self.stack.push(StackEntry {
-                            tree: subtree,
-                            index: new_index,
-                            seek_dimension: self.seek_dimension.clone(),
-                            sum_dimension: self.sum_dimension.clone(),
-                        });
-                    }
-                    break;
-                }
-            }
-        }
-    }
-
     fn descend_to_last_item(&mut self, mut subtree: &'a SumTree<T>) {
         self.did_seek = true;
         loop {
@@ -744,7 +701,7 @@ where
 
     fn next(&mut self) -> Option<Self::Item> {
         if !self.did_seek {
-            self.descend_to_first_item(self.tree, |_| true);
+            self.next();
         }
 
         if let Some(item) = self.item() {
@@ -769,13 +726,7 @@ where
 {
     pub fn new(tree: &'a SumTree<T>, filter_node: F) -> Self {
         let mut cursor = tree.cursor::<(), U>();
-        if filter_node(&tree.summary()) {
+        cursor.next_internal(&filter_node);
-            cursor.descend_to_first_item(tree, &filter_node);
-        } else {
-            cursor.did_seek = true;
-            cursor.at_end = true;
-        }
-
         Self {
             cursor,
             filter_node,

zed/src/sum_tree/mod.rs

@@ -73,7 +73,6 @@ impl<T: Item> SumTree<T> {
     #[allow(unused)]
     pub fn items(&self) -> Vec<T> {
         let mut cursor = self.cursor::<(), ()>();
-        cursor.descend_to_first_item(self, |_| true);
         cursor.cloned().collect()
     }
 
@@ -566,6 +565,7 @@ mod tests {
         for seed in 0..100 {
             use rand::{distributions, prelude::*};
 
+            dbg!(seed);
             let rng = &mut StdRng::seed_from_u64(seed);
 
             let mut tree = SumTree::<u8>::new();