4 files changed, 19 insertions, 12 deletions
diff --git a/Src/Zle/zle_main.c b/Src/Zle/zle_main.c
index a2b20df25..f5aec8403 100644
--- a/Src/Zle/zle_main.c
+++ b/Src/Zle/zle_main.c
@@ -157,10 +157,10 @@ mod_export char *statusline;
 /**/
 int stackhist, stackcs;
 
-/* != 0 if we are making undo records */
+/* position in undo stack from when the current vi change started */
 
 /**/
-int undoing;
+zlong vistartchange;
 
 /* current modifier status */
 
@@ -1080,8 +1080,7 @@ zlecore(void)
 	    if (invicmdmode() && zlecs > findbol() &&
 		(zlecs == zlell || zleline[zlecs] == ZWC('\n')))
 		DECCS();
-	    if (undoing)
-		handleundo();
+	    handleundo();
 	} else {
 	    errflag = 1;
 	    break;
@@ -1190,7 +1189,7 @@ zleread(char **lp, char **rp, int flags, int context, char *init, char *finish)
     zlereadflags = flags;
     zlecontext = context;
     histline = curhist;
-    undoing = 1;
+    vistartchange = 0;
     zleline = (ZLE_STRING_T)zalloc(((linesz = 256) + 2) * ZLE_CHAR_SIZE);
     *zleline = ZWC('\0');
     virangeflag = lastcmd = done = zlecs = zlell = mark = 0;
@@ -1198,6 +1197,7 @@ zleread(char **lp, char **rp, int flags, int context, char *init, char *finish)
     viinsbegin = 0;
     statusline = NULL;
     selectkeymap("main", 1);
+    initundo();
     /*
      * If main is linked to the viins keymap, we need to register
      * explicitly that we're now in vi insert mode as there's
@@ -1222,7 +1222,6 @@ zleread(char **lp, char **rp, int flags, int context, char *init, char *finish)
 	    stackhist = -1;
 	}
     }
-    initundo();
     if (isset(PROMPTCR))
 	putc('\r', shout);
     if (tmout)
diff --git a/Src/Zle/zle_tricky.c b/Src/Zle/zle_tricky.c
index 9d163ad9e..2689d0fa1 100644
--- a/Src/Zle/zle_tricky.c
+++ b/Src/Zle/zle_tricky.c
@@ -611,8 +611,7 @@ docomplete(int lst)
     active = 1;
     comprecursive = 0;
     makecommaspecial(0);
-    if (undoing)
-	setlastline();
+    setlastline();
 
     /* From the C-code's point of view, we can only use compctl as a default
      * type of completion. Load it if it hasn't been loaded already and
diff --git a/Src/Zle/zle_utils.c b/Src/Zle/zle_utils.c
index b82e54ce5..61ae85c0b 100644
--- a/Src/Zle/zle_utils.c
+++ b/Src/Zle/zle_utils.c
@@ -1354,7 +1354,10 @@ handlesuffix(UNUSED(char **args))
 
 /* head of the undo list, and the current position */
 
-static struct change *changes, *curchange;
+/**/
+struct change *curchange;
+
+static struct change *changes;
 
 /* list of pending changes, not yet in the undo system */
 
diff --git a/Src/Zle/zle_vi.c b/Src/Zle/zle_vi.c
index 31f293387..9e9cc2f34 100644
--- a/Src/Zle/zle_vi.c
+++ b/Src/Zle/zle_vi.c
@@ -107,7 +107,7 @@ startvitext(int im)
 {
     startvichange(im);
     selectkeymap("main", 1);
-    undoing = 0;
+    vistartchange = (curchange && curchange->prev) ? curchange->prev->changeno : 0;
     viinsbegin = zlecs;
 }
 
@@ -399,7 +399,7 @@ vichange(UNUSED(char **args))
 	forekill(c2 - zlecs, CUT_RAW);
 	selectkeymap("main", 1);
 	viinsbegin = zlecs;
-	undoing = 0;
+	vistartchange = curchange->prev->changeno;
     }
     return ret;
 }
@@ -584,7 +584,13 @@ vicmdmode(UNUSED(char **args))
 {
     if (invicmdmode() || selectkeymap("vicmd", 0))
 	return 1;
-    undoing = 1;
+    struct change *current = curchange->prev;
+    while (current && current->changeno > vistartchange+1) {
+	current->flags |= CH_PREV;
+	current = current->prev;
+	if (!current) break;
+	current->flags |= CH_NEXT;
+    }
     vichgflag = 0;
     if (zlecs != findbol())
 	DECCS();