From f73a1b29824feaf16477b9df547fb28d9db81923 Mon Sep 17 00:00:00 2001
From: Kyle Smith <kyle.smith@salsify.com>
Date: Wed, 1 Feb 2023 08:09:12 -0500
Subject: [PATCH] Use jumplist for undo, set jumplist before preview.

---
 helix-term/src/commands.rs       | 11 +++--
 helix-term/src/commands/typed.rs | 69 ++++++++++++++++++--------------
 helix-view/src/editor.rs         |  7 +++-
 3 files changed, 50 insertions(+), 37 deletions(-)

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 1bdaafbdc..892046967 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -4313,15 +4313,18 @@ fn jump_forward(cx: &mut Context) {
 }
 
 fn jump_backward(cx: &mut Context) {
-    let count = cx.count();
-    let config = cx.editor.config();
-    let (view, doc) = current!(cx.editor);
+    jump_backward_impl(cx.editor, cx.count());
+}
+
+fn jump_backward_impl(editor: &mut Editor, count: usize) {
+    let config = editor.config();
+    let (view, doc) = current!(editor);
     let doc_id = doc.id();
 
     if let Some((id, selection)) = view.jumps.backward(view.id, doc, count) {
         view.doc = *id;
         let selection = selection.clone();
-        let (view, doc) = current!(cx.editor); // refetch doc
+        let (view, doc) = current!(editor); // refetch doc
 
         if doc.id() != doc_id {
             view.add_to_history(doc_id);
diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index d66496370..43efea05a 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -1410,53 +1410,60 @@ fn tutor(
     Ok(())
 }
 
+fn undo_goto_line_number_preview(cx: &mut compositor::Context) {
+    if cx.editor.goto_line_number_preview {
+        log::info!("undoing goto_line_number preview, jumping backwards");
+        // if let Some(line_number) = cx.editor.last_line_number {
+        // goto_line_without_jumplist(cx.editor, NonZeroUsize::new(line_number));
+
+        // Remove the jump we added during the preview session.
+        jump_backward_impl(cx.editor, 1);
+
+        // let (view, doc) = current!(cx.editor);
+        // view.ensure_cursor_in_view(doc, line_number);
+        cx.editor.goto_line_number_preview = false;
+    }
+}
+
+fn start_goto_line_number_preview(cx: &mut compositor::Context) {
+    if !cx.editor.goto_line_number_preview {
+        let (view, doc) = current!(cx.editor);
+
+        // Allow the user to jump back to the previous location before invoking
+        // `goto_line_number`.
+        push_jump(view, doc);
+
+        cx.editor.goto_line_number_preview = true;
+    }
+}
+
 pub(super) fn goto_line_number(
     cx: &mut compositor::Context,
     args: &[Cow<str>],
     event: PromptEvent,
 ) -> anyhow::Result<()> {
     match event {
-        PromptEvent::Abort => {
-            if let Some(line_number) = cx.editor.last_line_number {
-                goto_line_without_jumplist(cx.editor, NonZeroUsize::new(line_number));
-                let (view, doc) = current!(cx.editor);
-                view.ensure_cursor_in_view(doc, line_number);
-                cx.editor.last_line_number = None;
-            }
-            return Ok(());
-        }
+        PromptEvent::Abort => undo_goto_line_number_preview(cx),
+
         PromptEvent::Validate => {
             ensure!(!args.is_empty(), "Line number required");
-
-            let (view, doc) = current!(cx.editor);
-            push_jump(view, doc);
-
-            cx.editor.last_line_number = None;
+            cx.editor.goto_line_number_preview = false;
         }
+
         PromptEvent::Update => {
+            // When a user hits backspace and there are no numbers left,
+            // we can bring them back to their original selection(s).
             if args.is_empty() {
-                if let Some(line_number) = cx.editor.last_line_number {
-                    // When a user hits backspace and there are no numbers left,
-                    // we can bring them back to their original line
-                    goto_line_without_jumplist(cx.editor, NonZeroUsize::new(line_number));
-                    let (view, doc) = current!(cx.editor);
-                    view.ensure_cursor_in_view(doc, line_number);
-                    cx.editor.last_line_number = None;
-                }
+                undo_goto_line_number_preview(cx);
                 return Ok(());
             }
+            if !cx.editor.goto_line_number_preview {
+                start_goto_line_number_preview(cx);
+            }
 
-            cx.editor.last_line_number.get_or_insert_with(|| {
-                let (view, doc) = current!(cx.editor);
-
-                let text = doc.text().slice(..);
-                let line = doc.selection(view.id).primary().cursor_line(text);
-
-                line + 1
-            });
+            let line = args[0].parse::<usize>()?;
 
             let (view, doc) = current!(cx.editor);
-            let line = args[0].parse::<usize>()?;
             view.ensure_cursor_in_view(doc, line);
             goto_line_without_jumplist(cx.editor, NonZeroUsize::new(line));
         }
diff --git a/helix-view/src/editor.rs b/helix-view/src/editor.rs
index 1029c14f3..d8361a77b 100644
--- a/helix-view/src/editor.rs
+++ b/helix-view/src/editor.rs
@@ -793,7 +793,10 @@ pub struct Editor {
     /// The currently applied editor theme. While previewing a theme, the previewed theme
     /// is set here.
     pub theme: Theme,
-    pub last_line_number: Option<usize>,
+
+    // Are we currently previewing a goto_line_number typed command (`:goto <linenum>`, `:<linenum>`)?
+    pub goto_line_number_preview: bool,
+
     pub status_msg: Option<(Cow<'static, str>, Severity)>,
     pub autoinfo: Option<Info>,
 
@@ -896,7 +899,7 @@ pub fn new(
             syn_loader,
             theme_loader,
             last_theme: None,
-            last_line_number: None,
+            goto_line_number_preview: false,
             registers: Registers::default(),
             clipboard_provider: get_clipboard_provider(),
             status_msg: None,