PATCH: vi repeats

[Oliver Kiddle <okiddle@xxxxxxxxxxx>]

This patch is to the handling of vi-repeat-change.

Firstly, the previous change is saved while the next change is
constructed. If the widget for the new change returns failure, the new
change is thrown away so that the previous change is not lost. I've used
a struct to collect together the variables relevant to a repeat.

Secondly, it is now possible to have repeat work for a user-defined
widget. By using zle -f vichange. vichange is not as such one of
the ZLE_ flags but it could easily be implemented that way if we
wanted. The code for tracking changes (both the initial case and for a
repeat) is now more general. Previously, we didn't need to worry about
startvichange being called more than once for a single change so we got
away with more rudimentary code for that.

There's also some more minor fixes such as the repeat of a command where
numeric arguments were multiplied, such as 3d2w which deletes 6 words.

Oliver

diff --git a/Doc/Zsh/zle.yo b/Doc/Zsh/zle.yo
index fe6a7e9..ff31448 100644
--- a/Doc/Zsh/zle.yo
+++ b/Doc/Zsh/zle.yo
@@ -479,6 +479,13 @@ tt(kill) for indicating that text has been killed into the cutbuffer.
 When repeatedly invoking a kill widget, text is appended to the cutbuffer
 instead of replacing it, but when wrapping such widgets, it is necessary
 to call `tt(zle -f kill)' to retain this effect.
+
+tt(vichange) for indicating that the widget represents a vi change that
+can be repeated as a whole with `tt(vi-repeat-change)'. The flag should be set
+early in the function before inspecting the value of tt(NUMERIC) or invoking
+other widgets. This has no effect for a widget invoked from insert mode. If
+insert mode is active when the widget finishes, the change extends until next
+returning to command mode.
 )
 cindex(completion widgets, creating)
 item(tt(-C) var(widget) var(completion-widget) var(function))(
diff --git a/Functions/Zle/vi-pipe b/Functions/Zle/vi-pipe
index c32e64a..1729cb6 100644
--- a/Functions/Zle/vi-pipe
+++ b/Functions/Zle/vi-pipe
@@ -12,6 +12,9 @@ setopt localoptions noksharrays
 autoload -Uz read-from-minibuffer
 local _save_cut="$CUTBUFFER" REPLY
 
+# mark this widget as a vi change so it can be repeated as a whole
+zle -f vichange
+
 # force movement to default to line mode
 (( REGION_ACTIVE )) || zle -U V
 # Use the standard vi-change to accept a vi motion.
diff --git a/Src/Zle/zle.h b/Src/Zle/zle.h
index e9b1428..8f92e56 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 24e8d19..04eb706 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 1652b7c..938dc0e 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 fbd40cd..898b552 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 2d46747..8a6904c 100644
--- a/Src/Zle/zle_thingy.c
+++ b/Src/Zle/zle_thingy.c
@@ -678,7 +678,10 @@ 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);
+		} 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 4007c11..c6df3d8 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 fc0e49b..e0923db 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 */
-
 /*
- * 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.
+ * lastvichg: last vi change buffer, for vi change repetition
+ * curvichg: current incomplete vi change
  */
 
 /**/
-int vichgbufsz, vichgbufptr, vichgflag;
+struct vichange lastvichg, curvichg;
 
 /*
- * 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().
+ * 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
  */
+
 /**/
-char *vichgbuf;
+int vichgflag;
+
+/*
+ * analogous to vichgflag for a repeated change with the value following
+ * a similar pattern (is 3 until first repeated widget starts)
+ */
+
+/**/
+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;
diff --git a/Test/X02zlevi.ztst b/Test/X02zlevi.ztst
index c195732..84df7c6 100644
--- a/Test/X02zlevi.ztst
+++ b/Test/X02zlevi.ztst
@@ -140,6 +140,12 @@
 >BUFFER: xxe xxx xxxee
 >CURSOR: 10
 
+  zletest $'one two three four five six seven eight\e.03d2wk.1.'
+0:with numeric args to both action and movement, they are multiplied
+>BUFFER: eight
+>seven eight
+>CURSOR: 0
+
   zletest $'yankee doodle\ebhDyy0"1P'
 0:paste register 1 to get last deletion
 >BUFFER:  doodleyankee
@@ -262,10 +268,8 @@
   print -u $ZTST_fd 'This test may hang the shell when it fails...'
   zletest $'worm\erdhd..'
 0:use of vi-repeat as the motion and repeat after a failed change
->BUFFER: word
+>BUFFER: wodd
 >CURSOR: 2
-F:For vi compatibility, "." should repeat the "rd" change after "d."
-F:Update result to ">BUFFER: wodd" if compatibility is repaired.
 
   zpty_run 'bindkey "^_" undo'
   zletest $'undoc\037e'