From 393f7b8bb9d83aef1847d39ad20c6a73849917b8 Mon Sep 17 00:00:00 2001
From: Eric Freese <ericdfreese@gmail.com>
Date: Mon, 14 May 2018 10:59:08 -0600
Subject: [PATCH] Fix vi-mode partial-accept

Issue #188. PR #324.

Thanks to @toadjaune and @IngoHeimbach.
---
 spec/integrations/vi_mode_spec.rb | 46 +++++++++++++++++++++++++++++++
 src/config.zsh                    |  2 ++
 src/widgets.zsh                   | 14 +++++++---
 zsh-autosuggestions.zsh           | 16 ++++++++---
 4 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/spec/integrations/vi_mode_spec.rb b/spec/integrations/vi_mode_spec.rb
index 95f9293..cf471b5 100644
--- a/spec/integrations/vi_mode_spec.rb
+++ b/spec/integrations/vi_mode_spec.rb
@@ -17,5 +17,51 @@ describe 'when using vi mode' do
       end
     end
   end
+
+  describe '`vi-forward-word-end`' do
+    it 'should accept through the end of the current word' do
+      with_history('foobar foo') do
+        session.
+          send_string('foo').
+          send_keys('escape').
+          send_keys('e'). # vi-forward-word-end
+          send_keys('a'). # vi-add-next
+          send_string('baz')
+
+        wait_for { session.content }.to eq('foobarbaz')
+      end
+    end
+  end
+
+  describe '`vi-forward-word`' do
+    it 'should accept through the first character of the next word' do
+      with_history('foobar foo') do
+        session.
+          send_string('foo').
+          send_keys('escape').
+          send_keys('w'). # vi-forward-word
+          send_keys('a'). # vi-add-next
+          send_string('az')
+
+        wait_for { session.content }.to eq('foobar faz')
+      end
+    end
+  end
+
+  describe '`vi-find-next-char`' do
+    it 'should accept through the next occurrence of the character' do
+      with_history('foobar foo') do
+        session.
+          send_string('foo').
+          send_keys('escape').
+          send_keys('f'). # vi-find-next-char
+          send_keys('o').
+          send_keys('a'). # vi-add-next
+          send_string('b')
+
+        wait_for { session.content }.to eq('foobar fob')
+      end
+    end
+  end
 end
 
diff --git a/src/config.zsh b/src/config.zsh
index 5c44af5..c4d07fd 100644
--- a/src/config.zsh
+++ b/src/config.zsh
@@ -49,6 +49,8 @@ ZSH_AUTOSUGGEST_PARTIAL_ACCEPT_WIDGETS=(
 	vi-forward-word-end
 	vi-forward-blank-word
 	vi-forward-blank-word-end
+	vi-find-next-char
+	vi-find-next-char-skip
 )
 
 # Widgets that should be ignored (globbing supported but must be escaped)
diff --git a/src/widgets.zsh b/src/widgets.zsh
index 2004518..87bb62e 100644
--- a/src/widgets.zsh
+++ b/src/widgets.zsh
@@ -153,7 +153,7 @@ _zsh_autosuggest_execute() {
 
 # Partially accept the suggestion
 _zsh_autosuggest_partial_accept() {
-	local -i retval
+	local -i retval cursor_loc
 
 	# Save the contents of the buffer so we can restore later if needed
 	local original_buffer="$BUFFER"
@@ -165,13 +165,19 @@ _zsh_autosuggest_partial_accept() {
 	_zsh_autosuggest_invoke_original_widget $@
 	retval=$?
 
+	# Normalize cursor location across vi/emacs modes
+	cursor_loc=$CURSOR
+	if [[ "$KEYMAP" = "vicmd" ]]; then
+		cursor_loc=$((cursor_loc + 1))
+	fi
+
 	# If we've moved past the end of the original buffer
-	if (( $CURSOR > $#original_buffer )); then
+	if (( $cursor_loc > $#original_buffer )); then
 		# Set POSTDISPLAY to text right of the cursor
-		POSTDISPLAY="$RBUFFER"
+		POSTDISPLAY="${BUFFER[$(($cursor_loc + 1)),$#BUFFER]}"
 
 		# Clip the buffer at the cursor
-		BUFFER="$LBUFFER"
+		BUFFER="${BUFFER[1,$cursor_loc]}"
 	else
 		# Restore the original buffer
 		BUFFER="$original_buffer"
diff --git a/zsh-autosuggestions.zsh b/zsh-autosuggestions.zsh
index bc84bfd..4c06cdb 100644
--- a/zsh-autosuggestions.zsh
+++ b/zsh-autosuggestions.zsh
@@ -85,6 +85,8 @@ ZSH_AUTOSUGGEST_PARTIAL_ACCEPT_WIDGETS=(
 	vi-forward-word-end
 	vi-forward-blank-word
 	vi-forward-blank-word-end
+	vi-find-next-char
+	vi-find-next-char-skip
 )
 
 # Widgets that should be ignored (globbing supported but must be escaped)
@@ -431,7 +433,7 @@ _zsh_autosuggest_execute() {
 
 # Partially accept the suggestion
 _zsh_autosuggest_partial_accept() {
-	local -i retval
+	local -i retval cursor_loc
 
 	# Save the contents of the buffer so we can restore later if needed
 	local original_buffer="$BUFFER"
@@ -443,13 +445,19 @@ _zsh_autosuggest_partial_accept() {
 	_zsh_autosuggest_invoke_original_widget $@
 	retval=$?
 
+	# Normalize cursor location across vi/emacs modes
+	cursor_loc=$CURSOR
+	if [[ "$KEYMAP" = "vicmd" ]]; then
+		cursor_loc=$((cursor_loc + 1))
+	fi
+
 	# If we've moved past the end of the original buffer
-	if (( $CURSOR > $#original_buffer )); then
+	if (( $cursor_loc > $#original_buffer )); then
 		# Set POSTDISPLAY to text right of the cursor
-		POSTDISPLAY="$RBUFFER"
+		POSTDISPLAY="${BUFFER[$(($cursor_loc + 1)),$#BUFFER]}"
 
 		# Clip the buffer at the cursor
-		BUFFER="$LBUFFER"
+		BUFFER="${BUFFER[1,$cursor_loc]}"
 	else
 		# Restore the original buffer
 		BUFFER="$original_buffer"