Fix (properly) the logic around prompt re-use & Host Command handling

src/engine.rs

@@ -27,7 +27,7 @@ use {
             FileBackedHistory, History, HistoryCursor, HistoryItem, HistoryItemId,
             HistoryNavigationQuery, HistorySessionId, SearchDirection, SearchQuery,
         },
-        painting::{Painter, PromptLines},
+        painting::{Painter, PainterSuspendedState, PromptLines},
         prompt::{PromptEditMode, PromptHistorySearchStatus},
         result::{ReedlineError, ReedlineErrorVariants},
         terminal_extensions::{bracketed_paste::BracketedPasteGuard, kitty::KittyProtocolGuard},
@@ -109,8 +109,9 @@ pub struct Reedline {
     history_cursor_on_excluded: bool,
     input_mode: InputMode,
 
-    // Yielded to the host program after a `ReedlineEvent::ExecuteHostCommand`, thus redraw in-place
-    executing_host_command: bool,
+    // State of the painter after a `ReedlineEvent::ExecuteHostCommand` was requested, used after
+    // execution to decide if we can re-use the previous prompt or paint a new one.
+    suspended_state: Option<PainterSuspendedState>,
 
     // Validator
     validator: Option<Box<dyn Validator>>,
@@ -210,7 +211,7 @@ impl Reedline {
             history_excluded_item: None,
             history_cursor_on_excluded: false,
             input_mode: InputMode::Regular,
-            executing_host_command: false,
+            suspended_state: None,
             painter,
             transient_prompt: None,
             edit_mode,
@@ -671,12 +672,14 @@ impl Reedline {
     /// Helper implementing the logic for [`Reedline::read_line()`] to be wrapped
     /// in a `raw_mode` context.
     fn read_line_helper(&mut self, prompt: &dyn Prompt) -> Result<Signal> {
-        if self.executing_host_command {
-            self.executing_host_command = false;
-        } else {
-            self.painter.initialize_prompt_position()?;
-            self.hide_hints = false;
+        self.painter
+            .initialize_prompt_position(self.suspended_state.as_ref())?;
+        if self.suspended_state.is_some() {
+            // Last editor was suspended to run a ExecuteHostCommand event,
+            // we are resuming operation now.
+            self.suspended_state = None;
         }
+        self.hide_hints = false;
 
         self.repaint(prompt)?;
 
@@ -773,8 +776,11 @@ impl Reedline {
             for event in reedline_events.drain(..) {
                 match self.handle_event(prompt, event)? {
                     EventStatus::Exits(signal) => {
-                        if !self.executing_host_command {
-                            // Move the cursor below the input area, for external commands or new read_line call
+                        // Check if we are merely suspended (to process an ExecuteHostCommand event)
+                        // or if we're about to quit the editor.
+                        if self.suspended_state.is_none() {
+                            // We are about to quit the editor, move the cursor below the input
+                            // area, for external commands or new read_line call
                             self.painter.move_cursor_to_end()?;
                         }
                         return Ok(signal);
@@ -851,8 +857,7 @@ impl Reedline {
                 Ok(EventStatus::Handled)
             }
             ReedlineEvent::ExecuteHostCommand(host_command) => {
-                // TODO: Decide if we need to do something special to have a nicer painter state on the next go
-                self.executing_host_command = true;
+                self.suspended_state = Some(self.painter.state_before_suspension());
                 Ok(EventStatus::Exits(Signal::Success(host_command)))
             }
             ReedlineEvent::Edit(commands) => {
@@ -1122,8 +1127,7 @@ impl Reedline {
                 }
             }
             ReedlineEvent::ExecuteHostCommand(host_command) => {
-                // TODO: Decide if we need to do something special to have a nicer painter state on the next go
-                self.executing_host_command = true;
+                self.suspended_state = Some(self.painter.state_before_suspension());
                 Ok(EventStatus::Exits(Signal::Success(host_command)))
             }
             ReedlineEvent::Edit(commands) => {

src/painting/mod.rs

@@ -3,7 +3,7 @@ mod prompt_lines;
 mod styled_text;
 mod utils;
 
-pub use painter::Painter;
+pub use painter::{Painter, PainterSuspendedState};
 pub(crate) use prompt_lines::PromptLines;
 pub use styled_text::StyledText;
 pub(crate) use utils::estimate_single_line_wraps;

src/painting/painter.rs

@@ -14,6 +14,7 @@ use {
         QueueableCommand,
     },
     std::io::{Result, Write},
+    std::ops::RangeInclusive,
 };
 #[cfg(feature = "external_printer")]
 use {crate::LineBuffer, crossterm::cursor::MoveUp};
@@ -49,6 +50,42 @@ fn skip_buffer_lines(string: &str, skip: usize, offset: Option<usize>) -> &str {
 /// the type used by crossterm operations
 pub type W = std::io::BufWriter<std::io::Stderr>;
 
+#[derive(Debug, PartialEq, Eq)]
+pub struct PainterSuspendedState {
+    previous_prompt_rows_range: RangeInclusive<u16>,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum PromptRowSelector {
+    UseExistingPrompt { start_row: u16 },
+    MakeNewPrompt { new_row: u16 },
+}
+
+// Selects the row where the next prompt should start on, taking into account and whether it should re-use a previous
+// prompt.
+fn select_prompt_row(
+    suspended_state: Option<&PainterSuspendedState>,
+    (column, row): (u16, u16), // NOTE: Positions are 0 based here
+) -> PromptRowSelector {
+    if let Some(painter_state) = suspended_state {
+        // The painter was suspended, try to re-use the last prompt position to avoid
+        // unnecessarily making new prompts.
+        if painter_state.previous_prompt_rows_range.contains(&row) {
+            // Cursor is still in the range of the previous prompt, re-use it.
+            let start_row = *painter_state.previous_prompt_rows_range.start();
+            return PromptRowSelector::UseExistingPrompt { start_row };
+        } else {
+            // There was some output or cursor is outside of the range of previous prompt make a
+            // fresh new prompt.
+        }
+    }
+
+    // Assumption: if the cursor is not on the zeroth column,
+    //   there is content we want to leave intact, thus advance to the next row.
+    let new_row = if column > 0 { row + 1 } else { row };
+    PromptRowSelector::MakeNewPrompt { new_row }
+}
+
 /// Implementation of the output to the terminal
 pub struct Painter {
     // Stdout
@@ -85,12 +122,26 @@ impl Painter {
         self.screen_height().saturating_sub(self.prompt_start_row)
     }
 
+    /// Returns the state necessary before suspending the painter (to run a host command event).
+    ///
+    /// This state will be used to re-initialize the painter to re-use last prompt if possible.
+    pub fn state_before_suspension(&self) -> PainterSuspendedState {
+        let start_row = self.prompt_start_row;
+        let final_row = start_row + self.last_required_lines;
+        PainterSuspendedState {
+            previous_prompt_rows_range: start_row..=final_row,
+        }
+    }
+
     /// Sets the prompt origin position and screen size for a new line editor
     /// invocation
     ///
     /// Not to be used for resizes during a running line editor, use
     /// [`Painter::handle_resize()`] instead
-    pub(crate) fn initialize_prompt_position(&mut self) -> Result<()> {
+    pub(crate) fn initialize_prompt_position(
+        &mut self,
+        suspended_state: Option<&PainterSuspendedState>,
+    ) -> Result<()> {
         // Update the terminal size
         self.terminal_size = {
             let size = terminal::size()?;
@@ -102,26 +153,26 @@ impl Painter {
                 size
             }
         };
-        // Cursor positions are 0 based here.
-        let (column, row) = cursor::position()?;
-        // Assumption: if the cursor is not on the zeroth column,
-        // there is content we want to leave intact, thus advance to the next row
-        let new_row = if column > 0 { row + 1 } else { row };
-        //  If we are on the last line and would move beyond the last line due to
-        //  the condition above, we need to make room for the prompt.
-        //  Otherwise printing the prompt would scroll of the stored prompt
-        //  origin, causing issues after repaints.
-        let new_row = if new_row == self.screen_height() {
-            self.print_crlf()?;
-            new_row.saturating_sub(1)
-        } else {
-            new_row
+        let prompt_selector = select_prompt_row(suspended_state, cursor::position()?);
+        self.prompt_start_row = match prompt_selector {
+            PromptRowSelector::UseExistingPrompt { start_row } => start_row,
+            PromptRowSelector::MakeNewPrompt { new_row } => {
+                // If we are on the last line and would move beyond the last line, we need to make
+                // room for the prompt.
+                // Otherwise printing the prompt would scroll off the stored prompt
+                // origin, causing issues after repaints.
+                if new_row == self.screen_height() {
+                    self.print_crlf()?;
+                    new_row.saturating_sub(1)
+                } else {
+                    new_row
+                }
+            }
         };
-        self.prompt_start_row = new_row;
         Ok(())
     }
 
-    /// Main pain painter for the prompt and buffer
+    /// Main painter for the prompt and buffer
     /// It queues all the actions required to print the prompt together with
     /// lines that make the buffer.
     /// Using the prompt lines object in this function it is estimated how the
@@ -461,7 +512,7 @@ impl Painter {
         self.stdout.queue(cursor::Show)?;
 
         self.stdout.flush()?;
-        self.initialize_prompt_position()
+        self.initialize_prompt_position(None)
     }
 
     pub(crate) fn clear_scrollback(&mut self) -> Result<()> {
@@ -470,11 +521,11 @@ impl Painter {
             .queue(crossterm::terminal::Clear(ClearType::Purge))?
             .queue(cursor::MoveTo(0, 0))?
             .flush()?;
-        self.initialize_prompt_position()
+        self.initialize_prompt_position(None)
     }
 
     // The prompt is moved to the end of the buffer after the event was handled
-    // If the prompt is in the middle of a multiline buffer, then the output to stdout
+    // If the prompt is in the middle of a multiline buffer, then the next output to stdout
     // could overwrite the buffer writing
     pub(crate) fn move_cursor_to_end(&mut self) -> Result<()> {
         let final_row = self.prompt_start_row + self.last_required_lines;
@@ -619,4 +670,35 @@ mod tests {
         assert_eq!(skip_buffer_lines(string, 0, Some(0)), "sentence1",);
         assert_eq!(skip_buffer_lines(string, 1, Some(0)), "sentence2",);
     }
+
+    #[test]
+    fn test_select_new_prompt_with_no_state_no_output() {
+        assert_eq!(
+            select_prompt_row(None, (0, 12)),
+            PromptRowSelector::MakeNewPrompt { new_row: 12 }
+        );
+    }
+
+    #[test]
+    fn test_select_new_prompt_with_no_state_but_output() {
+        assert_eq!(
+            select_prompt_row(None, (3, 12)),
+            PromptRowSelector::MakeNewPrompt { new_row: 13 }
+        );
+    }
+
+    #[test]
+    fn test_select_existing_prompt() {
+        let state = PainterSuspendedState {
+            previous_prompt_rows_range: 11..=13,
+        };
+        assert_eq!(
+            select_prompt_row(Some(&state), (0, 12)),
+            PromptRowSelector::UseExistingPrompt { start_row: 11 }
+        );
+        assert_eq!(
+            select_prompt_row(Some(&state), (3, 12)),
+            PromptRowSelector::UseExistingPrompt { start_row: 11 }
+        );
+    }
 }