module/.autocomplete.async

#!/bin/zsh
zmodload -Fa zsh/zpty b:zpty
zmodload -Fa zsh/parameter p:funcstack p:functions p:parameters
zmodload -Fa zsh/system p:sysparams
zmodload -Fa zsh/zselect b:zselect
zmodload -Fa zsh/terminfo b:echoti
zmodload -Fa zsh/zutil b:zparseopts
builtin autoload -Uz add-zle-hook-widget

typeset -gHi _autocomplete__buffer_start_line=1
typeset -g ZSH_AUTOSUGGEST_USE_ASYNC=yes

zle -N history-incremental-search-backward .autocomplete.async.history-incremental-search
zle -N history-incremental-search-forward .autocomplete.async.history-incremental-search
bindkey -s -M menuselect '^R' '^G^R'
bindkey -s -M menuselect '^S' '^G^S'

.autocomplete.async.precmd() {
  typeset -gHF _autocomplete__async_avg_duration
  zstyle -s :autocomplete: min-delay _autocomplete__async_avg_duration
  (( _autocomplete__async_avg_duration += 0.1 ))

  zle -N .autocomplete.async.complete.fd-widget
  zle -N .autocomplete.async.timeout.fd-widget

  # Name starting with `.` avoids getting wrapped by syntax highlighting.
  zle -C ._list_choices list-choices .autocomplete.async.list-choices.completion-widget

  if [[ -v functions[_zsh_highlight_call_widget] ]]; then
    _zsh_highlight_call_widget() {
      .autocomplete.zle-flags $WIDGET
      builtin zle "$@"
    }
  fi

  if [[ -v functions[_zsh_autosuggest_highlight_apply] ]]; then
    local action
    for action in clear modify fetch accept partial_accept execute enable disable toggle; do
      eval "_zsh_autosuggest_widget_$action() {
        .autocomplete.zle-flags \$WIDGET  # Set flags according to widget name.
        _zsh_autosuggest_$action \"\$@\"
      }"
    done
    _zsh_autosuggest_widget_suggest() {
      .autocomplete.zle-flags  # Maintain previously set flags.
      _zsh_autosuggest_suggest "$@"
    }
  fi

  add-zsh-hook precmd .autocomplete.async.update-location

  add-zle-hook-widget line-init .autocomplete.async.reset-context
  add-zle-hook-widget line-init .autocomplete.async.complete
  add-zle-hook-widget line-pre-redraw .autocomplete.async.complete
  add-zle-hook-widget line-finish .autocomplete.async.clear
  add-zle-hook-widget isearch-update .autocomplete.async.isearch-update
  add-zle-hook-widget isearch-exit .autocomplete.async.isearch-exit

  add-zsh-hook zshexit .autocomplete.async.stop
}

.autocomplete.async.max-lines() {
  local -i min_lines max_lines lines_below_buffer
  zstyle -s ":autocomplete:${curcontext}:" list-lines min_lines || min_lines=16
  (( max_lines = LINES - BUFFERLINES ))
  (( lines_below_buffer = max_lines - _autocomplete__buffer_start_line ))
  return $(( max( min( min_lines, max_lines ), lines_below_buffer ) ))
}

.autocomplete.async.update-location() {
  emulate -L zsh; setopt $_autocomplete__options

  # Ask the terminal which line the cursor is on.
  local -a reply=()
  print -nr $'\e[6n'
  IFS=$'\e[;' read -Asd R

  # Add the number of newlines the prompt will print.
  local -i pslines=${#${(%%)PS1}//[^$'\n']}
  _autocomplete__buffer_start_line=$(( reply[-2] + pslines ))
}

.autocomplete.async.history-incremental-search() {
  if [[ $curcontext == $WIDGET* ]]; then
    unset curcontext
  else
    typeset -gH curcontext=${WIDGET}:::
  fi
  .autocomplete.async.start
}

.autocomplete.async.reset-context() {
  typeset -gH curcontext=''
  zstyle -s :autocomplete: default-context curcontext
  return 0
}

.autocomplete.async.isearch-update() {
  typeset -gHi _autocomplete__isearch=1
}

.autocomplete.async.isearch-exit() {
  .autocomplete.zle-flags $LASTWIDGET
  unset _autocomplete__isearch
}

.autocomplete.async.complete() {
  .autocomplete.zle-flags $LASTWIDGET

  (( KEYS_QUEUED_COUNT || PENDING )) &&
    return 0

  [[ $LASTWIDGET == .autocomplete.async.complete.fd-widget ]] &&
    return 0

  region_highlight=()
  [[ -v functions[_zsh_highlight] ]] &&
    _zsh_highlight
  typeset -gH _autocomplete__highlight=( "$region_highlight[@]" )
  [[ -v functions[_zsh_autosuggest_highlight_apply] ]] &&
    _zsh_autosuggest_highlight_apply

  if (( REGION_ACTIVE )) ||
      [[ -v _autocomplete__isearch && $LASTWIDGET == *(incremental|isearch)* ]]; then
    zle -Rc
    return 0
  fi

  if [[ $LASTWIDGET == ('list-expand'|'_complete_help') ||
      -n $MENUSELECT && $LASTWIDGET == .autocomplete.async.timeout.fd-widget ]]; then
    .autocomplete.async.stop
    return 0
  fi

  [[ $_lastcomp[insert] == *unambiguous ]] &&
    zle .auto-suffix-retain  # Makes the cursor stay in the right place.

  .autocomplete.async.start
  return 0
}

.autocomplete.async.stop() {
  .autocomplete.async.kill _autocomplete__async_complete_fd _autocomplete__async_complete_pid
  .autocomplete.async.kill _autocomplete__async_timeout_fd _autocomplete__async_timeout_pid
  return 0
}

.autocomplete.async.clear() {
  zle -Rc
  unset _autocomplete__isearch
  .autocomplete.async.stop
  .autocomplete.async.reset-context
  return 0
}

.autocomplete.async.kill() {
  emulate -L zsh; setopt $_autocomplete__options

  local -i fd=${(P)1} pid=${(P)2}
  if (( pid )); then
    [[ -o MONITOR ]] &&
      (( pid *= -1 ))
    kill -KILL $pid 2>/dev/null
  fi
  if (( fd )) && { : <&$fd } 2>/dev/null; then
    zle -F $fd 2>/dev/null
    exec {fd}<&-
  fi
  unset $1 $2
}

.autocomplete.async.start() {
  setopt $_autocomplete__options

  local +h -F SECONDS=0

  .autocomplete.async.stop

  local -F min_delay; zstyle -s :autocomplete: min-delay min_delay ||
    min_delay=0

  exec {_autocomplete__async_complete_fd}< <(
    print $sysparams[pid]
    {
      local REPLY
      zpty _autocomplete__zpty .autocomplete.async.complete.inner '$LBUFFER' '$RBUFFER'
      zpty -w _autocomplete__zpty $'\t'

      local line
      zpty -r _autocomplete__zpty line '*'$'\0'$'\0'
      zpty -r _autocomplete__zpty line '*'$'\0'$'\0'

      # `zselect -t` w/out other args wants an int > 0
      local -i timeout=$(( 100.0 * (min_delay - SECONDS) ))
      (( timeout > 0 )) &&
        zselect -t $timeout
    } always {
      zpty -d _autocomplete__zpty
      (( TRY_BLOCK_INTERRUPT > 0 )) ||
        print -rNC1 - "$SECONDS" "$line"
    }
  )
  read -u $_autocomplete__async_complete_fd _autocomplete__async_complete_pid
  zle -Fw "$_autocomplete__async_complete_fd" .autocomplete.async.complete.fd-widget

  exec {_autocomplete__async_timeout_fd}< <(
    print $sysparams[pid]
    {
      local -i timeout=$((
        100.0 * (min_delay - SECONDS + 10.0 * _autocomplete__async_avg_duration)
      ))
      (( timeout > 0 )) &&
        zselect -t $timeout
    } always {
      (( TRY_BLOCK_INTERRUPT ))
      print $?
    }
  )
  read -u $_autocomplete__async_timeout_fd _autocomplete__async_timeout_pid
  zle -Fw "$_autocomplete__async_timeout_fd" .autocomplete.async.timeout.fd-widget

  [[ -v functions[_zsh_autosuggest_highlight_apply] && -n $MENUSELECT ]] &&
      unset POSTDISPLAY

  # There's a weird bug in Zsh < 5.8, where where ^C stops working unless we force a fork.
  # See https://github.com/zsh-users/zsh-autosuggestions/issues/364
  command true
}

.autocomplete.async.complete.inner() {
  setopt $_autocomplete__options
  setopt NO_completeinword; [[ -z $curcontext ]] ||
    setopt completeinword

  setopt promptsubst
  PS1=''
  PS4=$_autocomplete__x

  local hooks=( chpwd periodic precmd preexec zshaddhistory zshexit )
  unset ${^hooks}_functions &>/dev/null
  $hooks[@] () { : }
  .autocomplete.no-op() { : }
  local hook; for hook in isearch-exit isearch-update line-pre-redraw line-init line-finish \
      history-line-set keymap-select; do
    zle -N zle-$hook .autocomplete.no-op
  done

  typeset -gH lbuffer=$1 rbuffer=$2

  complete-word() {
    LBUFFER=$lbuffer
    RBUFFER=$rbuffer
    zle list-choices -w
  }

  message() {
    typeset -gH _message_=$mesg
  }

  list-choices() {
    local curcontext=${curcontext:-${WIDGET}:::}
    unset 'compstate[vared]'
    local +h -a comppostfuncs=( message )
    {
      .autocomplete.async.list-choices.complete 2>&$_autocomplete__log_fd
    } always {
      print -rNC1 - '' '' "$compstate[nmatches]" "${(q+)_message_}" ''
    }
    compstate[insert]=''
    compstate[list]=''
    compstate[list_max]=0
  }

  zle -N complete-word
  zle -C list-choices list-choices list-choices
  bindkey '^I' complete-word
  local __tmp__
  vared __tmp__
}

.autocomplete.async.complete.fd-widget() {
  .autocomplete.zle-flags ||
    return 0

  setopt $_autocomplete__options
  setopt NO_completeinword; [[ -z $curcontext ]] ||
    setopt completeinword

  .autocomplete.async.kill _autocomplete__async_timeout_fd _autocomplete__async_timeout_pid

  local -F seconds
  local -i nmatches
  local message rest
  {
    [[ $2 == (|hup) ]] &&
      IFS=$'\0' read -ru $1 seconds nmatches message rest
  } always {
    .autocomplete.async.kill _autocomplete__async_complete_fd _autocomplete__async_complete_pid
  }

  (( _autocomplete__async_avg_duration = .1 * seconds + .9 * _autocomplete__async_avg_duration ))

  # If a widget can't be called, ZLE always returns true.
  # Thus, we return false on purpose, so we can check if our widget got called.
  setopt promptsubst
  local +h PS4=$_autocomplete__x
  if ! zle ._list_choices -w "$nmatches" "${(Q)message}" 2>&$_autocomplete__log_fd; then
    region_highlight=( "$_autocomplete__highlight[@]" )
    [[ -v functions[_zsh_autosuggest_highlight_apply] ]] &&
      _zsh_autosuggest_highlight_apply

    # Refresh if and only if our widget got called. Otherwise, Zsh will crash (eventually).
    zle -R
  else
    .autocomplete.async.stop
  fi
  return 0
}

.autocomplete.async.list-choices.completion-widget() {
  local -i nmatches=$1
  local message=$2

  local min_input; zstyle -s :autocomplete: min-input min_input ||
    min_input=0
  local ignored; zstyle -s :autocomplete: ignored-input ignored

  local +h -a comppostfuncs=( .autocomplete.async.list-choices.post "$comppostfuncs[@]" )
  if [[ -n $curcontext ]]; then
    _main_complete
  elif (( $#words[@] == 1 && $#words[CURRENT] < min_input )) ||
      [[ -n $words[CURRENT] && $words[CURRENT] == $~ignored ]]; then
    zle -Rc
  elif (( nmatches == 0 )); then
    if [[ -n $message ]]; then
      local msg
      zformat -f msg "$message"
      compadd -x "$msg"
    else
      zle -Rc
    fi
  else
    local curcontext=list-choices:::
    .autocomplete.async.list-choices.complete
  fi

  return 2
}

.autocomplete.async.list-choices.post() {
  compstate[insert]=
  unset MENUSELECT MENUMODE
}

.autocomplete.async.list-choices.complete() {
  .autocomplete.async.max-lines
  local -i _async_max_lines=$?
  {
    () {
      emulate -L zsh; setopt $_autocomplete__options
    functions[compadd]=$functions[.autocomplete.async.compadd]
    } "$@"
    _main_complete
  } always {
    () {
      emulate -L zsh; setopt $_autocomplete__options
    [[ -v functions[compadd] ]] &&
      unfunction compadd
    [[ -v functions[comptags] ]] &&
      unfunction comptags
    }
  }
}

.autocomplete.async.compadd() {
  if [[ $funcstack[2] == _autocomplete.history_lines ]]; then
    .autocomplete.compadd "$@"
    return
  fi

  local -a _opts_=()
  zparseopts -E -A _opts_ - D: E: x: X:

  if [[ -v _opts_[-E] || -v _opts_[-x] ]]; then
    .autocomplete.compadd "$@"
    return
  fi

  local -a _matches_=()
  local -i header=$+_opts_[-X] footer=1 allinsert=2
  local -i lines_available=$((
    max(0, _async_max_lines - header - footer - allinsert - compstate[list_lines])
  ))
  local -i ret number_of_new_matches lines_of_new_matches

  if [[ -v _opts_[-D] ]]; then
    .autocomplete.compadd "$@"; ret=$?

    [[ $funcstack[2] == _describe ]] ||
      return ret

    local array_name=$_opts_[-D]
    local -a _matches_=( ${(PA)array_name} )

    (( ${_matches_[(I)*:*]} > 0 )) ||
      return ret

    local -aU uniques=( ${_matches_[@]#*:} )
    number_of_new_matches=$#_matches_
    lines_of_new_matches=$#uniques[@]
    (( lines_available -= _autocomplete__reserved_lines ))
  else
    _autocomplete__reserved_lines=0
    local -i fd list_lines=0 nmatches=0
    {
      exec {fd}< <(
        {
          _opts_=()
          zparseopts -D -E -A _opts_ - A: D: O: X:
          .autocomplete.compadd -O _matches_ "$@"

          # Pre-emptively trim all matches that will definitely not fit on screen.
          local -i surplus=$(( $#_matches_ - COLUMNS * _async_max_lines / 3 ))
          (( surplus > 0 )) &&
            shift -p $surplus _matches_

          setopt localoptions listtypes
          zparseopts -D -E -A _opts_ - a
          _autocomplete.compadd_opts_len "$@"
          .autocomplete.compadd "${(@)@[1,?]}" -a - _matches_
        } always {
          print $compstate[list_lines]
          print $compstate[nmatches]
          print -rNC1 - "$_matches_[@]"
        }
      )
      IFS=$'\0' read -u $fd list_lines
      IFS=$'\0' read -u $fd nmatches
      IFS=$'\0' read -u $fd -A _matches_
    } always {
      exec {fd}<&-
    }
    _matches_=( ${_matches_[@]:#} )
    number_of_new_matches=$(( nmatches - compstate[nmatches] ))
    lines_of_new_matches=$(( list_lines - compstate[list_lines] ))
  fi

  if (( lines_of_new_matches <= lines_available )); then
    if [[ $funcstack[2] == _describe && -v _opts_[-D] ]]; then
      (( _autocomplete__reserved_lines += $lines_of_new_matches ))
      return ret
    fi
    .autocomplete.compadd "$@"
    return
  fi

  local info hint=$'%{\e[02;39m%}' kbd=$'%{\e[22;39m%}' end=$'%{\e[0m%}'
  zstyle -s ":autocomplete:${curcontext}:too-many-matches" message info ||
    info="${hint}(partial list; press ${kbd}Down$hint to expand)$end"

  if [[ -v _opts_[-D] ]]; then
    if (( $#uniques > lines_available )); then
      builtin compadd -x $info
      shift -p $(( $#uniques - max(0, lines_available - 1) )) uniques
      set -A $array_name ${(M@)${(PA)array_name}:#*:(${(~j:|:)uniques})}
    fi
    (( _autocomplete__reserved_lines += $#uniques ))
    (( ${(P@)#array_name} > 0 ))
    return
  fi

  zparseopts -D -E -A _opts_ - a d: k l U
  local -a dopt=()
  if [[ -v _opts_[-d] ]]; then
    if [[ -v _opts_[-l] ]]; then
      dopt+=( -ld )
    else
      dopt+=( -d )
    fi
    dopt+=( "$_opts_[-d]" )
  fi

  if (( lines_of_new_matches > 0 )); then
    local -F matches_per_line=$(( 1.0 * number_of_new_matches / lines_of_new_matches ))
    if (( matches_per_line < 1 )); then
      dopt[1]=-ld
      array_name=$_opts_[-d]
      if [[ -z $array_name ]]; then
        array_name=displ
        local -a displ=( "$_matches_[@]" )
        dopt[2]=displ
      fi
      set -A $array_name ${(r:COLUMNS-1:@)${(P@)array_name}//$'\n'/\n}
      matches_per_line=1
    fi
    if (( $#_matches_ > matches_per_line * lines_available )); then
      builtin compadd -x $info

      # Round this down, before subtracting.
      local -i fit=$(( matches_per_line * max(0, lines_available - 1) ))

      shift -p $(( $#_matches_ - fit )) _matches_
      (( $#_matches_ > 0 )) ||
        comptags() {
          [[ $funcstack[3] == _autocomplete.history_lines ]] ||
            return 1
          builtin comptags "$@"
        }
    fi
  fi
  _autocomplete.compadd_opts_len "$@"
  .autocomplete.compadd "${(@)@[1,?]}" $dopt -a _matches_
}

.autocomplete.async.timeout.fd-widget() {
  .autocomplete.zle-flags ||
    return 0

  local -i timedout=0
  {
    [[ $2 == (|hup) ]] &&
      read -ru $1 timedout
  } always {
    .autocomplete.async.kill _autocomplete__async_timeout_fd _autocomplete__async_timeout_pid
  }

  (( timedout )) ||
    return 0

  (( _autocomplete__async_avg_duration *= 1.1 ))
  .autocomplete.async.kill _autocomplete__async_complete_fd _autocomplete__async_complete_pid
  return 0
}