Start on randomized test and add SuggestionMapSnapshot::chunks

crates/editor/src/display_map/suggestion_map.rs

@@ -1,11 +1,15 @@
-use std::ops::{Add, AddAssign, Sub};
-
-use crate::{ToOffset, ToPoint};
-
-use super::fold_map::{FoldEdit, FoldOffset, FoldSnapshot};
+use super::{
+    fold_map::{FoldChunks, FoldEdit, FoldOffset, FoldSnapshot},
+    TextHighlights,
+};
+use crate::ToPoint;
 use gpui::fonts::HighlightStyle;
-use language::{Bias, Edit, Patch, Rope};
+use language::{Bias, Chunk, Edit, Patch, Point, Rope};
 use parking_lot::Mutex;
+use std::{
+    cmp,
+    ops::{Add, AddAssign, Range, Sub},
+};
 
 pub type SuggestionEdit = Edit<SuggestionOffset>;
 
@@ -34,15 +38,26 @@ impl AddAssign for SuggestionOffset {
     }
 }
 
+#[derive(Copy, Clone, Debug, Default, Eq, Ord, PartialOrd, PartialEq)]
+pub struct SuggestionPoint(pub Point);
+
 #[derive(Clone)]
 pub struct Suggestion<T> {
     position: T,
     text: Rope,
+    highlight_style: HighlightStyle,
 }
 
 pub struct SuggestionMap(Mutex<SuggestionSnapshot>);
 
 impl SuggestionMap {
+    pub fn new(fold_snapshot: FoldSnapshot) -> Self {
+        Self(Mutex::new(SuggestionSnapshot {
+            fold_snapshot,
+            suggestion: None,
+        }))
+    }
+
     pub fn replace<T>(
         &mut self,
         new_suggestion: Option<Suggestion<T>>,
@@ -61,6 +76,7 @@ impl SuggestionMap {
             Suggestion {
                 position: fold_offset,
                 text: new_suggestion.text,
+                highlight_style: new_suggestion.highlight_style,
             }
         });
 
@@ -128,7 +144,7 @@ impl SuggestionMap {
                     ..SuggestionOffset(fold_edit.new.end.0 + suggestion_new_len),
             });
         }
-        snapshot.folds_snapshot = fold_snapshot;
+        snapshot.fold_snapshot = fold_snapshot;
 
         (snapshot.clone(), suggestion_edits)
     }
@@ -136,6 +152,202 @@ impl SuggestionMap {
 
 #[derive(Clone)]
 pub struct SuggestionSnapshot {
-    folds_snapshot: FoldSnapshot,
+    fold_snapshot: FoldSnapshot,
     suggestion: Option<Suggestion<FoldOffset>>,
 }
+
+impl SuggestionSnapshot {
+    pub fn max_point(&self) -> SuggestionPoint {
+        if let Some(suggestion) = self.suggestion.as_ref() {
+            let suggestion_point = suggestion.position.to_point(&self.fold_snapshot);
+            let mut max_point = suggestion_point.0;
+            max_point += suggestion.text.max_point();
+            max_point += self.fold_snapshot.max_point().0 - suggestion_point.0;
+            SuggestionPoint(max_point)
+        } else {
+            SuggestionPoint(self.fold_snapshot.max_point().0)
+        }
+    }
+
+    pub fn len(&self) -> SuggestionOffset {
+        if let Some(suggestion) = self.suggestion.as_ref() {
+            let mut len = suggestion.position.0;
+            len += suggestion.text.len();
+            len += self.fold_snapshot.len().0 - suggestion.position.0;
+            SuggestionOffset(len)
+        } else {
+            SuggestionOffset(self.fold_snapshot.len().0)
+        }
+    }
+
+    pub fn chunks<'a>(
+        &'a self,
+        range: Range<SuggestionOffset>,
+        language_aware: bool,
+        text_highlights: Option<&'a TextHighlights>,
+    ) -> Chunks<'a> {
+        if let Some(suggestion) = self.suggestion.as_ref() {
+            let suggestion_range =
+                suggestion.position.0..suggestion.position.0 + suggestion.text.len();
+
+            let prefix_chunks = if range.start.0 < suggestion_range.start {
+                Some(self.fold_snapshot.chunks(
+                    FoldOffset(range.start.0)
+                        ..cmp::min(FoldOffset(suggestion_range.start), FoldOffset(range.end.0)),
+                    language_aware,
+                    text_highlights,
+                ))
+            } else {
+                None
+            };
+
+            let clipped_suggestion_range = cmp::max(range.start.0, suggestion_range.start)
+                ..cmp::min(range.end.0, suggestion_range.end);
+            let suggestion_chunks = if clipped_suggestion_range.start < clipped_suggestion_range.end
+            {
+                let start = clipped_suggestion_range.start - suggestion_range.start;
+                let end = clipped_suggestion_range.end - suggestion_range.start;
+                Some(suggestion.text.chunks_in_range(start..end))
+            } else {
+                None
+            };
+
+            let suffix_chunks = if range.end.0 > suggestion_range.end {
+                let start = cmp::max(suggestion_range.end, range.start.0) - suggestion_range.len();
+                let end = range.end.0 - suggestion_range.len();
+                Some(self.fold_snapshot.chunks(
+                    FoldOffset(start)..FoldOffset(end),
+                    language_aware,
+                    text_highlights,
+                ))
+            } else {
+                None
+            };
+
+            Chunks {
+                prefix_chunks,
+                suggestion_chunks,
+                suffix_chunks,
+                highlight_style: suggestion.highlight_style,
+            }
+        } else {
+            Chunks {
+                prefix_chunks: Some(self.fold_snapshot.chunks(
+                    FoldOffset(range.start.0)..FoldOffset(range.end.0),
+                    language_aware,
+                    text_highlights,
+                )),
+                suggestion_chunks: None,
+                suffix_chunks: None,
+                highlight_style: Default::default(),
+            }
+        }
+    }
+
+    #[cfg(test)]
+    pub fn text(&self) -> String {
+        self.chunks(Default::default()..self.len(), false, None)
+            .map(|chunk| chunk.text)
+            .collect()
+    }
+}
+
+pub struct Chunks<'a> {
+    prefix_chunks: Option<FoldChunks<'a>>,
+    suggestion_chunks: Option<text::Chunks<'a>>,
+    suffix_chunks: Option<FoldChunks<'a>>,
+    highlight_style: HighlightStyle,
+}
+
+impl<'a> Iterator for Chunks<'a> {
+    type Item = Chunk<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(chunks) = self.prefix_chunks.as_mut() {
+            if let Some(chunk) = chunks.next() {
+                return Some(chunk);
+            } else {
+                self.prefix_chunks = None;
+            }
+        }
+
+        if let Some(chunks) = self.suggestion_chunks.as_mut() {
+            if let Some(chunk) = chunks.next() {
+                return Some(Chunk {
+                    text: chunk,
+                    syntax_highlight_id: None,
+                    highlight_style: Some(self.highlight_style),
+                    diagnostic_severity: None,
+                    is_unnecessary: false,
+                });
+            } else {
+                self.suggestion_chunks = None;
+            }
+        }
+
+        if let Some(chunks) = self.suffix_chunks.as_mut() {
+            if let Some(chunk) = chunks.next() {
+                return Some(chunk);
+            } else {
+                self.suffix_chunks = None;
+            }
+        }
+
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{display_map::fold_map::FoldMap, MultiBuffer};
+    use gpui::MutableAppContext;
+    use rand::{prelude::StdRng, Rng};
+
+    #[gpui::test(iterations = 100)]
+    fn test_random_suggestions(cx: &mut MutableAppContext, mut rng: StdRng) {
+        cx.set_global(Settings::test(cx));
+        let operations = env::var("OPERATIONS")
+            .map(|i| i.parse().expect("invalid `OPERATIONS` variable"))
+            .unwrap_or(10);
+
+        let len = rng.gen_range(0..30);
+        let buffer = if rng.gen() {
+            let text = util::RandomCharIter::new(&mut rng)
+                .take(len)
+                .collect::<String>();
+            MultiBuffer::build_simple(&text, cx)
+        } else {
+            MultiBuffer::build_random(&mut rng, cx)
+        };
+        let buffer_snapshot = buffer.read(cx).snapshot(cx);
+        log::info!("Buffer text: {:?}", buffer_snapshot.text());
+
+        let (mut fold_map, _) = FoldMap::new(buffer_snapshot.clone());
+        let (fold_snapshot, _) = fold_map.read(buffer_snapshot, vec![]);
+        let suggestion_map = SuggestionMap::new(fold_snapshot.clone());
+
+        for _ in 0..operations {
+            let mut buffer_edits = Vec::new();
+            match rng.gen_range(0..=100) {
+                0..=59 => {
+                    for (fold_snapshot, fold_edits) in fold_map.randomly_mutate(&mut rng) {
+                        suggestion_map.sync(fold_snapshot, fold_edits);
+                    }
+                }
+                _ => buffer.update(cx, |buffer, cx| {
+                    let subscription = buffer.subscribe();
+                    let edit_count = rng.gen_range(1..=5);
+                    buffer.randomly_mutate(&mut rng, edit_count, cx);
+                    buffer_snapshot = buffer.snapshot(cx);
+                    let edits = subscription.consume().into_inner();
+                    log::info!("editing {:?}", edits);
+                    buffer_edits.extend(edits);
+                }),
+            };
+
+            log::info!("buffer text: {:?}", buffer_snapshot.text());
+            log::info!("folds text: {:?}", buffer_snapshot.text());
+        }
+    }
+}