feat: add `undo!`

Allows using `undo!` to undo changes and remove them from the undo-tree. Can only be used for moving backwards in the same undo branch.

5 files changed, 113 insertions, 5 deletions

diff --git a/runtime/doc/undo.txt b/runtime/doc/undo.txt
index a853aea995..39b26476af 100644
--- a/runtime/doc/undo.txt
+++ b/runtime/doc/undo.txt
@@ -22,6 +22,14 @@ u			Undo [count] changes.
 :u[ndo] {N}		Jump to after change number {N}.  See |undo-branches|
 			for the meaning of {N}.
 
+:u[ndo]!		Undo one change and remove it from undo history.
+								*E5767*
+:u[ndo]! {N}		Like ":u[ndo] {N}", but forget all changes in the
+			current undo branch up until {N}. You may only use
+			":undo! {N}" to move backwards in the same undo
+			branch, not to redo or switch to a different undo
+			branch.
+
 							*CTRL-R*
 CTRL-R			Redo [count] changes which were undone.
 
diff --git a/src/nvim/ex_cmds.lua b/src/nvim/ex_cmds.lua
index c2d40c8bb7..427e018141 100644
--- a/src/nvim/ex_cmds.lua
+++ b/src/nvim/ex_cmds.lua
@@ -2947,7 +2947,7 @@ module.cmds = {
   },
   {
     command='undo',
-    flags=bit.bor(RANGE, COUNT, ZEROR, TRLBAR, CMDWIN),
+    flags=bit.bor(BANG, RANGE, COUNT, ZEROR, TRLBAR, CMDWIN),
     addr_type='ADDR_OTHER',
     func='ex_undo',
   },
diff --git a/src/nvim/ex_docmd.c b/src/nvim/ex_docmd.c
index fc0bb48803..b67754ffe5 100644
--- a/src/nvim/ex_docmd.c
+++ b/src/nvim/ex_docmd.c
@@ -76,6 +76,7 @@
 #include "nvim/terminal.h"
 #include "nvim/ui.h"
 #include "nvim/undo.h"
+#include "nvim/undo_defs.h"
 #include "nvim/version.h"
 #include "nvim/vim.h"
 #include "nvim/window.h"
@@ -8231,10 +8232,39 @@ static void ex_bang(exarg_T *eap)
 /// ":undo".
 static void ex_undo(exarg_T *eap)
 {
-  if (eap->addr_count == 1) {       // :undo 123
-    undo_time(eap->line2, false, false, true);
-  } else {
-    u_undo(1);
+  if (eap->addr_count != 1) {
+    if (eap->forceit) {
+      u_undo_and_forget(1);         // :undo!
+    } else {
+      u_undo(1);                    // :undo
+    }
+    return;
+  }
+
+  long step = eap->line2;
+
+  if (eap->forceit) {             // undo! 123
+    // change number for "undo!" must be lesser than current change number
+    if (step >= curbuf->b_u_seq_cur) {
+      emsg(_(e_undobang_cannot_redo_or_move_branch));
+      return;
+    }
+    // ensure that target change number is in same branch
+    // while also counting the amount of undoes it'd take to reach target
+    u_header_T *uhp;
+    int count = 0;
+
+    for (uhp = curbuf->b_u_curhead ? curbuf->b_u_curhead : curbuf->b_u_newhead;
+         uhp != NULL && uhp->uh_seq > step;
+         uhp = uhp->uh_next.ptr, ++count) {
+    }
+    if (step != 0 && (uhp == NULL || uhp->uh_seq < step)) {
+      emsg(_(e_undobang_cannot_redo_or_move_branch));
+      return;
+    }
+    u_undo_and_forget(count);
+  } else {                        // :undo 123
+    undo_time(step, false, false, true);
   }
 }
 
diff --git a/src/nvim/globals.h b/src/nvim/globals.h
index 2b85b6a208..e07a0e22ca 100644
--- a/src/nvim/globals.h
+++ b/src/nvim/globals.h
@@ -1013,6 +1013,9 @@ EXTERN char e_line_number_out_of_range[] INIT(= N_("E1247: Line number out of ra
 
 EXTERN char e_highlight_group_name_too_long[] INIT(= N_("E1249: Highlight group name too long"));
 
+EXTERN char e_undobang_cannot_redo_or_move_branch[]
+INIT(= N_("E5767: Cannot use :undo! to redo or move to a different undo branch"));
+
 EXTERN char top_bot_msg[] INIT(= N_("search hit TOP, continuing at BOTTOM"));
 EXTERN char bot_top_msg[] INIT(= N_("search hit BOTTOM, continuing at TOP"));
 
diff --git a/test/functional/editor/undo_spec.lua b/test/functional/editor/undo_spec.lua
index a023ca3d90..a041428cdc 100644
--- a/test/functional/editor/undo_spec.lua
+++ b/test/functional/editor/undo_spec.lua
@@ -2,9 +2,18 @@ local helpers = require('test.functional.helpers')(after_each)
 
 local clear = helpers.clear
 local command = helpers.command
+local eval = helpers.eval
 local expect = helpers.expect
+local eq = helpers.eq
 local feed = helpers.feed
+local feed_command = helpers.feed_command
 local insert = helpers.insert
+local funcs = helpers.funcs
+
+local function lastmessage()
+  local messages = funcs.split(funcs.execute('messages'), '\n')
+  return messages[#messages]
+end
 
 describe('u CTRL-R g- g+', function()
   before_each(clear)
@@ -59,3 +68,61 @@ describe('u CTRL-R g- g+', function()
     undo_and_redo(4, 'g-', 'g+', '1')
   end)
 end)
+
+describe(':undo! command', function()
+  before_each(function()
+    clear()
+    feed('i1 little bug in the code<Esc>')
+    feed('o1 little bug in the code<Esc>')
+    feed('oTake 1 down, patch it around<Esc>')
+    feed('o99 little bugs in the code<Esc>')
+  end)
+  it('works', function()
+    feed_command('undo!')
+    expect([[
+      1 little bug in the code
+      1 little bug in the code
+      Take 1 down, patch it around]])
+    feed('<C-r>')
+    eq('Already at newest change', lastmessage())
+  end)
+  it('works with arguments', function()
+    feed_command('undo! 2')
+    expect([[
+      1 little bug in the code
+      1 little bug in the code]])
+    feed('<C-r>')
+    eq('Already at newest change', lastmessage())
+  end)
+  it('correctly sets alternative redo', function()
+    feed('uo101 little bugs in the code<Esc>')
+    feed_command('undo!')
+    feed('<C-r>')
+    expect([[
+      1 little bug in the code
+      1 little bug in the code
+      Take 1 down, patch it around
+      99 little bugs in the code]])
+
+    feed('uuoTake 2 down, patch them around<Esc>')
+    feed('o101 little bugs in the code<Esc>')
+    feed_command('undo! 2')
+    feed('<C-r><C-r>')
+    expect([[
+      1 little bug in the code
+      1 little bug in the code
+      Take 1 down, patch it around
+      99 little bugs in the code]])
+  end)
+  it('fails when attempting to redo or move to different undo branch', function()
+    feed_command('undo! 4')
+    eq('E5767: Cannot use :undo! to redo or move to a different undo branch', eval('v:errmsg'))
+    feed('u')
+    feed_command('undo! 4')
+    eq('E5767: Cannot use :undo! to redo or move to a different undo branch', eval('v:errmsg'))
+    feed('o101 little bugs in the code<Esc>')
+    feed('o101 little bugs in the code<Esc>')
+    feed_command('undo! 4')
+    eq('E5767: Cannot use :undo! to redo or move to a different undo branch', eval('v:errmsg'))
+  end)
+end)