optimize cursor movement a bit

4 files changed, 157 insertions, 86 deletions

diff --git a/src/grid.rs b/src/grid.rs
index ee608ca..1cc4daf 100644
--- a/src/grid.rs
+++ b/src/grid.rs
@@ -189,30 +189,25 @@ impl Grid {
         .unwrap();
 
         let mut prev_attrs = crate::attrs::Attrs::default();
-        let mut final_col = 0;
-        for row in self.visible_rows() {
-            let (new_attrs, new_col) = row.write_contents_formatted(
+        let mut prev_pos = Pos::default();
+        let mut wrapping = false;
+        for (i, row) in self.visible_rows().enumerate() {
+            let i = i.try_into().unwrap();
+            let (new_pos, new_attrs) = row.write_contents_formatted(
                 contents,
                 0,
                 self.size.cols,
+                i,
+                wrapping,
+                prev_pos,
                 prev_attrs,
             );
-            if let Some(col) = new_col {
-                final_col = col;
-            }
-            if !row.wrapped() {
-                write!(contents, "{}", crate::term::CRLF::new()).unwrap();
-            }
+            prev_pos = new_pos;
             prev_attrs = new_attrs;
+            wrapping = row.wrapped();
         }
 
-        let mut final_row = self.size.rows;
-        while contents.ends_with(b"\r\n") {
-            contents.truncate(contents.len() - 2);
-            final_row -= 1;
-        }
-
-        if final_row != self.pos.row || final_col != self.pos.col {
+        if prev_pos != self.pos {
             write!(contents, "{}", crate::term::MoveTo::new(self.pos))
                 .unwrap();
         }
@@ -220,36 +215,30 @@ impl Grid {
 
     pub fn write_contents_diff(&self, contents: &mut Vec<u8>, prev: &Self) {
         write!(contents, "{}", crate::term::Attrs::default()).unwrap();
+
+        let mut prev_pos = prev.pos;
         let mut prev_attrs = crate::attrs::Attrs::default();
-        let mut final_row = prev.pos.row;
-        let mut final_col = prev.pos.col;
-        for (idx, (row, prev_row)) in
+        let mut wrapping = false;
+        for (i, (row, prev_row)) in
             self.visible_rows().zip(prev.visible_rows()).enumerate()
         {
-            let idx = idx.try_into().unwrap();
-            let (new_attrs, new_col) = row.write_contents_diff(
+            let i = i.try_into().unwrap();
+            let (new_pos, new_attrs) = row.write_contents_diff(
                 contents,
                 prev_row,
-                |contents| {
-                    write!(
-                        contents,
-                        "{}",
-                        crate::term::MoveTo::new(Pos { row: idx, col: 0 })
-                    )
-                    .unwrap();
-                },
                 0,
                 self.size.cols,
+                i,
+                wrapping,
+                prev_pos,
                 prev_attrs,
             );
-            if let Some(col) = new_col {
-                final_row = idx;
-                final_col = col;
-            }
+            prev_pos = new_pos;
             prev_attrs = new_attrs;
+            wrapping = row.wrapped();
         }
 
-        if self.pos.row != final_row || self.pos.col != final_col {
+        if prev_pos != self.pos {
             write!(contents, "{}", crate::term::MoveTo::new(self.pos))
                 .unwrap();
         }
diff --git a/src/lib.rs b/src/lib.rs
index d14b78f..7f391c6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -26,6 +26,7 @@
 #![warn(clippy::nursery)]
 #![allow(clippy::missing_const_for_fn)]
 #![allow(clippy::single_match)]
+#![allow(clippy::too_many_arguments)]
 
 mod attrs;
 mod cell;
diff --git a/src/row.rs b/src/row.rs
index 8008212..f7d4485 100644
--- a/src/row.rs
+++ b/src/row.rs
@@ -15,6 +15,10 @@ impl Row {
         }
     }
 
+    fn cols(&self) -> u16 {
+        self.cells.len().try_into().unwrap()
+    }
+
     pub fn clear(&mut self, bgcolor: crate::attrs::Color) {
         for cell in &mut self.cells {
             cell.clear(bgcolor);
@@ -99,12 +103,16 @@ impl Row {
         contents: &mut Vec<u8>,
         start: u16,
         width: u16,
+        row: u16,
+        wrapping: bool,
+        pos: crate::grid::Pos,
         attrs: crate::attrs::Attrs,
-    ) -> (crate::attrs::Attrs, Option<u16>) {
+    ) -> (crate::grid::Pos, crate::attrs::Attrs) {
         let mut prev_was_wide = false;
+        let mut prev_pos = pos;
         let mut prev_attrs = attrs;
 
-        let mut cols = None;
+        let mut new_pos = crate::grid::Pos { row, col: start };
         for cell in self
             .cells()
             .skip(start as usize)
@@ -115,52 +123,92 @@ impl Row {
                 continue;
             }
 
-            let attrs = cell.attrs();
-            if &prev_attrs != attrs {
-                attrs.write_escape_code_diff(contents, &prev_attrs);
-                prev_attrs = *attrs;
-            }
+            prev_was_wide = cell.is_wide();
 
-            if cell.has_contents() {
-                // using write! here is significantly slower, for some reason
-                // write!(contents, "{}", cell.contents()).unwrap();
-                contents.extend(cell.contents().as_bytes());
-            } else if cell.bgcolor() == crate::attrs::Color::Default {
-                write!(contents, "{}", crate::term::MoveRight::default())
-                    .unwrap();
+            let has_contents = cell.has_contents();
+            if !has_contents && cell.bgcolor() == crate::attrs::Color::Default
+            {
+                new_pos.col += 1;
             } else {
-                write!(
-                    contents,
-                    "{}{}",
-                    crate::term::EraseChar::default(),
-                    crate::term::MoveRight::default()
-                )
-                .unwrap();
-            }
+                if new_pos != prev_pos {
+                    if new_pos.row == prev_pos.row + 1 {
+                        if !wrapping
+                            || prev_pos.col != self.cols()
+                            || new_pos.col != 0
+                        {
+                            write!(
+                                contents,
+                                "{}{}",
+                                crate::term::CRLF::default(),
+                                crate::term::MoveRight::new(new_pos.col)
+                            )
+                            .unwrap();
+                        }
+                    } else if prev_pos.row == new_pos.row {
+                        write!(
+                            contents,
+                            "{}",
+                            crate::term::MoveRight::new(
+                                new_pos.col - prev_pos.col
+                            )
+                        )
+                        .unwrap();
+                    } else {
+                        write!(
+                            contents,
+                            "{}",
+                            crate::term::MoveTo::new(new_pos)
+                        )
+                        .unwrap();
+                    }
+                    prev_pos = new_pos;
+                }
 
-            prev_was_wide = cell.is_wide();
-            cols =
-                Some(cols.unwrap_or(0) + if prev_was_wide { 2 } else { 1 });
+                let attrs = cell.attrs();
+                if &prev_attrs != attrs {
+                    attrs.write_escape_code_diff(contents, &prev_attrs);
+                    prev_attrs = *attrs;
+                }
+
+                if has_contents {
+                    // using write! here is significantly slower, for some
+                    // reason
+                    // write!(contents, "{}", cell.contents()).unwrap();
+                    contents.extend(cell.contents().as_bytes());
+                    let width = if prev_was_wide { 2 } else { 1 };
+                    prev_pos.col += width;
+                    new_pos.col += width;
+                } else {
+                    write!(
+                        contents,
+                        "{}",
+                        crate::term::EraseChar::default(),
+                    )
+                    .unwrap();
+                    new_pos.col += 1;
+                }
+            }
         }
 
-        (prev_attrs, cols)
+        (prev_pos, prev_attrs)
     }
 
-    pub fn write_contents_diff<F: for<'a> Fn(&'a mut Vec<u8>)>(
+    pub fn write_contents_diff(
         &self,
         contents: &mut Vec<u8>,
         prev: &Self,
-        initial_pos: F,
         start: u16,
         width: u16,
+        row: u16,
+        wrapping: bool,
+        pos: crate::grid::Pos,
         attrs: crate::attrs::Attrs,
-    ) -> (crate::attrs::Attrs, Option<u16>) {
+    ) -> (crate::grid::Pos, crate::attrs::Attrs) {
         let mut prev_was_wide = false;
+        let mut prev_pos = pos;
         let mut prev_attrs = attrs;
-        let mut skip = 0;
-        let mut cols = None;
-        let mut initial_pos = Some(initial_pos);
 
+        let mut new_pos = crate::grid::Pos { row, col: start };
         for (cell, prev_cell) in self
             .cells()
             .zip(prev.cells())
@@ -172,18 +220,45 @@ impl Row {
                 continue;
             }
 
+            prev_was_wide = cell.is_wide();
+
             if cell == prev_cell {
-                prev_was_wide = cell.is_wide();
-                skip += if prev_was_wide { 2 } else { 1 };
+                new_pos.col += if prev_was_wide { 2 } else { 1 };
             } else {
-                if let Some(f) = initial_pos.take() {
-                    f(contents)
-                }
-                if skip > 0 {
-                    write!(contents, "{}", crate::term::MoveRight::new(skip))
+                if new_pos != prev_pos {
+                    if new_pos.row == prev_pos.row + 1 {
+                        if !wrapping
+                            || prev_pos.col != self.cols()
+                            || new_pos.col != 0
+                        {
+                            write!(
+                                contents,
+                                "{}{}",
+                                crate::term::CRLF::default(),
+                                crate::term::MoveRight::new(new_pos.col)
+                            )
+                            .unwrap();
+                        }
+                    } else if prev_pos.row == new_pos.row
+                        && prev_pos.col < new_pos.col
+                    {
+                        write!(
+                            contents,
+                            "{}",
+                            crate::term::MoveRight::new(
+                                new_pos.col - prev_pos.col
+                            )
+                        )
                         .unwrap();
-                    cols = Some(cols.unwrap_or(0) + skip);
-                    skip = 0;
+                    } else {
+                        write!(
+                            contents,
+                            "{}",
+                            crate::term::MoveTo::new(new_pos)
+                        )
+                        .unwrap();
+                    }
+                    prev_pos = new_pos;
                 }
 
                 let attrs = cell.attrs();
@@ -197,24 +272,22 @@ impl Row {
                     // reason
                     // write!(contents, "{}", cell.contents()).unwrap();
                     contents.extend(cell.contents().as_bytes());
+                    let width = if prev_was_wide { 2 } else { 1 };
+                    prev_pos.col += width;
+                    new_pos.col += width;
                 } else {
                     write!(
                         contents,
-                        "{}{}",
+                        "{}",
                         crate::term::EraseChar::default(),
-                        crate::term::MoveRight::default()
                     )
                     .unwrap();
+                    new_pos.col += 1;
                 }
-
-                prev_was_wide = cell.is_wide();
-                cols = Some(
-                    cols.unwrap_or(0) + if prev_was_wide { 2 } else { 1 },
-                );
             }
         }
 
-        (prev_attrs, cols)
+        (prev_pos, prev_attrs)
     }
 
     fn content_width(&self, start: u16, formatting: bool) -> u16 {
diff --git a/src/screen.rs b/src/screen.rs
index 5150e3a..b774c10 100644
--- a/src/screen.rs
+++ b/src/screen.rs
@@ -207,12 +207,16 @@ impl Screen {
         start: u16,
         width: u16,
     ) -> impl Iterator<Item = Vec<u8>> + '_ {
-        self.grid().visible_rows().map(move |row| {
+        self.grid().visible_rows().enumerate().map(move |(i, row)| {
+            let i = i.try_into().unwrap();
             let mut contents = vec![];
             row.write_contents_formatted(
                 &mut contents,
                 start,
                 width,
+                i,
+                false,
+                crate::grid::Pos { row: i, col: start },
                 crate::attrs::Attrs::default(),
             );
             contents
@@ -263,14 +267,18 @@ impl Screen {
         self.grid()
             .visible_rows()
             .zip(prev.grid().visible_rows())
-            .map(move |(row, prev_row)| {
+            .enumerate()
+            .map(move |(i, (row, prev_row))| {
+                let i = i.try_into().unwrap();
                 let mut contents = vec![];
                 row.write_contents_diff(
                     &mut contents,
                     prev_row,
-                    |_| (),
                     start,
                     width,
+                    i,
+                    false,
+                    crate::grid::Pos { row: i, col: start },
                     crate::attrs::Attrs::default(),
                 );
                 contents