less broken multiline editing

1 files changed, 18 insertions, 1 deletions

diff --git a/src/syntax_highlighting.rs b/src/syntax_highlighting.rs
index 043509a..19f44ab 100644
--- a/src/syntax_highlighting.rs
+++ b/src/syntax_highlighting.rs
@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use crate::{
-    BString, ansi, bstr,
+    BString, ansi, bstr, cursor,
     parse::{Highlight, HighlightKind},
 };
 
@@ -15,6 +15,8 @@ pub enum SetColorError {
     NoSuchKeyword,
 }
 
+const CLEAR_TO_END: &bstr = b"\x1b[K";
+
 impl Highlighter {
     pub fn new() -> Self {
         let mut this = Self {
@@ -51,6 +53,7 @@ impl Highlighter {
     pub fn pretty_print(
         &self,
         bytes: &[u8],
+        cursor_offset: usize,
         colors: Vec<Highlight>,
         stdout: &mut dyn std::io::Write,
     ) -> std::io::Result<()> {
@@ -104,7 +107,16 @@ impl Highlighter {
 
         let mut current_color = self.color(HighlightKind::None);
 
+        let mut saved_pos = false;
+
+        stdout.write_all(CLEAR_TO_END)?;
+
         for (i, x) in bytes.iter().cloned().enumerate() {
+            if i == cursor_offset {
+                cursor::f_save(stdout)?;
+                saved_pos = true;
+            }
+
             while let Some(color_boundary) = coloring.first().cloned()
                 && color_boundary.loc <= i
             {
@@ -126,6 +138,7 @@ impl Highlighter {
 
             if x == b'\n' {
                 stdout.write_all(b"\r\n")?;
+                stdout.write_all(CLEAR_TO_END)?;
             } else {
                 stdout.write_all(&[x])?;
             }
@@ -133,6 +146,10 @@ impl Highlighter {
 
         stdout.write_all(self.color(HighlightKind::None))?;
 
+        if saved_pos {
+            cursor::f_restore(stdout)?;
+        }
+
         Ok(())
     }
 }