.autocomplete.async

#!/bin/zsh
zmodload -F zsh/zpty b:zpty
zmodload -F zsh/parameter p:funcstack p:functions p:parameters
zmodload -F zsh/system b:sysopen p:sysparams
zmodload -F zsh/zselect b:zselect
zmodload -F zsh/terminfo b:echoti p:terminfo
zmodload -F 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
private -a log_functions=()

builtin zle -N history-incremental-search-backward .autocomplete.async.history-incremental-search
builtin zle -N history-incremental-search-forward .autocomplete.async.history-incremental-search

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

  # Start names with `.` to avoid getting wrapped by syntax highlighting.
  builtin zle -N .autocomplete.async.pty.zle-widget
  builtin zle -C .autocomplete.async.pty.completion-widget list-choices \
      .autocomplete.async.pty.completion-widget
  builtin zle -N .autocomplete.async.complete.fd-widget
  builtin 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
    private action=
    for action in clear modify fetch accept partial_accept execute enable disable toggle; do
      # Set flags according to widget name.
      eval "_zsh_autosuggest_widget_$action() {
        .autocomplete.zle-flags \$WIDGET
        _zsh_autosuggest_$action \"\$@\"
      }"
    done
    _zsh_autosuggest_widget_suggest() {
      .autocomplete.zle-flags  # Maintain previously set flags.
      _zsh_autosuggest_suggest "$@"
    }
  fi

  .autocomplete.patch _message

  add-zle-hook-widget line-init .autocomplete.async.read-cursor-position
  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.read-cursor-position() {
  emulate -L zsh
  # Cannot set strict shell options here, because other programs' async
  # callbacks can interrupt the code and then get executed in the same context.
  # See https://github.com/marlonrichert/zsh-autocomplete/issues/342

  (( KEYS_QUEUED_COUNT || PENDING )) &&
      return

  # $MC_SID is set when we're in a subshell started by Midnight Commander.
  # See https://github.com/marlonrichert/zsh-autocomplete/issues/269
  if [[ -v MC_SID || ! ( -v terminfo[u6] && -v terminfo[u7] ) ]]; then
    local -i max_lines=
    builtin zstyle -s ":autocomplete:${curcontext}:" list-lines max_lines ||
        max_lines=16
    typeset -gHi _autocomplete__buffer_start_line=$(( min( max( LINES - max_lines, 1 ), LINES ) ))
    return 0
  fi

  # Cursor position report (u6) always contains two %d markers.
  private -a CPR=( "${(s:%d:)$( echoti u6 )}" )

  # If there is a %i before the first %d, then decrement both %d values by 1.
  private -i i=${${${(M)CPR[1]%'%i'}:+1}:-0}
  CPR[1]=${CPR[1]%'%i'}

  local REPLY=
  private Y=
  echoti u7  # Send cursor position request.

  # Parse the cursor position report.
  read -rsk $#CPR[1]  # Discard the preamble.
  while [[ $REPLY != $CPR[2] ]]; do
    read -rsk
    Y+=$REPLY
  done
  Y="${Y%$CPR[2]}"  # Discard the separator.

  # Flush the rest.
  while [[ $REPLY != $CPR[3] ]]; do
    read -rsk
  done

  typeset -gHi _autocomplete__buffer_start_line=$(( min( max( Y - i, 1 ), LINES ) ))
}

.autocomplete.async.history-incremental-search() {
  if [[ $curcontext == $WIDGET* ]]; then
    unset curcontext
  else
    typeset -gH curcontext=${WIDGET}:::
  fi
  [[ -o sharehistory ]] &&
      fc -RI  # Get new history events from file.
  .autocomplete.async.start
}

.autocomplete.async.reset-context() {
  local context
  builtin zstyle -s :autocomplete: default-context context
  typeset -gH curcontext=$context
  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

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

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

  .autocomplete.async.stop

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

  [[ $LASTWIDGET ==
        (_complete_help|list-expand|(|.)(describe-key-briefly|what-cursor-position|where-is)) ]] &&
      return

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

  .autocomplete.async.start
  return 0
}

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

.autocomplete.async.stop() {
  local fd=$_autocomplete__async_complete_fd
  unset _autocomplete__async_complete_fd
  unset _autocomplete__mesg _autocomplete__comp_mesg _autocomplete__words _autocomplete__current
  if [[ $fd == <-> ]]; then
    builtin zle -F $fd 2> /dev/null
    exec {fd}<&-
  fi
}

.autocomplete.async.start() {
  local fd=
  sysopen -r -o cloexec -u fd <(
    typeset -F SECONDS=0
    setopt promptsubst
    PS4=$_autocomplete__ps4
    .autocomplete.async.start.inner
  )
  builtin zle -Fw "$fd" .autocomplete.async.complete.fd-widget
  typeset -gH _autocomplete__async_complete_fd=$fd

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

.autocomplete.async.start.inner() {
  {
    private hooks=( chpwd periodic precmd preexec zshaddhistory zshexit )
    builtin unset ${^hooks}_functions &> /dev/null
    $hooks[@] () { : }

    private hook=
    for hook in \
        zle-{isearch-{exit,update},line-{pre-redraw,init,finish},history-line-set,keymap-select}
    do
      builtin zle -N $hook .autocomplete.async.pty.no-op
    done
    {
      local REPLY=
      zpty AUTOCOMPLETE .autocomplete.async.pty
      private -i fd=$REPLY

      zpty -w AUTOCOMPLETE $'\t'

      local -F min_delay=
      builtin zstyle -s :autocomplete: min-delay min_delay ||
          min_delay=0.01
      zselect -t "$(( [#10] 100 * max( 0, min_delay - SECONDS ) ))"

      local header=
      zpty -r AUTOCOMPLETE header $'*\C-B'

      local -a reply=()
      local text=
      zselect -rt "$((
          [#10] 100 * max( 0, 100 * _autocomplete__async_avg_duration - SECONDS )
      ))" "$fd" &&
          zpty -r AUTOCOMPLETE text $'*\C-C'
    } always {
      zpty -d AUTOCOMPLETE
    }
  } always {
    # Always produce output, so we always reach the callback, so we can close the fd and unset
    # $_autocomplete__async_complete_fd (if necessary).
    print -rNC1 -- "$SECONDS" "${text%$'\0\C-C'}"
  }
} 2>>| $_autocomplete__log_file
log_functions+=( .autocomplete.async.start.inner )

.autocomplete.async.pty() {
  typeset -gH _autocomplete__lbuffer="$LBUFFER" _autocomplete__rbuffer="$RBUFFER"

  builtin bindkey $'\t' .autocomplete.async.pty.zle-widget
  local __tmp__=
  builtin vared __tmp__
} 2>>| $_autocomplete__log_file
log_functions+=( .autocomplete.async.pty )

.autocomplete.async.pty.no-op() {
  :
}

.autocomplete.async.pty.zle-widget() {
  .autocomplete.async.pty.zle-widget.inner "$@"
}

.autocomplete.async.pty.zle-widget.inner() {
  # The completion widget sometimes returns without calling its function. So, we need to print all
  # our control characters here, to ensure we don't end up waiting endlessly to read them.
  {
    print -n -- '\C-B'
    LBUFFER=$_autocomplete__lbuffer
    RBUFFER=$_autocomplete__rbuffer

    setopt $_autocomplete__comp_opts[@]
    [[ -n $curcontext ]] &&
        setopt $_autocomplete__ctxt_opts[@]
    builtin zle .autocomplete.async.pty.completion-widget -w 2> /dev/null
  } always {
    print -n -- '\C-C'
    builtin kill $sysparams[pid]
  }
} 2>>| $_autocomplete__log_file
log_functions+=( .autocomplete.async.pty.zle-widget.inner )

.autocomplete.async.pty.completion-widget() {
  .autocomplete.async.pty.completion-widget.inner "$@"
}

.autocomplete.async.pty.completion-widget.inner() {
  if .autocomplete.async.insufficient-input; then
    print -rNC1 -- "0" "" ""
    return
  fi
  if .autocomplete.async.same-state; then
    print -rNC1 -- "$_lastcomp[list_lines]" "$_autocomplete__mesg" "$_autocomplete__comp_mesg[@]"
    return
  fi
  unset _autocomplete__mesg _autocomplete__comp_mesg
  {
    local curcontext=${curcontext:-list-choices:::}
    unset 'compstate[vared]'

    _message() {
      compadd() {
        typeset -gHa _autocomplete__comp_mesg=( "$@" )
        builtin compadd "$@"
      }
      zformat() {
        builtin zformat "$@"
        typeset -gHa _autocomplete__comp_mesg=( "$gopt[@]" -x "$format" )
      }
      .autocomplete._message "$@"
      unfunction zformat
      functions[compadd]="$functions[.autocomplete.compadd]"
    }

    local +h -a comppostfuncs=( .autocomplete.async.pty.message )
    _main_complete
  } always {
    print -rNC1 -- "$compstate[list_lines]" "$_autocomplete__mesg" "$_autocomplete__comp_mesg[@]"
  }
} 2>>| $_autocomplete__log_file
log_functions+=( .autocomplete.async.pty.completion-widget.inner )

.autocomplete.async.pty.message() {
  typeset -gH _autocomplete__mesg=$mesg
  return 0
}

.autocomplete.async.complete.fd-widget() {
  setopt promptsubst
  local +h PS4=$_autocomplete__ps4
  .autocomplete.async.complete.fd-widget.inner "$@"
}

.autocomplete.async.complete.fd-widget.inner() {
  local -i fd=$1
  {
    builtin zle -F $fd # Unhook ourselves immediately, so we don't get called more than once.

    # Ensure our input will not be stopped.
    unset _autocomplete__async_complete_fd

    .autocomplete.zle-flags ||
        return 0

    local -a reply=()
    IFS=$'\0' read -rAu $fd
    shift -p reply
  } always {
    exec {fd}<&-
  }

  unset _autocomplete__mesg _autocomplete__comp_mesg

  # 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 $_autocomplete__comp_opts[@]
  [[ -n $curcontext ]] &&
      setopt $_autocomplete__ctxt_opts[@]
  if ! builtin zle ._list_choices -w "$reply[@]" 2>>| $_autocomplete__log_file; 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).
    builtin zle -R
  else
    .autocomplete.async.stop
  fi
  return 0
} 2>>| $_autocomplete__log_file
log_functions+=( .autocomplete.async.complete.fd-widget.inner )

.autocomplete.async.insufficient-input() {
  local min_input=
  builtin zstyle -s :autocomplete:${curcontext:-list-choices} min-input min_input ||
      min_input=0
  local ignored=
  builtin zstyle -s :autocomplete:${curcontext:-list-choices} ignored-input ignored
  if (( ${#words[@]} == 1 && ${#words[CURRENT]} < min_input )) ||
      [[ -n $words[CURRENT] && $words[CURRENT] == $~ignored ]]; then
    compstate[list]=
    true
  else
    false
  fi
}
# functions -T .autocomplete.async.insufficient-input

.autocomplete.async.same-state() {
  [[ $_autocomplete__words == $words && $_autocomplete__current == $CURRENT ]]
}

.autocomplete.async.list-choices.completion-widget() {
  unset _autocomplete__mesg _autocomplete__comp_mesg _autocomplete__words _autocomplete__current

  if [[ $1 != <->.<-> || $2 != <-> ]]; then
    compstate[list]=
    return
  fi

  .autocomplete.async.insufficient-input &&
      return

  private -F _seconds_=$1
  private -i _list_lines_=$2
  private _mesg_=$3
  shift 3

  local +h -a comppostfuncs=( .autocomplete.async.list-choices.post "$comppostfuncs[@]" )
  if [[ -n $compstate[old_list] ]] && .autocomplete.async.same-state; then
    compstate[old_list]=keep
  elif [[ $_list_lines_ == 1 && -n $1 ]]; then
    builtin compadd "$@"
  elif [[ $_list_lines_ == 1 && -n $_mesg_ ]]; then
    builtin compadd -x "$_mesg_"
  else
    typeset -gHF _autocomplete__async_avg_duration=$((
        .1 * _seconds_ + .9 * _autocomplete__async_avg_duration
    ))
    if [[ -n $curcontext ]]; then
      _main_complete
    else
      local curcontext=list-choices:::
      .autocomplete.async.list-choices.main-complete
    fi
  fi

  typeset -gH _autocomplete__mesg=$_mesg_
  typeset -gHa _autocomplete__comp_mesg=( "$@" )
  typeset -gHa _autocomplete__words=( "$words[@]" )
  typeset -gHi _autocomplete__current=$CURRENT

  return 2  # Don't return 1, to prevent beeping.
} 2>>| $_autocomplete__log_file
log_functions+=( .autocomplete.async.list-choices.completion-widget )

.autocomplete.async.list-choices.post() {
  [[ -v _autocomplete__partial_list ]] &&
      builtin compadd -J -last- -x '%F{0}%K{14}(MORE)%f%k'
  compstate[insert]=
  unset MENUSELECT MENUMODE
}

.autocomplete.async.list-choices.main-complete() {
  {
    local -i min_lines=
    builtin zstyle -s ":autocomplete:${curcontext}:" list-lines min_lines ||
        min_lines=16
    min_lines=$(( min( LINES - ( 1 + BUFFERLINES ), min_lines ) ))
    private -i lines_below_buffer=$(( LINES - ( _autocomplete__buffer_start_line + BUFFERLINES ) ))
    local -i _autocomplete__max_lines=$(( max( min_lines, lines_below_buffer ) ))
    local -i _autocomplete__reserved_lines=0
    () {
      emulate -L zsh
      setopt $_autocomplete__func_opts[@]
      functions[compadd]=$functions[.autocomplete.async.compadd]
      # functions -T compadd _describe
    } "$@"
    _main_complete "$@"
  } always {
    unfunction compadd comptags 2> /dev/null
  }
}

.autocomplete.async.compadd() {
  local -A _opts_=()
  local -aU _dopt_=()
  zparseopts -A _opts_ -E -- D: E: J: O: V: x: X: d:=_dopt_

  if [[ ( -v _opts_[-x] && $# == 2 ) ]]; then
    .autocomplete.compadd "$@"
    return
  elif [[ $funcstack[(I)_describe] -gt 0 && ! -v _opts_[-D] ]]; then
    _autocomplete__reserved_lines=0
  fi

  private _displ_array_= _matches_array_=
  private -i _ret_=1 _new_nmatches_=-1 _new_list_lines_=-1
  private -i _avail_list_lines_=$((
      max( _autocomplete__max_lines - 1 - ${${_opts_[(i)-[Xx]]}:+1} - compstate[list_lines], 0 )
  ))

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

     (( funcstack[(I)_describe] )) ||
        return _ret_

    _displ_array_=$_opts_[-D]

    (( ${${(P)_displ_array_}[(I)*:*]} )) ||
        return _ret_

    _matches_array_=$_opts_[-O]
    if [[ -z $_matches_array_ ]]; then
      local -a _matches_=( "${(P)_displ_array_}" )
      _matches_array_=_matches_
    fi

    (( _avail_list_lines_ -= _autocomplete__reserved_lines ))
    _new_nmatches_=${(PA)#_displ_array_}
    _new_list_lines_=${#${(@u)${(PA)_displ_array_}[@]#*:}}
    (( _autocomplete__reserved_lines += _new_list_lines_ ))
  else
    _autocomplete__reserved_lines=0

    # (Note: Can't put comments inside expansions, without setting INTERACTIVE_COMMENTS.)
    private -a _out_=()
    _out_=( "${(0)"$(
        .autocomplete.compadd "$@"
        print -rNC1 -- "$compstate[list_lines]" "$compstate[nmatches]"
    )"}" )
    (( _new_list_lines_ = max( 0, $_out_[1] - $compstate[list_lines] ) ))
    (( _new_nmatches_ =   max( 0, $_out_[2] - $compstate[nmatches]   ) ))
  fi

  if (( _new_list_lines_ <= _avail_list_lines_ )); then
    if ! [[ -v _opts_[-D] ]]; then
      .autocomplete.compadd "$@"
      _ret_=$?
    fi
    return _ret_
  fi

  local -a _groupname_=()
  zparseopts -A _opts_ -D -E - k U d:=_dopt_ l=_dopt_ J:=_groupname_ V:=_groupname_
  set -- "$_groupname_[@]" "$@"

  if ! [[ -v _opts_[-D] ]]; then
    local -a _matches_=()
    .autocomplete.compadd -O _matches_ "$@" # Collect all matching completions.
    _matches_array_=_matches_
  fi
  _displ_array_=$_dopt_[-1]

  private -i _nmatches_per_line_=$(( 1.0 * _new_nmatches_ / _new_list_lines_ ))
  if [[ _nmatches_per_line_ -lt 1 && ! -v _opts_[-D] ]]; then
    # If we need more than one line per match, then make each match fit on one line.
    if [[ -z $_displ_array_ ]]; then
      local -a displ=( "${(P@)_matches_array}" )
      _dopt_=( -d displ )
      _displ_array_=displ
    fi
    _dopt_=( -l "$_dopt_[@]" )
    set -A $_displ_array_ ${(r:COLUMNS-1:@)${(P@)_displ_array_}//$'\n'/\n}
    _nmatches_per_line_=1
  fi

  # Need to round this down _before_ subtracting or it will be effectively rounded up.
  private -i _fit_=$(( max( _avail_list_lines_, 0 ) * _nmatches_per_line_ ))
  private -i _trim_=$(( max( 0, _new_nmatches_ - _fit_ ) ))

  if (( _trim_ )); then
    shift -p $(( min( _trim_, ${(PA)#_matches_array_} ) )) $_matches_array_
    [[ -n $_displ_array_ ]] &&
        shift -p $(( min( _trim_, ${(PA)#_displ_array_} ) )) $_displ_array_
    typeset -gH _autocomplete__partial_list=$curtag

    # If we have to trim off matches, that means we've run out of screen space.
    # So, disable all further completion.
    comptags() {
      return 1
    }
  fi

  if ! [[ -v _opts_[-D] ]]; then
    _autocomplete.compadd_opts_len "$@"
    .autocomplete.compadd $_dopt_ "$@[1,?]" -a -- $_matches_array_
    _ret_=$?
  fi

  return _ret_
}

typeset -gHa _autocomplete__log_functions=( $log_functions[@] )