#compdef dtrace

local curcontext="$curcontext" ret=1
local -a state line args

case $OSTYPE in
  ^darwin*)
    args=(
      '-32[generate 32-bit D programs and ELF files]'
      '-64[generate 64-bit D programs and ELF files]'
      '(-A -h -l -V)-G[generate an ELF file containing embedded dtrace program]'
    )
  ;;
  darwin*)
    args=(
      '*-W[wait for the named process to launch]:name'
    )
  ;;
esac

_arguments -C -s $args \
  '-a[claim anonymous tracing state]' \
  '(-G -h -l -V)-A[generate driver.conf(4) directives for anonymous tracing]' \
  '-b+[set trace buffer size]:size' \
  '*-c+[run specified command and exit upon its completion]:command' \
  '-C[run cpp(1) preprocessor on script files]' \
  '*-D+[define symbol when invoking preprocessor]:name' \
  '-e[exit after compiling request but prior to enabling probes]' \
  '-f+[enable or list probes matching the specified function name]:function:->functions' \
  '-F[coalesce trace output by function]' \
  '-H[print included files when invoking preprocessor]' \
  '(-A -G -l -V)-h[generate a header file]' \
  '*-i+[enable or list probes matching the specified probe id]:probe-id' \
  '-I+[add include directory to preprocessor search path]:path:_directories' \
  '(-A -G -h -V)-l[list probes instead of enabling them]' \
  '-L+[add directory to search path for DTrace libraries]:path:_directories' \
  '*-m+[trace or list probes matching the specified module name]:module:->modules' \
  '*-n+[trace or list probes matching the specified probe name]:name:->probes' \
  '-o+[set output file]:output file:_files' \
  '-p+[grab specified process-ID and cache its symbol tables]:pid:_pids' \
  '*-P+[trace or list probes matching the specified provider name]:provider:->providers' \
  '-q[set quiet mode (only output explicitly traced data)]' \
  '*-s+[enable or list probes according to the specified D script]:script file:_files' \
  '-S[print D compiler intermediate code]' \
  '*-U+[undefine symbol when invoking preprocessor]:name' \
  '-v[set verbose mode (report program stability attributes)]' \
  '(-A -G -h -l)-V[report DTrace API version]' \
  '-w[permit destructive actions]' \
  '*-x+[enable or modify compiler and tracing options]: : _values option
    "aggrate[rate of aggregation reading]\:time"
    "aggsize[aggregation buffer size]\:size"
    "ustackframes[number of user stack frames]\:scalar"
    "bufpolicy[specify the buffer policy for the principal buffer]\:policy\:(fill switch ring)"
    "bufresize[buffer resizing policy]\:policy\:(auto manual)"
    "bufsize[size of the per-CPU principal buffer]\:size"
    "cleanrate[cleaning rate (hz)]\:time"
    "cpu[specify the CPU on which to enable tracing]\:scalar"
    "defaultargs[allow references to unspecified macro arguments]"
    "destructive[allow destructive actions]"
    "dynvarsize[size of the dynamic variable space]\:size"
    "flowindent[turn on flow indentation]"
    "grabanon[claim anonymous state]"
    "jstackframes[number of default stack frames for jstack()]\:scalar"
    "jstackstrsize[default string space size for jstack()]\:scalar"
    "nspec[number of speculations]\:scalar"
    "quiet[set quiet mode]"
    "specsize[size of the speculation buffer]\:size"
    "strsize[maximum size of strings]\:size"
    "stackframes[maximum number of kernelspace stack frames to unwind for stack()]\:scalar"
    "stackindent[whitespace characters to use when indenting stack() and ustack() output]\:scalar"
    "statusrate[rate of status checking]\:time"
    "switchrate[rate of buffer switching]\:time"
    "ustackframes[maximum number of userspace stack frames to unwind for ustack()]\:scalar"' \
  '-X+[specify ISO C conformance settings for preprocessor]:ISO C conformance:((a\:"ISO plus K&R extensions (default)" c\:"Strictly conformant ISO C" s\:"K&R C only" t\:"ISO plus K&R extensions"))' \
  '-Z[permit probe descriptions that match zero probes]' && return

if [[ -n $state ]]; then
  local -a fields=( providers modules functions probes )
  local -a suf=( : : : )
  typeset -ga _cache_dtrace_probes
  (( $#_cache_dtrace_probes )) || _cache_dtrace_probes=(
      ${${${${(f)"$(_call_program dtrace-probes dtrace -l)"}[2,-1]}:#*[:[]}/(#b) #<-> #([^ ]#) #([^ ]#) #([^ ]#) ##([^ ]##)/${match[1]}:${match[3]:+$match[2]}:${match[3]:-$match[2]}:$match[4]}
  ) # filtering out those containing : and [, they occur on macos but may just be unmangled forms and not directly usable
  suf[${fields[(i)$state]}]=( "${compstate[quote][-1]} " ) # field matching the state gets a space suffix
  while compset -P 1 '*:'; do
    shift fields # each already listed field reduces one being possible
  done
  _tags dtrace-${^fields[1,(r)$state]}
  while _tags; do
    _requested dtrace-probes expl 'probe name' compadd -S "$suf[4]" - ${${(M)_cache_dtrace_probes:#((#s)|*:)${IPREFIX}[^:]#(#e)}##*:} && ret=0
    _requested dtrace-functions expl 'function' compadd -S "$suf[3]" - ${${${(M)_cache_dtrace_probes:#((#s)|*:)${IPREFIX}[^:]#:[^:]#(#e)}%:*}##*:} && ret=0
    _requested dtrace-modules expl 'module' compadd -S "$suf[2]" - ${${${(M)_cache_dtrace_probes:#((#s)|*:)${IPREFIX}[^:]#:[^:]#:[^:]#(#e)}#*:}%%:*} && ret=0
    _requested dtrace-providers expl 'provider' compadd -S "$suf[1]" - ${_cache_dtrace_probes%%:*} && ret=0

    (( ret )) || return 0
  done
fi

return 1