Improve line annotation API

The line annotation as implemented in helix-editor#5420 had two shortcomings:
* It required the height of virtual text lines to be known ahead time
* It checked for line anchors at every grapheme

The first problem made the API impractical to use in practice because
almost all virtual text needs to be softwrapped. For example inline
diagnostics should be softwrapped to avoid cutting off the diagnostic
message (as no scrolling is possible). While more complex virtual text
like side by side diffs must dynamically calculate the number of empty
lines two align two documents (which requires taking account both
softwrap and virtual text). To address this, the API has been
refactored to use a trait.

The second issue caused some performance overhead and unnecessarily
complicated the `DocumentFormatter`. It was addressed by only calling
the trait mentioned above at line breaks (instead of always). This
allows offers additional flexibility to annotations as it offers
the flexibility to align lines (needed for side by side diffs).

helix-core/src/doc_formatter.rs

@@ -11,7 +11,7 @@
 
 use std::borrow::Cow;
 use std::fmt::Debug;
-use std::mem::{replace, take};
+use std::mem::replace;
 
 #[cfg(test)]
 mod test;
@@ -166,9 +166,6 @@ pub struct DocumentFormatter<'t> {
     line_pos: usize,
     exhausted: bool,
 
-    /// Line breaks to be reserved for virtual text
-    /// at the next line break
-    virtual_lines: usize,
     inline_anntoation_graphemes: Option<(Graphemes<'t>, Option<Highlight>)>,
 
     // softwrap specific
@@ -209,7 +206,6 @@ impl<'t> DocumentFormatter<'t> {
             graphemes: RopeGraphemes::new(text.slice(block_char_idx..)),
             char_pos: block_char_idx,
             exhausted: false,
-            virtual_lines: 0,
             indent_level: None,
             peeked_grapheme: None,
             word_buf: Vec::with_capacity(64),
@@ -250,7 +246,6 @@ impl<'t> DocumentFormatter<'t> {
             if let Some((grapheme, highlight)) = self.next_inline_annotation_grapheme(char_pos) {
                 (grapheme.into(), GraphemeSource::VirtualText { highlight })
             } else if let Some(grapheme) = self.graphemes.next() {
-                self.virtual_lines += self.annotations.annotation_lines_at(self.char_pos);
                 let codepoints = grapheme.len_chars() as u32;
 
                 let overlay = self.annotations.overlay_at(char_pos);
@@ -293,8 +288,11 @@ impl<'t> DocumentFormatter<'t> {
             0
         };
 
+        let virtual_lines =
+            self.annotations
+                .virtual_lines_at(self.char_pos, self.visual_pos, self.line_pos);
         self.visual_pos.col = indent_carry_over as usize;
-        self.visual_pos.row += 1 + take(&mut self.virtual_lines);
+        self.visual_pos.row += 1 + virtual_lines;
         let mut i = 0;
         let mut word_width = 0;
         let wrap_indicator = UnicodeSegmentation::graphemes(&*self.text_fmt.wrap_indicator, true)
@@ -404,24 +402,32 @@ impl<'t> Iterator for DocumentFormatter<'t> {
             self.advance_grapheme(self.visual_pos.col, self.char_pos)?
         };
 
-        let visual_pos = self.visual_pos;
-        let char_pos = self.char_pos;
+        let grapheme = FormattedGrapheme {
+            raw: grapheme.grapheme,
+            source: grapheme.source,
+            visual_pos: self.visual_pos,
+            line_idx: self.line_pos,
+            char_idx: self.char_pos,
+        };
+
         self.char_pos += grapheme.doc_chars();
-        let line_idx = self.line_pos;
-        if grapheme.grapheme == Grapheme::Newline {
-            self.visual_pos.row += 1;
-            self.visual_pos.row += take(&mut self.virtual_lines);
+        if !grapheme.is_virtual() {
+            self.annotations.process_virtual_text_anchors(&grapheme);
+        }
+        if grapheme.raw == Grapheme::Newline {
+            // move to end of newline char
+            self.visual_pos.col += 1;
+            let virtual_lines =
+                self.annotations
+                    .virtual_lines_at(self.char_pos, self.visual_pos, self.line_pos);
+            self.visual_pos.row += 1 + virtual_lines;
             self.visual_pos.col = 0;
-            self.line_pos += 1;
+            if !grapheme.is_virtual() {
+                self.line_pos += 1;
+            }
         } else {
             self.visual_pos.col += grapheme.width();
         }
-        Some(FormattedGrapheme {
-            raw: grapheme.grapheme,
-            source: grapheme.source,
-            visual_pos,
-            line_idx,
-            char_idx: char_pos,
-        })
+        Some(grapheme)
     }
 }

helix-core/src/text_annotations.rs

@@ -1,8 +1,12 @@
 use std::cell::Cell;
+use std::cmp::Ordering;
+use std::fmt::Debug;
 use std::ops::Range;
+use std::ptr::NonNull;
 
+use crate::doc_formatter::FormattedGrapheme;
 use crate::syntax::Highlight;
-use crate::Tendril;
+use crate::{Position, Tendril};
 
 /// An inline annotation is continuous text shown
 /// on the screen before the grapheme that starts at
@@ -75,19 +79,98 @@ impl Overlay {
     }
 }
 
-/// Line annotations allow for virtual text between normal
-/// text lines. They cause `height` empty lines to be inserted
-/// below the document line that contains `anchor_char_idx`.
+/// Line annotations allow inserting virtual text lines between normal text
+/// lines.  These lines can be filled with text in the rendering code as their
+/// contents have no effect beyond visual appearance.
 ///
-/// These lines can be filled with text in the rendering code
-/// as their contents have no effect beyond visual appearance.
+/// The height of virtual text is usually not known ahead of time as virtual
+/// text often requires softwrapping. Furthermore the height of some virtual
+/// text like side-by-side diffs depends on the height of the text (again
+/// influenced by softwrap) and other virtual text. Therefore line annotations
+/// are computed on the fly instead of ahead of time like other annotations.
 ///
-/// To insert a line after a document line simply set
-/// `anchor_char_idx` to `doc.line_to_char(line_idx)`
-#[derive(Debug, Clone)]
-pub struct LineAnnotation {
-    pub anchor_char_idx: usize,
-    pub height: usize,
+/// The core of this trait `insert_virtual_lines` function. It is called at the
+/// end of every  visual line and allows the `LineAnnotation` to insert empty
+/// virtual lines. Apart from that the `LineAnnotation` trait has multiple
+/// methods that allow it to track anchors in the document.
+///
+/// When a new traversal of a document starts `reset_pos` is called. Afterwards
+/// the other functions are called with indices that are larger then the
+/// one passed to `reset_pos`. This allows performing a binary search (use
+/// `partition_point`) in `reset_pos` once and then to only look at the next
+/// anchor during each method call.
+///
+/// The `reset_pos`, `skip_conceal` and `process_anchor` functions all return a
+/// `char_idx` anchor. This anchor is stored when transversing the document and
+/// when the grapheme at the anchor is traversed the `process_anchor` function
+/// is called.
+///
+/// # Note
+///
+/// All functions only receive immutable references to `self`.
+/// `LineAnnotation`s that want to store an internal position or
+/// state of some kind should use `Cell`. Using interior mutability for
+/// caches is preferable as otherwise a lot of lifetimes become invariant
+/// which complicates APIs a lot.
+pub trait LineAnnotation {
+    /// Resets the internal position to `char_idx`. This function is called
+    /// when a new traversal of a document starts.
+    ///
+    /// All `char_idx` passed to `insert_virtual_lines` are strictly monotonically increasing
+    /// with the first `char_idx` greater or equal to the `char_idx`
+    /// passed to this function.
+    ///
+    /// # Returns
+    ///
+    /// The `char_idx` of the next anchor this `LineAnnotation` is interested in,
+    /// replaces the currently registered anchor. Return `usize::MAX` to ignore
+    fn reset_pos(&mut self, _char_idx: usize) -> usize {
+        usize::MAX
+    }
+
+    /// Called when a text is concealed that contains an anchor registered by this `LineAnnotation`.
+    /// In this case the line decorations  **must** ensure that virtual text anchored within that
+    /// char range is skipped.
+    ///
+    /// # Returns
+    ///
+    /// The `char_idx` of the next anchor this `LineAnnotation` is interested in,
+    /// **after the end of conceal_end_char_idx**
+    /// replaces the currently registered anchor. Return `usize::MAX` to ignore
+    fn skip_concealed_anchors(&mut self, conceal_end_char_idx: usize) -> usize {
+        self.reset_pos(conceal_end_char_idx)
+    }
+
+    /// Process an anchor (horizontal position is provided) and returns the next anchor.
+    ///
+    /// # Returns
+    ///
+    /// The `char_idx` of the next anchor this `LineAnnotation` is interested in,
+    /// replaces the currently registered anchor. Return `usize::MAX` to ignore
+    fn process_anchor(&mut self, _grapheme: &FormattedGrapheme) -> usize {
+        usize::MAX
+    }
+
+    /// This function is called at the end of a visual line to insert virtual text
+    ///
+    /// # Returns
+    ///
+    /// The number of additional virtual lines to reserve
+    ///
+    /// # Note
+    ///
+    /// The `line_end_visual_pos` parameter indicates the visual vertical distance
+    /// from the start of block where the traversal starts.  This includes the offset
+    /// from other `LineAnnotations`. This allows inline annotations to consider
+    /// the height of the text and "align" two different documents (like for side
+    /// by side diffs).  These annotations that want to "align" two documents should
+    /// therefore be added last so that other virtual text is also considered while aligning
+    fn insert_virtual_lines(
+        &mut self,
+        line_end_char_idx: usize,
+        line_end_visual_pos: Position,
+        doc_line: usize,
+    ) -> Position;
 }
 
 #[derive(Debug)]
@@ -143,23 +226,78 @@ fn reset_pos<A, M>(layers: &[Layer<A, M>], pos: usize, get_pos: impl Fn(&A) -> u
     }
 }
 
+/// Safety: We store LineAnnotation in a NonNull pointer. This is necessary to work
+/// around an unfortunate inconsistency in rusts variance system that unnnecesarily
+/// makes the lifetime invariant if implemented with safe code. This makes the
+/// DocFormatter API very cumbersome/basically impossible to work with.
+///
+/// Normally object types `dyn Foo + 'a` are covariant so if we used `Box<dyn LineAnnotation + 'a>` below
+/// everything would be alright. However we want to use `Cell<Box<dyn LineAnnotation + 'a>>`
+/// to be able to call the mutable function on `LineAnnotation`. The problem is that
+/// some types like `Cell` make all their arguments invariant. This is important for soundness
+/// normally for the same reasons that `&'a mut T` is invariant over `T`
+/// (see <https://doc.rust-lang.org/nomicon/subtyping.html>). However for `&'a mut` (`dyn Foo + 'b`)
+/// there is a specical rule in the language to make `'b` covariant (otherwise trait objects would be
+/// super annoying to use). See  <https://users.rust-lang.org/t/solved-variance-of-dyn-trait-a> for
+/// why this is sound. Sadly that rule doesn't apply to `Cell<Box<(dyn Foo + 'a)>`
+/// (or other invariant types like `UnsafeCell` or `*mut (dyn Foo + 'a)`).
+///
+/// We sidestep the problem by using `NonNull` which is covariant. In the
+/// special case of trait objects this is sound (easily checked by adding a
+/// `PhantomData<&'a mut Foo + 'a)>` field). We don't need an explicit `Cell`
+/// type here because we never hand out any refereces to the trait objects. That
+/// means any reference to the pointer can create a valid multable reference
+/// that is covariant over `'a` (or in other words it's a raw pointer, as long as
+/// we don't hand out references we are free to do whatever we want).
+struct RawBox<T: ?Sized>(NonNull<T>);
+
+impl<T: ?Sized> RawBox<T> {
+    /// Safety: Only a single mutable reference
+    /// created by this function may exist at a given time.
+    #[allow(clippy::mut_from_ref)]
+    unsafe fn get(&self) -> &mut T {
+        &mut *self.0.as_ptr()
+    }
+}
+impl<T: ?Sized> From<Box<T>> for RawBox<T> {
+    fn from(box_: Box<T>) -> Self {
+        // obviously safe because Box::into_raw never returns null
+        unsafe { Self(NonNull::new_unchecked(Box::into_raw(box_))) }
+    }
+}
+
+impl<T: ?Sized> Drop for RawBox<T> {
+    fn drop(&mut self) {
+        unsafe { drop(Box::from_raw(self.0.as_ptr())) }
+    }
+}
+
 /// Annotations that change that is displayed when the document is render.
 /// Also commonly called virtual text.
-#[derive(Default, Debug, Clone)]
+#[derive(Default)]
 pub struct TextAnnotations<'a> {
     inline_annotations: Vec<Layer<'a, InlineAnnotation, Option<Highlight>>>,
     overlays: Vec<Layer<'a, Overlay, Option<Highlight>>>,
-    line_annotations: Vec<Layer<'a, LineAnnotation, ()>>,
+    line_annotations: Vec<(Cell<usize>, RawBox<dyn LineAnnotation + 'a>)>,
+}
+
+impl Debug for TextAnnotations<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("TextAnnotations")
+            .field("inline_annotations", &self.inline_annotations)
+            .field("overlays", &self.overlays)
+            .finish_non_exhaustive()
+    }
 }
 
 impl<'a> TextAnnotations<'a> {
     /// Prepare the TextAnnotations for iteration starting at char_idx
     pub fn reset_pos(&self, char_idx: usize) {
         reset_pos(&self.inline_annotations, char_idx, |annot| annot.char_idx);
         reset_pos(&self.overlays, char_idx, |annot| annot.char_idx);
-        reset_pos(&self.line_annotations, char_idx, |annot| {
-            annot.anchor_char_idx
-        });
+        for (next_anchor, layer) in &self.line_annotations {
+            next_anchor.set(unsafe { layer.get().reset_pos(char_idx) });
+        }
     }
 
     pub fn collect_overlay_highlights(
@@ -219,8 +357,9 @@ impl<'a> TextAnnotations<'a> {
     ///
     /// The line annotations **must be sorted** by their `char_idx`.
     /// Multiple line annotations with the same `char_idx` **are not allowed**.
-    pub fn add_line_annotation(&mut self, layer: &'a [LineAnnotation]) -> &mut Self {
-        self.line_annotations.push((layer, ()).into());
+    pub fn add_line_annotation(&mut self, layer: Box<dyn LineAnnotation + 'a>) -> &mut Self {
+        self.line_annotations
+            .push((Cell::new(usize::MAX), layer.into()));
         self
     }
 
@@ -250,21 +389,35 @@ impl<'a> TextAnnotations<'a> {
         overlay
     }
 
-    pub(crate) fn annotation_lines_at(&self, char_idx: usize) -> usize {
-        self.line_annotations
-            .iter()
-            .map(|layer| {
-                let mut lines = 0;
-                while let Some(annot) = layer.annotations.get(layer.current_index.get()) {
-                    if annot.anchor_char_idx == char_idx {
-                        layer.current_index.set(layer.current_index.get() + 1);
-                        lines += annot.height
-                    } else {
-                        break;
+    pub(crate) fn process_virtual_text_anchors(&self, grapheme: &FormattedGrapheme) {
+        for (next_anchor, layer) in &self.line_annotations {
+            loop {
+                match next_anchor.get().cmp(&grapheme.char_idx) {
+                    Ordering::Less => next_anchor
+                        .set(unsafe { layer.get().skip_concealed_anchors(grapheme.char_idx) }),
+                    Ordering::Equal => {
+                        next_anchor.set(unsafe { layer.get().process_anchor(grapheme) })
                     }
-                }
-                lines
-            })
-            .sum()
+                    Ordering::Greater => break,
+                };
+            }
+        }
+    }
+
+    pub(crate) fn virtual_lines_at(
+        &self,
+        char_idx: usize,
+        line_end_visual_pos: Position,
+        doc_line: usize,
+    ) -> usize {
+        let mut virt_off = Position::new(0, 0);
+        for (_, layer) in &self.line_annotations {
+            virt_off += unsafe {
+                layer
+                    .get()
+                    .insert_virtual_lines(char_idx, line_end_visual_pos + virt_off, doc_line)
+            };
+        }
+        virt_off.row
     }
 }