1 files changed, 48 insertions, 13 deletions

diff --git a/src/nvim/ex_docmd.c b/src/nvim/ex_docmd.c
index 085948f7b1..987fb439d7 100644
--- a/src/nvim/ex_docmd.c
+++ b/src/nvim/ex_docmd.c
@@ -1339,6 +1339,20 @@ static int parse_count(exarg_T *eap, char **errormsg, bool validate)
           || ascii_iswhite(*p))) {
     long n = getdigits_long(&eap->arg, false, -1);
     eap->arg = skipwhite(eap->arg);
+
+    if (eap->args != NULL) {
+      assert(eap->argc > 0 && eap->arg >= eap->args[0]);
+      // If eap->arg is still pointing to the first argument, just make eap->args[0] point to the
+      // same location. This is needed for usecases like vim.cmd.sleep('10m'). If eap->arg is
+      // pointing outside the first argument, shift arguments by 1.
+      if (eap->arg < eap->args[0] + eap->arglens[0]) {
+        eap->arglens[0] -= (size_t)(eap->arg - eap->args[0]);
+        eap->args[0] = eap->arg;
+      } else {
+        shift_cmd_args(eap);
+      }
+    }
+
     if (n <= 0 && (eap->argt & EX_ZEROR) == 0) {
       if (errormsg != NULL) {
         *errormsg = _(e_zerocount);
@@ -1512,6 +1526,30 @@ end:
   return retval;
 }
 
+// Shift Ex-command arguments to the right.
+static void shift_cmd_args(exarg_T *eap)
+{
+  assert(eap->args != NULL && eap->argc > 0);
+
+  char **oldargs = eap->args;
+  size_t *oldarglens = eap->arglens;
+
+  eap->argc--;
+  eap->args = eap->argc > 0 ? xcalloc(eap->argc, sizeof(char *)) : NULL;
+  eap->arglens = eap->argc > 0 ? xcalloc(eap->argc, sizeof(size_t)) : NULL;
+
+  for (size_t i = 0; i < eap->argc; i++) {
+    eap->args[i] = oldargs[i + 1];
+    eap->arglens[i] = oldarglens[i + 1];
+  }
+
+  // If there are no arguments, make eap->arg point to the end of string.
+  eap->arg = (eap->argc > 0 ? eap->args[0] : (oldargs[0] + oldarglens[0]));
+
+  xfree(oldargs);
+  xfree(oldarglens);
+}
+
 static int execute_cmd0(int *retv, exarg_T *eap, char **errormsg, bool preview)
 {
   // If filename expansion is enabled, expand filenames
@@ -1551,15 +1589,7 @@ static int execute_cmd0(int *retv, exarg_T *eap, char **errormsg, bool preview)
                                    eap->args[0] + eap->arglens[0],
                                    (eap->argt & EX_BUFUNL) != 0, false, false);
       eap->addr_count = 1;
-      // Shift each argument by 1
-      for (size_t i = 0; i < eap->argc - 1; i++) {
-        eap->args[i] = eap->args[i + 1];
-      }
-      // Make the last argument point to the NUL terminator at the end of string
-      eap->args[eap->argc - 1] = eap->args[eap->argc - 1] + eap->arglens[eap->argc - 1];
-      eap->argc -= 1;
-
-      eap->arg = eap->args[0];
+      shift_cmd_args(eap);
     }
     if (eap->line2 < 0) {  // failed
       return FAIL;
@@ -1658,6 +1688,11 @@ int execute_cmd(exarg_T *eap, CmdParseInfo *cmdinfo, bool preview)
     (void)hasFolding(eap->line2, NULL, &eap->line2);
   }
 
+  // Use first argument as count when possible
+  if (parse_count(eap, &errormsg, true) == FAIL) {
+    goto end;
+  }
+
   // Execute the command
   execute_cmd0(&retv, eap, &errormsg, preview);
 
@@ -3887,7 +3922,7 @@ static char *repl_cmdline(exarg_T *eap, char *src, size_t srclen, char *repl, ch
     } else {
       // Otherwise, argument gets shifted alongside the replaced text.
       // The amount of the shift is equal to the difference of the old and new string length.
-      eap->args[j] = new_cmdline + (eap->args[j] - *cmdlinep) + (len - srclen);
+      eap->args[j] = new_cmdline + ((eap->args[j] - *cmdlinep) + (ptrdiff_t)(len - srclen));
     }
   }
 
@@ -5867,7 +5902,7 @@ static void ex_undo(exarg_T *eap)
 {
   if (eap->addr_count != 1) {
     if (eap->forceit) {
-      u_undo_and_forget(1);         // :undo!
+      u_undo_and_forget(1, true);   // :undo!
     } else {
       u_undo(1);                    // :undo
     }
@@ -5894,7 +5929,7 @@ static void ex_undo(exarg_T *eap)
       emsg(_(e_undobang_cannot_redo_or_move_branch));
       return;
     }
-    u_undo_and_forget(count);
+    u_undo_and_forget(count, true);
   } else {                        // :undo 123
     undo_time(step, false, false, true);
   }
@@ -6088,7 +6123,7 @@ static void ex_redrawstatus(exarg_T *eap)
   } else {
     status_redraw_curbuf();
   }
-  if (msg_scrolled && (State & MODE_CMDLINE)) {
+  if (msg_scrolled && !msg_use_msgsep() && (State & MODE_CMDLINE)) {
     return;  // redraw later
   }