39986, 39989: improve handling of vi-repeat-change

Save previous vi change and throw away a new change that fails.
Add zle -f vichange to allow shell widget to be a single change.
Fix repeat of command where numeric arguments were multiplied.

7 files changed, 123 insertions, 84 deletions

diff --git a/Src/Zle/zle.h b/Src/Zle/zle.h
index e9b14281d..8f92e5611 100644
--- a/Src/Zle/zle.h
+++ b/Src/Zle/zle.h
@@ -284,6 +284,20 @@ struct change {
 #define CH_NEXT (1<<0)   /* next structure is also part of this change */
 #define CH_PREV (1<<1)   /* previous structure is also part of this change */
 
+/* vi change handling for vi-repeat-change */
+
+/*
+ * Examination of the code suggests vichgbuf is consistently tied
+ * to raw byte input, so it is left as a character array rather
+ * than turned into wide characters.  In particular, when we replay
+ * it we use ungetbytes().
+ */
+struct vichange {
+    struct modifier mod; /* value of zmod associated with vi change */
+    char *buf;           /* bytes for keys that make up the vi command */
+    int bufsz, bufptr;   /* allocated and in use sizes of buf */
+};
+
 /* known thingies */
 
 #define Th(X) (&thingies[X])
diff --git a/Src/Zle/zle_keymap.c b/Src/Zle/zle_keymap.c
index 24e8d1964..04eb70675 100644
--- a/Src/Zle/zle_keymap.c
+++ b/Src/Zle/zle_keymap.c
@@ -1634,7 +1634,7 @@ getkeymapcmd(Keymap km, Thingy *funcp, char **strp)
 	unmetafy(keybuf + lastlen, &keybuflen);
 	ungetbytes(keybuf+lastlen, keybuflen);
 	if(vichgflag)
-	    vichgbufptr -= keybuflen;
+	    curvichg.bufptr -= keybuflen;
 	keybuf[keybuflen = lastlen] = 0;
     }
     *funcp = func;
diff --git a/Src/Zle/zle_main.c b/Src/Zle/zle_main.c
index 1652b7cd4..938dc0e29 100644
--- a/Src/Zle/zle_main.c
+++ b/Src/Zle/zle_main.c
@@ -924,13 +924,13 @@ getbyte(long do_keytmout, int *timeout)
 	ret = STOUC(cc);
     }
     /*
-     * vichgbuf is raw bytes, not wide characters, so is dealt
+     * curvichg.buf is raw bytes, not wide characters, so is dealt
      * with here.
      */
     if (vichgflag) {
-	if (vichgbufptr == vichgbufsz)
-	    vichgbuf = realloc(vichgbuf, vichgbufsz *= 2);
-	vichgbuf[vichgbufptr++] = ret;
+	if (curvichg.bufptr == curvichg.bufsz)
+	    curvichg.buf = realloc(curvichg.buf, curvichg.bufsz *= 2);
+	curvichg.buf[curvichg.bufptr++] = ret;
     }
     errno = old_errno;
     return lastchar = ret;
@@ -1260,6 +1260,7 @@ zleread(char **lp, char **rp, int flags, int context, char *init, char *finish)
     *zleline = ZWC('\0');
     virangeflag = lastcmd = done = zlecs = zlell = mark = yankb = yanke = 0;
     vichgflag = 0;
+    viinrepeat = 0;
     viinsbegin = 0;
     statusline = NULL;
     selectkeymap("main", 1);
@@ -1389,6 +1390,8 @@ int
 execzlefunc(Thingy func, char **args, int set_bindk)
 {
     int r = 0, ret = 0, remetafy = 0;
+    int nestedvichg = vichgflag;
+    int isrepeat = (viinrepeat == 3);
     Widget w;
     Thingy save_bindk = bindk;
 
@@ -1398,6 +1401,8 @@ execzlefunc(Thingy func, char **args, int set_bindk)
 	unmetafy_line();
 	remetafy = 1;
     }
+    if (isrepeat)
+	viinrepeat = 2;
 
     if (func->flags & DISABLED) {
 	/* this thingy is not the name of a widget */
@@ -1523,6 +1528,25 @@ execzlefunc(Thingy func, char **args, int set_bindk)
     CCRIGHT();
     if (remetafy)
 	metafy_line();
+
+    /* if this widget constituted the vi change, end it */
+    if (vichgflag == 2 && !nestedvichg) {
+	if (invicmdmode()) {
+	    if (ret) {
+		free(curvichg.buf);
+	    } else {
+		if (lastvichg.buf)
+		    free(lastvichg.buf);
+		lastvichg = curvichg;
+	    }
+	    vichgflag = 0;
+	    curvichg.buf = NULL;
+	} else
+	    vichgflag = 1; /* vi change continues while in insert mode */
+    }
+    if (isrepeat)
+        viinrepeat = !invicmdmode();
+
     return ret;
 }
 
@@ -2230,7 +2254,7 @@ finish_(UNUSED(Module m))
     cleanup_keymaps();
     deletehashtable(thingytab);
 
-    zfree(vichgbuf, vichgbufsz);
+    zfree(lastvichg.buf, lastvichg.bufsz);
     zfree(kungetbuf, kungetsz);
     free_isrch_spots();
     if (rdstrs)
diff --git a/Src/Zle/zle_misc.c b/Src/Zle/zle_misc.c
index fbd40cd03..898b552de 100644
--- a/Src/Zle/zle_misc.c
+++ b/Src/Zle/zle_misc.c
@@ -609,8 +609,10 @@ viputbefore(UNUSED(char **args))
     int n = zmult;
 
     startvichange(-1);
-    if (n < 0 || zmod.flags & MOD_NULL)
+    if (n < 0)
 	return 1;
+    if (zmod.flags & MOD_NULL)
+	return 0;
     if (zmod.flags & MOD_VIBUF)
 	kctbuf = &vibuf[zmod.vibuf];
     else
@@ -630,8 +632,10 @@ viputafter(UNUSED(char **args))
     int n = zmult;
 
     startvichange(-1);
-    if (n < 0 || zmod.flags & MOD_NULL)
+    if (n < 0)
 	return 1;
+    if (zmod.flags & MOD_NULL)
+	return 0;
     if (zmod.flags & MOD_VIBUF)
 	kctbuf = &vibuf[zmod.vibuf];
     else
diff --git a/Src/Zle/zle_thingy.c b/Src/Zle/zle_thingy.c
index 2d4674785..c7092854a 100644
--- a/Src/Zle/zle_thingy.c
+++ b/Src/Zle/zle_thingy.c
@@ -678,7 +678,16 @@ bin_zle_flags(char *name, char **args, UNUSED(Options ops), UNUSED(char func))
 		else if (!strcmp(*flag, "keepsuffix"))
 		    w->flags |= ZLE_KEEPSUFFIX;
 		*/
-		else {
+	        else if (!strcmp(*flag, "vichange")) {
+		    if (invicmdmode()) {
+			startvichange(-1);
+			if (zmod.flags & (MOD_MULT|MOD_TMULT)) {
+			    Param pm = (Param) paramtab->getnode(paramtab, "NUMERIC");
+			    if (pm && pm->node.flags & PM_SPECIAL)
+				pm->node.flags &= ~PM_UNSET;
+			}
+		    }
+		} else {
 		    zwarnnam(name, "invalid flag `%s' given to zle -f", *flag);
 		    ret = 1;
 		}
diff --git a/Src/Zle/zle_utils.c b/Src/Zle/zle_utils.c
index 4007c1112..c6df3d89c 100644
--- a/Src/Zle/zle_utils.c
+++ b/Src/Zle/zle_utils.c
@@ -1704,6 +1704,7 @@ mergeundo(void)
 	current->flags |= CH_PREV;
 	current->prev->flags |= CH_NEXT;
     }
+    vistartchange = -1;
 }
 
 /*
diff --git a/Src/Zle/zle_vi.c b/Src/Zle/zle_vi.c
index fc0e49b32..e0923db3e 100644
--- a/Src/Zle/zle_vi.c
+++ b/Src/Zle/zle_vi.c
@@ -45,46 +45,41 @@ int wordflag;
 /**/
 int vilinerange;
 
-/* last vi change buffer, for vi change repetition */
+/*
+ * lastvichg: last vi change buffer, for vi change repetition
+ * curvichg: current incomplete vi change
+ */
+
+/**/
+struct vichange lastvichg, curvichg;
 
 /*
- * vichgbufsz: Allocated size of vichgbuf.
- * vichgbufptr: Length in use.
- * vichgflag: true whilst inputting a vi normal mode; causes it to be
- *   accumulated in vichgbuf, incrementing vichgbufptr.
+ * true whilst a vi change is active causing keys to be
+ * accumulated in curvichg.buf
+ * first set to 2 and when the initial widget finishes, reduced to 1 if
+ * in insert mode implying that the change continues until returning to
+ * normal mode
  */
 
 /**/
-int vichgbufsz, vichgbufptr, vichgflag;
+int vichgflag;
 
 /*
- * The bytes that make up the current vi command.  See vichgbuf* above.
- *
- * Examination of the code suggests vichgbuf is consistently tied
- * to raw byte input, so it is left as a character array rather
- * than turned into wide characters.  In particular, when we replay
- * it we use ungetbytes().
+ * analogous to vichgflag for a repeated change with the value following
+ * a similar pattern (is 3 until first repeated widget starts)
  */
+
 /**/
-char *vichgbuf;
+int viinrepeat;
 
 /* point where vi insert mode was last entered */
 
 /**/
 int viinsbegin;
 
-/* value of zmod associated with vi change */
-static struct modifier lastmod;
-
-/*
- * inrepeat: current widget is the vi change being repeated
- * vichgrepeat: nested widget call within a repeat
- */
-static int inrepeat, vichgrepeat;
-
 /**
  * im: >= 0: is an insertmode
- *    -1: skip setting insert mode
+ *    -1: skip setting insert/overwrite mode
  *    -2: entering viins at start of editing from clean --- don't use
  *        inrepeat or keybuf, synthesise an entry to insert mode.
  * Note that zmult is updated so this should be called before zmult is used.
@@ -94,30 +89,27 @@ static int inrepeat, vichgrepeat;
 void
 startvichange(int im)
 {
-    if (im != -1) {
-	vichgflag = 1;
-	if (im > -1)
-	    insmode = im;
-    }
-    if (inrepeat && im != -2) {
-	zmod = lastmod;
-	inrepeat = vichgflag = 0;
-	vichgrepeat = 1;
-    } else {
-	lastmod = zmod;
-	if (vichgbuf)
-	    free(vichgbuf);
-	vichgbuf = (char *)zalloc(vichgbufsz = 16 + keybuflen);
+    if (im > -1)
+	insmode = im;
+    if (viinrepeat && im != -2) {
+	zmod = lastvichg.mod;
+	vichgflag = 0;
+    } else if (!vichgflag) {
+	curvichg.mod = zmod;
+	if (curvichg.buf)
+	    free(curvichg.buf);
+	curvichg.buf = (char *)zalloc(curvichg.bufsz = 16 + keybuflen);
 	if (im == -2) {
-	    vichgbuf[0] =
+	    vichgflag = 1;
+	    curvichg.buf[0] =
 		zlell ? (insmode ? (zlecs < zlell ? 'i' : 'a') : 'R') : 'o';
-	    vichgbuf[1] = '\0';
-	    vichgbufptr = 1;
+	    curvichg.buf[1] = '\0';
+	    curvichg.bufptr = 1;
 	} else {
-	    strcpy(vichgbuf, keybuf);
-	    unmetafy(vichgbuf, &vichgbufptr);
+	    vichgflag = 2;
+	    strcpy(curvichg.buf, keybuf);
+	    unmetafy(curvichg.buf, &curvichg.bufptr);
 	}
-	vichgrepeat = 0;
     }
 }
 
@@ -226,10 +218,13 @@ getvirange(int wf)
 	     */
 	    if ((k2 == bindk) ? dovilinerange() : execzlefunc(k2, zlenoargs, 1))
 		ret = -1;
-	    if(vichgrepeat)
+	    if (viinrepeat)
 		zmult = mult1;
-	    else
+	    else {
 		zmult = mult1 * zmod.tmult;
+		if (vichgflag == 2)
+		    curvichg.mod.mult = zmult;
+            }
 	} while(prefixflag && !ret);
 	wordflag = 0;
 	selectlocalmap(NULL);
@@ -401,7 +396,6 @@ videlete(UNUSED(char **args))
 	    vifirstnonblank(zlenoargs);
 	}
     }
-    vichgflag = 0;
     return ret;
 }
 
@@ -518,7 +512,6 @@ viyank(UNUSED(char **args))
 	cut(zlecs, c2 - zlecs, CUT_YANK);
 	ret = 0;
     }
-    vichgflag = 0;
     /* cursor now at the start of the range yanked. For line mode
      * restore the column position */
     if (vilinerange && lastcol != -1) {
@@ -593,8 +586,7 @@ vireplace(UNUSED(char **args))
  * a change, we always read the argument normally, even if the count    *
  * was bad.  When recording a change for repeating, and a bad count is  *
  * given, we squash the repeat buffer to avoid repeating the partial    *
- * command; we've lost the previous change, but that can't be avoided   *
- * without a rewrite of the repeat code.                                */
+ * command.                                                             */
 
 /**/
 int
@@ -644,18 +636,12 @@ vireplacechars(UNUSED(char **args))
 
     /* check argument range */
     if (n < 1 || fail) {
-	if(vichgrepeat)
+	if (viinrepeat)
 	    vigetkey();
-	if(vichgflag) {
-	    free(vichgbuf);
-	    vichgbuf = NULL;
-	    vichgflag = 0;
-	}
 	return 1;
     }
     /* get key */
     if((ch = vigetkey()) == ZLEEOF) {
-	vichgflag = 0;
 	return 1;
     }
     /* do change */
@@ -682,7 +668,6 @@ vireplacechars(UNUSED(char **args))
 	    zleline[zlecs++] = ch;
 	zlecs--;
     }
-    vichgflag = 0;
     return 0;
 }
 
@@ -693,7 +678,16 @@ vicmdmode(UNUSED(char **args))
     if (invicmdmode() || selectkeymap("vicmd", 0))
 	return 1;
     mergeundo();
-    vichgflag = 0;
+    insmode = unset(OVERSTRIKE);
+    if (vichgflag == 1) {
+	vichgflag = 0;
+	if (lastvichg.buf)
+	    free(lastvichg.buf);
+	lastvichg = curvichg;
+	curvichg.buf = NULL;
+    }
+    if (viinrepeat == 1)
+        viinrepeat = 0;
     if (zlecs != findbol())
 	DECCS();
     return 0;
@@ -748,7 +742,6 @@ vioperswapcase(UNUSED(char **args))
 	vifirstnonblank();
 #endif
     }
-    vichgflag = 0;
     return ret;
 }
 
@@ -771,7 +764,6 @@ viupcase(UNUSED(char **args))
 	zlecs = oldcs;
 	ret = 0;
     }
-    vichgflag = 0;
     return ret;
 }
 
@@ -794,7 +786,6 @@ vidowncase(UNUSED(char **args))
 	zlecs = oldcs;
 	ret = 0;
     }
-    vichgflag = 0;
     return ret;
 }
 
@@ -803,23 +794,23 @@ int
 virepeatchange(UNUSED(char **args))
 {
     /* make sure we have a change to repeat */
-    if (!vichgbuf || vichgflag || virangeflag)
+    if (!lastvichg.buf || vichgflag || virangeflag)
 	return 1;
     /* restore or update the saved count and buffer */
     if (zmod.flags & MOD_MULT) {
-	lastmod.mult = zmod.mult;
-	lastmod.flags |= MOD_MULT;
+	lastvichg.mod.mult = zmod.mult;
+	lastvichg.mod.flags |= MOD_MULT;
     }
     if (zmod.flags & MOD_VIBUF) {
-	lastmod.vibuf = zmod.vibuf;
-	lastmod.flags = (lastmod.flags & ~MOD_VIAPP) |
+	lastvichg.mod.vibuf = zmod.vibuf;
+	lastvichg.mod.flags = (lastvichg.mod.flags & ~MOD_VIAPP) |
 	    MOD_VIBUF | (zmod.flags & MOD_VIAPP);
-    } else if (lastmod.flags & MOD_VIBUF &&
-	    lastmod.vibuf >= 27 && lastmod.vibuf <= 34)
-	lastmod.vibuf++; /* for "1 to "8 advance to next buffer */
+    } else if (lastvichg.mod.flags & MOD_VIBUF &&
+	    lastvichg.mod.vibuf >= 27 && lastvichg.mod.vibuf <= 34)
+	lastvichg.mod.vibuf++; /* for "1 to "8 advance to next buffer */
     /* repeat the command */
-    inrepeat = 1;
-    ungetbytes(vichgbuf, vichgbufptr);
+    viinrepeat = 3;
+    ungetbytes(lastvichg.buf, lastvichg.bufptr);
     return 0;
 }
 
@@ -835,10 +826,8 @@ viindent(UNUSED(char **args))
 	region_active = 2;
     /* get the range */
     if ((c2 = getvirange(0)) == -1) {
-	vichgflag = 0;
 	return 1;
     }
-    vichgflag = 0;
     /* must be a line range */
     if (!vilinerange) {
 	zlecs = oldcs;
@@ -873,10 +862,8 @@ viunindent(UNUSED(char **args))
 	region_active = 2;
     /* get the range */
     if ((c2 = getvirange(0)) == -1) {
-	vichgflag = 0;
 	return 1;
     }
-    vichgflag = 0;
     /* must be a line range */
     if (!vilinerange) {
 	zlecs = oldcs;