Make zsh zle restore characters on backspace in vi overwrite-mode

I am using Zsh in Vi-mode.

When $KEYMAP == vicmd (i.e. command-mode), I want hitting backspace to move the cursor to the left by one character, without deleting anything. [working]

When $KEYMAP == viins && $ZLE_STATE == *insert* (i.e. insert-mode), I want hitting backspace to move the cursor to the left by one character, deleting the immediately preceding character on the line. [working]

When $KEYMAP == viins && $ZLE_STATE == *overwrite* (i.e. overwrite-mode / replace-mode), I want hitting backspace to move the cursor to the left by one character, restoring the immediately preceding character on the line with the one that had originally been there prior to entering into overwrite-mode. [NOT working]

Here is an example:

# [COMMAND MODE] We start with the following string on the command line:
$ Hello, world!
     ^
     cursor position

# [REPLACE MODE] Now, I hit "R" to enter replace-mode and I type "stuff".
$ Helstufforld!
          ^
          cursor position

# [REPLACE MODE] Finally, I hit backspace 3 times.
$ Helst, world!
       ^
       cursor position

The above example shows what I want to happen when I hit backspace while in overwrite-mode; however, what really happens is the following:

# [COMMAND MODE] We start with the following string on the command line:
$ Hello, world!
     ^
     cursor position

# [REPLACE MODE] Now, I hit "R" to enter replace-mode and I type "stuff".
$ Helstufforld!
          ^
          cursor position

# [REPLACE MODE] Finally, I hit backspace 3 times.
$ Helstworld!
       ^
       cursor position

Notice how, when hitting backspace in the second example, rather than restoring the original 3 characters that were just overwritten (i.e. ", w"), instead the last 3 characters that replaced these characters (i.e. "uff") were deleted, and the characters to the right of the cursor were shifted to the left.

How do I get the behavior that I want?

Okay, so I ended up hacking together a solution to the problem I was having, which I will post here in case anyone else encounters the same issue.

Solution

Put this in your .zshrc:

readonly ZLE_VI_MODE_CMD=0
readonly ZLE_VI_MODE_INS=1
readonly ZLE_VI_MODE_REP=2
readonly ZLE_VI_MODE_OTH=3

function zle-vi-mode {
    if [[ $KEYMAP == vicmd ]]; then
        echo -n $ZLE_VI_MODE_CMD
    elif [[ $KEYMAP == (viins|main) ]] && [[ $ZLE_STATE == *insert* ]]; then
        echo -n $ZLE_VI_MODE_INS
    elif [[ $KEYMAP == (viins|main) ]] && [[ $ZLE_STATE == *overwrite* ]]; then
        echo -n $ZLE_VI_MODE_REP
    else
        echo -n $ZLE_VI_MODE_OTH
    fi
}

function zle-backward-delete-char-fix {
    case "$(zle-vi-mode)" in
        $ZLE_VI_MODE_REP)
            if [[ $CURSOR -le $MARK ]]; then
                CURSOR=$(( $(($CURSOR-1)) > 0 ? $(($CURSOR-1)) : 0 ))
                MARK=$CURSOR
            else
                zle undo
            fi
            ;;
        *)
            zle backward-delete-char
            ;;
    esac
}

zle -N zle-backward-delete-char-fix

## Change cursor shape according to the current Vi-mode.
function zle-line-init zle-keymap-select {
    case "$(zle-vi-mode)" in
        $ZLE_VI_MODE_CMD) echo -ne '\e[2 q' ;; # cursor -> block
        $ZLE_VI_MODE_INS) echo -ne '\e[6 q' ;; # cursor -> vertical bar
        $ZLE_VI_MODE_REP)
            echo -ne '\e[4 q' # cursor -> underline
            MARK=$CURSOR
            ;;
        *)
            ;;
    esac
}

zle -N zle-line-init
zle -N zle-keymap-select

bindkey -v

bindkey '^?' zle-backward-delete-char-fix
bindkey '^h' zle-backward-delete-char-fix

The above code will, additionally, cause your cursor shape to change depending on what vi-mode you are currently in (i.e. since this is a copy/paste from my .zshrc, and that's what I like). If you don't want this, and just want the plain fix, replace the zle-init-line / zle-keymap-select function with the following:

function zle-line-init zle-keymap-select {
    case "$(zle-vi-mode)" in
        $ZLE_VI_MODE_REP)
            MARK=$CURSOR
            ;;
        *)
            ;;
    esac
}