3 files changed, 38 insertions, 5 deletions
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"));