From 6e15a2d3de6f6f394cc29ee832fbf728f9ecc94c Mon Sep 17 00:00:00 2001
From: Maciej Dems <maciej.dems@p.lodz.pl>
Date: Wed, 8 Jul 2020 11:04:56 +0200
Subject: [PATCH] Preserve ZLE_KILL and ZLE_YANK flags

---
 src/bind.zsh            |  6 ++++++
 src/config.zsh          | 28 +++++++++++++++++++++++++++-
 zsh-autosuggestions.zsh | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/src/bind.zsh b/src/bind.zsh
index 1dde137..3380617 100644
--- a/src/bind.zsh
+++ b/src/bind.zsh
@@ -51,8 +51,14 @@ _zsh_autosuggest_bind_widget() {
 	# correctly. $WIDGET cannot be trusted because other plugins call
 	# zle without the `-w` flag (e.g. `zle self-insert` instead of
 	# `zle self-insert -w`).
+	# Preserve the ZLE_KILL | ZLE_YANK flags for builtin widgets.
 	eval "_zsh_autosuggest_bound_${bind_count}_${(q)widget}() {
 		_zsh_autosuggest_widget_$autosuggest_action $prefix$bind_count-${(q)widget} \$@
+		case ${(q)widget} in
+			(${(j:|:)ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS}) zle -f 'kill';;
+			(${(j:|:)ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS}) zle -f 'yank';;
+			(${(j:|:)ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS}) zle -f 'yankbefore';;
+		esac
 	}"
 
 	# Create the bound widget
diff --git a/src/config.zsh b/src/config.zsh
index 5a0ebd8..9c54585 100644
--- a/src/config.zsh
+++ b/src/config.zsh
@@ -82,7 +82,6 @@ typeset -g ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX=autosuggest-orig-
 		run-help
 		set-local-history
 		which-command
-		yank
 		yank-pop
 		zle-\*
 	)
@@ -91,3 +90,30 @@ typeset -g ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX=autosuggest-orig-
 # Pty name for capturing completions for completion suggestion strategy
 (( ! ${+ZSH_AUTOSUGGEST_COMPLETIONS_PTY_NAME} )) &&
 typeset -g ZSH_AUTOSUGGEST_COMPLETIONS_PTY_NAME=zsh_autosuggest_completion_pty
+
+# Widgets that should preserve ZLE_KILL flag
+(( ! ${+ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS} )) && {
+	typeset -ga ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS
+	ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS=(
+		kill-\*
+		backward-kill-\*
+	)
+}
+
+# Widgets that should preserve ZLE_YANK flag
+(( ! ${+ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS} )) && {
+	typeset -ga ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS
+	ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS=(
+		bracketed-paste
+		vi-put-after
+		yank
+	)
+}
+
+# Widgets that should preserve ZLE_YANKBEFORE flag
+(( ! ${+ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS} )) && {
+	typeset -ga ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS
+	ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS=(
+		vi-put-before
+	)
+}
diff --git a/zsh-autosuggestions.zsh b/zsh-autosuggestions.zsh
index b19cac7..f43ae5b 100644
--- a/zsh-autosuggestions.zsh
+++ b/zsh-autosuggestions.zsh
@@ -108,7 +108,6 @@ typeset -g ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX=autosuggest-orig-
 		run-help
 		set-local-history
 		which-command
-		yank
 		yank-pop
 		zle-\*
 	)
@@ -118,6 +117,33 @@ typeset -g ZSH_AUTOSUGGEST_ORIGINAL_WIDGET_PREFIX=autosuggest-orig-
 (( ! ${+ZSH_AUTOSUGGEST_COMPLETIONS_PTY_NAME} )) &&
 typeset -g ZSH_AUTOSUGGEST_COMPLETIONS_PTY_NAME=zsh_autosuggest_completion_pty
 
+# Widgets that should preserve ZLE_KILL flag
+(( ! ${+ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS} )) && {
+	typeset -ga ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS
+	ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS=(
+		kill-\*
+		backward-kill-\*
+	)
+}
+
+# Widgets that should preserve ZLE_YANK flag
+(( ! ${+ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS} )) && {
+	typeset -ga ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS
+	ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS=(
+		bracketed-paste
+		vi-put-after
+		yank
+	)
+}
+
+# Widgets that should preserve ZLE_YANKBEFORE flag
+(( ! ${+ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS} )) && {
+	typeset -ga ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS
+	ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS=(
+		vi-put-before
+	)
+}
+
 #--------------------------------------------------------------------#
 # Utility Functions                                                  #
 #--------------------------------------------------------------------#
@@ -181,8 +207,14 @@ _zsh_autosuggest_bind_widget() {
 	# correctly. $WIDGET cannot be trusted because other plugins call
 	# zle without the `-w` flag (e.g. `zle self-insert` instead of
 	# `zle self-insert -w`).
+	# Preserve the ZLE_KILL | ZLE_YANK flags for builtin widgets.
 	eval "_zsh_autosuggest_bound_${bind_count}_${(q)widget}() {
 		_zsh_autosuggest_widget_$autosuggest_action $prefix$bind_count-${(q)widget} \$@
+		case ${(q)widget} in
+			(${(j:|:)ZSH_AUTOSUGGEST_ZLE_KILL_WIDGETS}) zle -f 'kill';;
+			(${(j:|:)ZSH_AUTOSUGGEST_ZLE_YANK_WIDGETS}) zle -f 'yank';;
+			(${(j:|:)ZSH_AUTOSUGGEST_ZLE_YANKBEFORE_WIDGETS}) zle -f 'yankbefore';;
+		esac
 	}"
 
 	# Create the bound widget