1 files changed, 81 insertions, 37 deletions

diff --git a/src/nvim/undo.c b/src/nvim/undo.c
index f16b4264d7..4b267a1627 100644
--- a/src/nvim/undo.c
+++ b/src/nvim/undo.c
@@ -99,6 +99,7 @@
 #include "nvim/quickfix.h"
 #include "nvim/screen.h"
 #include "nvim/sha256.h"
+#include "nvim/state.h"
 #include "nvim/strings.h"
 #include "nvim/types.h"
 #include "nvim/os/os.h"
@@ -690,7 +691,7 @@ char *u_get_undo_file_name(const char *const buf_ffname, const bool reading)
         int ret;
         char *failed_dir;
         if ((ret = os_mkdir_recurse(dir_name, 0755, &failed_dir)) != 0) {
-          EMSG3(_("E926: Unable to create directory \"%s\" for undo file: %s"),
+          EMSG3(_("E5003: Unable to create directory \"%s\" for undo file: %s"),
                 failed_dir, os_strerror(ret));
           xfree(failed_dir);
         } else {
@@ -712,7 +713,7 @@ char *u_get_undo_file_name(const char *const buf_ffname, const bool reading)
 
     // When reading check if the file exists.
     if (undo_file_name != NULL
-        && (!reading || os_file_exists((char_u *)undo_file_name))) {
+        && (!reading || os_path_exists((char_u *)undo_file_name))) {
       break;
     }
     xfree(undo_file_name);
@@ -1094,7 +1095,7 @@ void u_write_undo(const char *const name, const bool forceit, buf_T *const buf,
 
   /* If the undo file already exists, verify that it actually is an undo
    * file, and delete it. */
-  if (os_file_exists((char_u *)file_name)) {
+  if (os_path_exists((char_u *)file_name)) {
     if (name == NULL || !forceit) {
       /* Check we can read it and it's an undo file. */
       fd = os_open(file_name, O_RDONLY, 0);
@@ -1668,7 +1669,7 @@ void u_undo(int count)
     undo_undoes = TRUE;
   else
     undo_undoes = !undo_undoes;
-  u_doit(count);
+  u_doit(count, false);
 }
 
 /*
@@ -1677,15 +1678,58 @@ void u_undo(int count)
  */
 void u_redo(int count)
 {
-  if (vim_strchr(p_cpo, CPO_UNDO) == NULL)
-    undo_undoes = FALSE;
-  u_doit(count);
+  if (vim_strchr(p_cpo, CPO_UNDO) == NULL) {
+    undo_undoes = false;
+  }
+
+  u_doit(count, false);
 }
 
-/*
- * Undo or redo, depending on 'undo_undoes', 'count' times.
- */
-static void u_doit(int startcount)
+/// Undo and remove the branch from the undo tree.
+/// Also moves the cursor (as a "normal" undo would).
+bool u_undo_and_forget(int count)
+{
+  if (curbuf->b_u_synced == false) {
+    u_sync(true);
+    count = 1;
+  }
+  undo_undoes = true;
+  u_doit(count, true);
+
+  if (curbuf->b_u_curhead == NULL) {
+    // nothing was undone.
+    return false;
+  }
+
+  // Delete the current redo header
+  // set the redo header to the next alternative branch (if any)
+  // otherwise we will be in the leaf state
+  u_header_T *to_forget = curbuf->b_u_curhead;
+  curbuf->b_u_newhead = to_forget->uh_next.ptr;
+  curbuf->b_u_curhead = to_forget->uh_alt_next.ptr;
+  if (curbuf->b_u_curhead) {
+    to_forget->uh_alt_next.ptr = NULL;
+    curbuf->b_u_curhead->uh_alt_prev.ptr = to_forget->uh_alt_prev.ptr;
+    curbuf->b_u_seq_cur = curbuf->b_u_curhead->uh_next.ptr ?
+        curbuf->b_u_curhead->uh_next.ptr->uh_seq : 0;
+  } else if (curbuf->b_u_newhead) {
+    curbuf->b_u_seq_cur = curbuf->b_u_newhead->uh_seq;
+  }
+  if (to_forget->uh_alt_prev.ptr) {
+    to_forget->uh_alt_prev.ptr->uh_alt_next.ptr = curbuf->b_u_curhead;
+  }
+  if (curbuf->b_u_newhead) {
+    curbuf->b_u_newhead->uh_prev.ptr = curbuf->b_u_curhead;
+  }
+  if (curbuf->b_u_seq_last == to_forget->uh_seq) {
+    curbuf->b_u_seq_last--;
+  }
+  u_freebranch(curbuf, to_forget, NULL);
+  return true;
+}
+
+/// Undo or redo, depending on `undo_undoes`, `count` times.
+static void u_doit(int startcount, bool quiet)
 {
   int count = startcount;
 
@@ -1721,7 +1765,7 @@ static void u_doit(int startcount)
         break;
       }
 
-      u_undoredo(TRUE);
+      u_undoredo(true);
     } else {
       if (curbuf->b_u_curhead == NULL || get_undolevel() <= 0) {
         beep_flush();           /* nothing to redo */
@@ -1741,7 +1785,7 @@ static void u_doit(int startcount)
       curbuf->b_u_curhead = curbuf->b_u_curhead->uh_prev.ptr;
     }
   }
-  u_undo_end(undo_undoes, FALSE);
+  u_undo_end(undo_undoes, false, quiet);
 }
 
 /*
@@ -2054,7 +2098,7 @@ void undo_time(long step, int sec, int file, int absolute)
       }
     }
   }
-  u_undo_end(did_undo, absolute);
+  u_undo_end(did_undo, absolute, false);
 }
 
 /*
@@ -2278,7 +2322,8 @@ static void u_undoredo(int undo)
   if (undo)
     /* We are below the previous undo.  However, to make ":earlier 1s"
      * work we compute this as being just above the just undone change. */
-    --curbuf->b_u_seq_cur;
+    curbuf->b_u_seq_cur = curhead->uh_next.ptr ?
+        curhead->uh_next.ptr->uh_seq : 0;
 
   /* Remember where we are for ":earlier 1f" and ":later 1f". */
   if (curhead->uh_save_nr != 0) {
@@ -2298,16 +2343,13 @@ static void u_undoredo(int undo)
 #endif
 }
 
-/*
- * If we deleted or added lines, report the number of less/more lines.
- * Otherwise, report the number of changes (this may be incorrect
- * in some cases, but it's better than nothing).
- */
-static void 
-u_undo_end (
-    int did_undo,                   /* just did an undo */
-    int absolute                   /* used ":undo N" */
-)
+/// If we deleted or added lines, report the number of less/more lines.
+/// Otherwise, report the number of changes (this may be incorrect
+/// in some cases, but it's better than nothing).
+static void u_undo_end(
+    int did_undo,     ///< just did an undo
+    int absolute,     ///< used ":undo N"
+    bool quiet)
 {
   char        *msgstr;
   u_header_T  *uhp;
@@ -2316,9 +2358,11 @@ u_undo_end (
   if ((fdo_flags & FDO_UNDO) && KeyTyped)
     foldOpenCursor();
 
-  if (global_busy           /* no messages now, wait until global is finished */
-      || !messaging())        /* 'lazyredraw' set, don't do messages now */
+  if (quiet
+      || global_busy        // no messages until global is finished
+      || !messaging()) {    // 'lazyredraw' set, don't do messages now
     return;
+  }
 
   if (curbuf->b_ml.ml_flags & ML_EMPTY)
     --u_newcount;
@@ -2507,20 +2551,20 @@ static void u_add_time(char_u *buf, size_t buflen, time_t tt)
  */
 void ex_undojoin(exarg_T *eap)
 {
-  if (curbuf->b_u_newhead == NULL)
-    return;                 /* nothing changed before */
+  if (curbuf->b_u_newhead == NULL) {
+    return;                 // nothing changed before
+  }
   if (curbuf->b_u_curhead != NULL) {
     EMSG(_("E790: undojoin is not allowed after undo"));
     return;
   }
-  if (!curbuf->b_u_synced)
-    return;                 /* already unsynced */
-  if (get_undolevel() < 0)
-    return;                 /* no entries, nothing to do */
-  else {
-    /* Go back to the last entry */
-    curbuf->b_u_curhead = curbuf->b_u_newhead;
-    curbuf->b_u_synced = false;      /* no entries, nothing to do */
+  if (!curbuf->b_u_synced) {
+    return;                 // already unsynced
+  }
+  if (get_undolevel() < 0) {
+    return;                 // no entries, nothing to do
+  } else {
+    curbuf->b_u_synced = false;  // Append next change to last entry
   }
 }