Author: Chris Noxz <chris@noxz.tech>

This patch adds the feature to visually indicate the current vi mode in the
prompt. It's required to add a placeholder in your PS1 such as '[xxx] '. The
indicator will be inserted between the first '[' and the first ']'. All status
indicators needs to be of the same visible size, such as 'cmd','ins' and 'rep'.

Happy hacking

diff --git a/vi.c b/vi.c
index 12202d5..bd96137 100644
--- a/vi.c
+++ b/vi.c
@@ -28,6 +28,10 @@ struct edstate {
 	int	cursor;		/* byte# in cbuf having the cursor */
 };
 
+static const char vprompt_normal[] = "\\[\x1b[1;31m\\]cmd\\[\x1b[0m\\]";
+static const char vprompt_insert[] = "\\[\x1b[1;32m\\]ins\\[\x1b[0m\\]";
+static const char vprompt_replace[] = "\\[\x1b[1;33m\\]rep\\[\x1b[0m\\]";
+static void set_insert(int);
 
 static int	vi_hook(int);
 static void	vi_reset(char *, size_t);
@@ -754,21 +758,21 @@ vi_cmd(int argcnt, const char *cmd)
 			if (es->linelen != 0)
 				while (isu8cont(es->cbuf[++es->cursor]))
 					continue;
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'A':
 			modified = 1; hnum = hlast;
 			del_range(0, 0);
 			es->cursor = es->linelen;
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'S':
 			es->cursor = domove(1, "^", 1);
 			del_range(es->cursor, es->linelen);
 			modified = 1; hnum = hlast;
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'Y':
@@ -811,7 +815,7 @@ vi_cmd(int argcnt, const char *cmd)
 			}
 			if (*cmd == 'c') {
 				modified = 1; hnum = hlast;
-				insert = INSERT;
+				set_insert(INSERT);
 			}
 			break;
 
@@ -841,7 +845,7 @@ vi_cmd(int argcnt, const char *cmd)
 		case 'C':
 			modified = 1; hnum = hlast;
 			del_range(es->cursor, es->linelen);
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'D':
@@ -870,13 +874,13 @@ vi_cmd(int argcnt, const char *cmd)
 
 		case 'i':
 			modified = 1; hnum = hlast;
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'I':
 			modified = 1; hnum = hlast;
 			es->cursor = domove(1, "^", 1);
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'j':
@@ -931,7 +935,7 @@ vi_cmd(int argcnt, const char *cmd)
 
 		case 'R':
 			modified = 1; hnum = hlast;
-			insert = REPLACE;
+			set_insert(REPLACE);
 			break;
 
 		case 's':
@@ -943,7 +947,7 @@ vi_cmd(int argcnt, const char *cmd)
 					if (argcnt-- == 0)
 						break;
 			del_range(es->cursor, cur);
-			insert = INSERT;
+			set_insert(INSERT);
 			break;
 
 		case 'v':
@@ -1087,7 +1091,7 @@ vi_cmd(int argcnt, const char *cmd)
 					es->cursor--;
 				return -1;
 			}
-			insert = INSERT;
+			set_insert(INSERT);
 			}
 			break;
 
@@ -1297,7 +1301,7 @@ redo_insert(int count)
 	if (es->cursor > 0)
 		while (isu8cont(es->cbuf[--es->cursor]))
 			continue;
-	insert = 0;
+	set_insert(0);
 	return 0;
 }
 
@@ -2031,7 +2035,7 @@ expand_word(int command)
 	if (rval == 0 && i > 0)
 		es->cursor += i;
 	modified = 1; hnum = hlast;
-	insert = INSERT;
+	set_insert(INSERT);
 	lastac = 0;
 	refresh(0);
 	return rval;
@@ -2135,7 +2139,7 @@ complete_word(int command, int count)
 	x_free_words(nwords, words);
 
 	modified = 1; hnum = hlast;
-	insert = INSERT;
+	set_insert(INSERT);
 	lastac = 0;	 /* prevent this from being redone... */
 	refresh(0);
 
@@ -2197,10 +2201,49 @@ x_vi_zotc(int c)
 	x_putc(c);
 }
 
+static void
+set_insert(int nmode)
+{
+	int update = nmode != insert;
+	insert = nmode;
+
+	if (update)
+		ed_mov_opt(0, wbuf[1 - win]);
+}
+
 static void
 vi_pprompt(int full)
 {
-	pprompt(prompt + (full ? 0 : prompt_skip), prompt_trunc);
+	const char *prefix = (
+		insert == 0 ? vprompt_normal :
+		insert == REPLACE ? vprompt_replace :
+		vprompt_insert);
+	char *nprompt = calloc(strlen(prefix) + strlen(prompt) + 1, sizeof(char));
+	int start_index = 0,
+		stop_index = 0,
+		length = 0;
+	char *start, *stop;
+
+	if ((start = strchr(prompt, '[')) && (stop = strchr(prompt, ']'))) {
+		start_index = (int)(start - prompt);
+		stop_index = (int)(stop - prompt);
+		length = (int)(stop_index - start_index - 1);
+	}
+
+	if (
+		promptlen(vprompt_insert, NULL) == length
+		&& promptlen(vprompt_normal, NULL) == length
+		&& promptlen(vprompt_replace, NULL) == length
+	) {
+		strncpy(nprompt, prompt, start_index + 1);
+		strcat(nprompt, prefix);
+		strncat(nprompt, prompt + stop_index, strlen(prompt) - stop_index);
+	} else {
+		strcat(nprompt, prompt);
+	}
+
+	pprompt(nprompt + (full ? 0 : prompt_skip), prompt_trunc);
+	free(nprompt);
 }
 
 static void
