#! /bin/sh
# next line is a comment to tcl, but not sh: \
    P=/opt/bin:$PATH ; for i in /opt/tcl* ; do P=$i/bin:$P ; done ; \
	PATH=$P ; export PATH ; exec tclsh $0 ${1+"$@"}

# pstree invokes the sys-v style ps, and then, IF the output contains
# columns PID and PPID, sorts them according to parentage.  Otherwise,
# the output is sent through unsorted.
#
# If we can find the CMD or COMMAND column, we indent that to indicate
# parentage.  Otherwise we indent from the left.
#
# See below for default on different OS's.

set env(PATH) "/usr/bin:/bin:/usr/sbin:/sbin:$env(PATH)"

proc usage {} {

    puts "Usage:
    $::prog \[select_args] \[\[-o] fmt_arg]
    or
    $::prog -

Sps stands for \"sorted ps\", and has four main functions:
    o  it runs ps and sorts the output to show parent-child relationships;
    o  it makes it easy to select process trees by pid or regexp pattern;
    o  this help page shows a lot of the useful output format codes;
    o  it knows sensible default arguments for several different OS's.

Sps passes most arguments through to ps for processing.  However, there
are a few special cases.

 o  If the argument -P ppp\[,...] is used, it means to show only the
    branch of the process tree (both ancestors and children) that includes
    `ppp'.  Here, ppp\[,...] is interpreted as a pid if ppp is a number;
    otherwise it is interpreted as a regexp pattern.  In the latter case, the
    pattern is applied to the entire output line from ps.  Matched lines are
    suffixed with \" *\", and the list of matched pid's is printed at the end of
    the output.  Example:
	 sps -P345,emacs
    prints the process trees for process 345 and all processes with \"emacs\"
    somewhere in the output from ps.  The sps exit code will be non-zero
    no processes match any entries from the ppp,... list.

 o  The argument -q means \"quiet\".  This only affects `-P ppp', in which
    case only the lines matching ppp are printed, but not the ps header nor
    the rest of the process tree branch.  The summary line at the bottom
    (with the pid's of the matched processes) is still printed.  Another q
    (`-qq' or `-q -q') increases the quietness: only the summary line of
    matched pid's is printed.

 o  If the only argument is \"-\", then stdin is assumed to
    contain a ps-style output, and sps simply sorts it.

 o  If one or more arguments are of the form ^xxx^yyy, then the arguments
    are used to edit the default ps command by replacing all occurrances
    of \"xxx\" with \"yyy\".  xxx is used as a regexp, so beware of wildcards.
    The edited command is printed to stderr. (N.B.  You can use \"^^\" to get
    the default command printed without anything being executed.)

 o  If one argument is \"args\", it is shorthand for the common use
    ^comm^args, ie generate the long \"args\" output (show command with all
    its arguments) instead of the short \"comm\" format (show the name of
    the command being executed).

Sps executes ps with the specified select_args, and then sorts the output
to show parent-child relationships.  The \"select_args\" are the ps arguments
that select which processes are displayed.  The default select_args vary with
the OS, and are `$::def_opt' for your system ($::system).
If PID and PPID columns aren't present, the output won't be sorted.  For
SysV-like ps commands, the additional option \[-o] fmt_arg is permitted; the
default fmt_arg is  \"$::format\".
Note: to recognize the fmt_arg argument, it must come last and contain
a comma.  The `-o' is optional.

You can override the main select_args and/or the fmt_arg values.
If you specify any arguments that start with `-', they are assumed to
be select_args, and override the default (yes, sps is smart enough to trim the
leading dash from Linux args).  If the last argument contains a comma, it is
assumed to be a fmt_arg, and overrides the default fmt_arg value.
N.B. On SunOS 4.x, ps has rather limited capabilities, and you should not 
normally override the default options.

A non-exhaustive list of useful fmt_arg fields include:
    user or ruser (effective or real UID)                   (posix)
    group or rgroup (effective or real GID)                 (posix)
    pid, ppid, pgid                                         (posix)
    pcpu (% cpu)                                            (posix)
    vsz (virtual mem size, KB)                              (posix)
    nice                                                    (posix)
    etime, stime, time (elapsed, start, and cpu time)       (posix)
    tty                                                     (posix)
    comm, args (argv\[0], argv\[0..argc-1])                   (posix)
    class, pri (scheduling class, priority)                 (solaris)
    osz (main mem image size, pages)                        (solaris)
    rss (resident set size, KB)                             (solaris)
    pmem (rss, as % of phys memory)                         (solaris)
    fname (1st 8 chars of basename of executable file)      (solaris)
    majflt,minflt,pagein (page faults,reclaims,#ops in)     (DU)
    msgrcv,msgsnd,nsignals (#msgs recv'd;# sent; #sigs)     (DU)
    addr,cls,flags,intpri (address, class, flags, priority) (HP-UX)
    prmid,prmgrp,state (PRM res. gp id, name; state)        (HP-UX)"
}

proc init {} {

    set ::quiet 0	; # default level of printing quietness.

    set sys [exec uname -s]
    set rev [exec uname -r]
    set ::system "$sys-$rev"

    # Architecture-specific adjustments
    set ::Allows_o 1	; # !0 means it has the -o format arg.
    set ::trimdash 0	; # !0 means to eliminate dashes from options.
    switch -glob -- $::system {
	SunOS-4.* {
	    set ::Allows_o 0
	    set ::def_opt "ajgx"
	    set ::pspath /usr/bin/ps
	}
	Linux-* {
	    set ::def_opt "-e"
	    set ::pspath /bin/ps
	}
	HP-UX-* {
	    set ::env(UNIX95) ""
	    set ::def_opt "-e"
	    set ::pspath /usr/bin/ps
	}
	default {
	    set ::def_opt "-e"
	    set ::pspath /usr/bin/ps
	}
    }
    set ::format "user,pid,ppid,vsz,pcpu,time,comm"
}

#
# Process the arglist.  Note: should NOT be called if argv == {-}.
#
proc doArgs {argv} {

    if {[string match -h* [lindex $argv 0]]} {
	usage
	exit
    }

    set subst_list {}	; # list of ^xxx^yyy args

    # Check for special -P nnn option.
    if {[set i [lsearch -regexp $argv "^-P.*"]] >= 0} {
	set arg [lindex $argv $i]
	if {[string length $arg] == 2} {
	    set ::distinguished [lindex $argv [expr $i + 1]]
	    set argv [concat [lrange $argv 0 [expr $i - 1]] \
				[lrange $argv [expr $i + 2] end]]
	} else {
	    set ::distinguished [string range $arg 2 end]
	    set argv [concat [lrange $argv 0 [expr $i - 1]] \
				[lrange $argv [expr $i + 1] end]]
	}
    }

    # Check for special -q and -qq option(s).
    while {[set i [lsearch -exact $argv "-q"]] >= 0} {
	incr ::quiet
	set argv [concat [lrange $argv 0 [expr $i - 1]] \
					    [lrange $argv [expr $i + 1] end]]
    }
    while {[set i [lsearch -exact $argv "-qq"]] >= 0} {
	incr ::quiet 2
	set argv [concat [lrange $argv 0 [expr $i - 1]] \
					    [lrange $argv [expr $i + 1] end]]
    }

    # Check for user-supplied options (other than fmt_args);
    # Check for ^xxx^yyy and elide them.
    set nopt 0
    set arglist {}
    foreach x $argv {
	if {[string compare $x "args"] == 0} {
	    # "args" is shorthand for ^comm^args
	    set x "^comm^args"
	}
	if {[regexp {^\^[^^]*\^} $x]} {
	    lappend subst_list [lrange [split $x "^"] 1 2]
	} else {
	    lappend arglist $x
	    if {[regexp {^-[^o]} $x]} {
		incr nopt
	    }
	}
    }
    set argv $arglist
    if !$nopt {
	# User didn't supply options.
	set argv [concat $::def_opt $argv]
    }

    # Check for user-supplied formats:
    if $::Allows_o {
	if {[string first "," [lindex $argv end]] == -1} {
	    # User didn't supply format string
	    lappend argv -o $::format
	} else {
	    # User did supply format string.  Ensure there's a -o before it.
	    set i [expr [llength $argv] - 2]
	    set argv [concat [lrange $argv 0 $i] -o [lindex $argv end]]
	}
    }

    if $::trimdash {
	set Argv {}
	foreach a $argv {
	    if {[string match -* $a]} {
		lappend Argv [string range $a 1 end]
	    } else {
		lappend Argv $a
	    }
	}
	set argv $Argv
    }
    # Trim out -o's w empty args -- this breaks some ps's.
    set Argv {}
    foreach a $argv b [lrange $argv 1 end] {
	if {"$a" == "-o" && [string match -* $b]} {
	    # Discard $a
	} else {
	    lappend Argv $a
	}
    }

    # Do the ^xxx^yyy editing.
    set ::cmd [concat "$::pspath" $Argv]
    if {[llength $subst_list]} {
	### puts "subst_list $subst_list"
	set unsubst_cmd $::cmd
	if {[lsearch -regexp $subst_list {^{}}] >= 0} {
	    puts stderr "Unsubstituted command:\n\t$unsubst_cmd"
	    puts stderr "Substitution strings can't have\
			    empty xxx in ^xxx^yyy."
	    exit 1
	}
	foreach xy $subst_list {
	    set x [lindex $xy 0]
	    set y [lindex $xy 1]
	    regsub -all $x $::cmd $y ::cmd
	}
    }
}

proc main {} {
    
    set ::prog [file tail $::argv0]

    #
    # Initialize the program with default values, etc.
    #
    init

    #
    # Special case: argv == {-} is handled differently from all the rest.
    #
    if {[llength $::argv] == 1 &&
    			[string compare [lindex $::argv 0] "-"] == 0} {
	# Parse stdin; don't execute command
	set handle stdin
	if {[gets $handle headline] < 0} {
	    puts stderr "$::prog: didn't even get one line from stdin"
	    exit 1
	}
    } else {
	#
	# Process all the user's options
	#
	doArgs $::argv

	#
	# Start the ps command
	#
	if {!$::quiet} {
	    puts stderr $::cmd
	}
	if {[catch {open "|$::cmd" r} handle]} {
	    puts stderr "$::prog: failed to exec ps: $handle"
	    exit 1
	}

	#
	# Get the "headline".
	#
	if {[gets $handle headline] < 0} {
	    # Didn't even get one line on stdout.
	    # Repeat the command to collect any stderr output...
	    catch {eval exec $::cmd 2>@stdout} result
	    if {"$result" == "" } {
		puts stderr \
		    "$::prog: didn't even get one line from ps; was executing"
		puts stderr "\t$::cmd"
	    } else {
		puts stderr "An error occurred while executing"
		puts stderr "\t$::cmd"
		puts stderr "The stderr output was:\n"
		puts -nonewline stderr "\t"
		puts [join [split $result "\n"] "\n\t"]
		exit 1
	    }
	}
    }

    #
    # Print the headline, if desired, and find the PID and PPID columns
    #
    process_headline $headline

    #
    # Special case if we don't have both PID and PPID.
    #
    if {$::col_PID == -1 || $::col_PPID == -1} {
	#
	# Don't have necessary info to construct parentage tree.
	# Just send data through to stdout.
	#
	while {[gets $handle line] >= 0} {
	    if {[catch {puts $line} result]} {
		if {[string compare [lindex $errorCode 0] POSIX] == 0} {
		    puts stderr [lindex $errorCode 2]
		} else {
		    puts stderr $result
		}
		exit 1
	    }
	}
	exit 0
    }

    #
    # Process the body of the ps output.
    #
    filter_body $handle

    #
    # If using -P ppp,..., then exit with non-zero status if no
    # processes were matched.
    #
    if {[info exists ::distinguished] && ! [info exists ::dist_and_present]} {
	exit 1
    } else {
	exit 0
    }
}

#
# Print the headline, if desired, and find the PID and PPID columns
#
proc process_headline {headline} {

    if {!$::quiet} {
	puts $headline
    }

    if {[set ::stridx_CMD [string first "CMD" $headline]] == -1} {
	set ::stridx_CMD [string first "COMMAND" $headline]
    }

    set ::headline_trimmed [string trim $headline]
    regsub -all "  +" $::headline_trimmed " " headline
    set columns [split $headline]
    set ::col_PID [lsearch $columns "PID"]
    set ::col_PPID [lsearch $columns "PPID"]
    set ::left_end [expr $::stridx_CMD - 1]
    if {$::left_end < 0} {
	set ::right_start 0
    } else {
	set ::right_start $::stridx_CMD
    }
}

proc filter_body {handle} {

    # OK.  We know the PID and PPID columns.  Filter away.
    set pslist {}
    foreach line [split [read -nonewline $handle] \n] {
	# Skip lines that look like the headline: assume that they are multiple
	# copies of it.
	set line_trimmed [string trim $line]
	if {[string compare $line_trimmed $::headline_trimmed] == 0} {continue}
	regsub -all "  +" $line_trimmed " " x
	set x [split $x]
	set pid [lindex $x $::col_PID]
	set ppid [lindex $x $::col_PPID]
	lappend pslist [list $pid $ppid $line]
	set ::psarray($pid) $line
    }

    # No output?
    if {[llength $pslist] == 0} {
	return
    }

    #
    # Sort into increasing ppid-order, subsort into increasing pid order,
    # then reorder into a tree
    #
    foreach entry [lsort -command sort_ppid_pid $pslist] {
	set pid [lindex $entry 0]
	set ppid [lindex $entry 1]
	lappend ::parent($pid) $ppid
	lappend ::children($ppid) $pid
    }

    #
    # Trim out all distinguished pid's or pats (ie those entered via
    # -P ppp,...) that aren't in its/their parent or children list.
    #
    if {[info exists ::distinguished]} {
	trim_distinguished
    }
    # Go through children().  Find ppid's that aren't also pid's.
    # Call them out as "toplevels".
    set toplevel {}
    foreach ppid [array names ::children] {
	if {![info exists ::psarray($ppid)]} {
	    lappend toplevel $ppid
	}
    }
    # If there aren't any toplevels, then pid 0 must be in the list,
    # as it is its own parent.
    if {![llength $toplevel]} {
	lappend toplevel 0
    }

    # If we have distinguished pid's, then only keep toplevels that are
    # either in the distinguished list (seen($pid) is set) or are parents
    # of same (trimmed($pid) is set).  (Note: if not all processes were listed,
    # e.g. ps -a is used instead of ps -e, then there can be processes
    # whose parents aren't listed in the ps output.  In that case, those
    # processes won't have been noticed when the distinguished branches
    # were selected.  Discard those unwanted processes now.)
    if {[info exists ::distinguished]} {
	set t {}
	foreach pid $toplevel {
	    if {[info exists seen($pid)] || [info exists ::trimmed($pid)]} {
		lappend t $pid
	    }
	}
	set toplevel $t
    }

    #
    # Print the filtered data, if the quiet level isn't too great.
    #
    if {$::quiet <= 1} {
	set nsibs [llength $toplevel]
	set thissib 0
	foreach pid $toplevel {
	    incr thissib
	    print_self_then_children "" 1 $thissib $nsibs $pid
	}
    }

    # If we looked for distinguished entries, print the list of pid's found.
    # Note that even if the list is empty, we print it so that a program
    # can simply take the last line as a list of matching pid's.
    if {[info exists ::distinguished]} {
	if {[info exists ::dist_and_present]} {
	    puts [lsort [array names ::dist_and_present]]
	} else {
	    puts ""
	}
    }
}

#
# Trim out all distinguished pid's or pats (ie those entered via
# -P ppp,...) that aren't in its/their parent or children list.
#
proc trim_distinguished {} {

    if {![info exists ::distinguished]} {
	return		; # sorry, there aren't any distinguished entries.
    }

    # Discard processes that aren't ancestors or children of the
    # distinguished processes.
    # First, make a list of distinguished pid's.
    set dist_pids {}
    foreach ppp [split $::distinguished " ,"] {
	if {[regexp "^\[ \t]*\$" $ppp]} {
	    # ignore whitespace
	    continue
	}  elseif {[regexp "\[0-9]+" $ppp]} {
	    lappend dist_pids $ppp
	} else {
	    # It's a pattern.  See if there are any entries that match it.
	    foreach pid [array names ::psarray] {
		if {[regexp $ppp $::psarray($pid)]} {
		    # It's a match
		    lappend dist_pids $pid
		}
	    }
	}
    }
    # Now harvest each distinguished pid's tree, and
    # modify the psarray to flag the distinguished pid entries.
    catch {unset seen}
    foreach pid $dist_pids {
	if {[info exists seen($pid)]} { continue }	; # skip duplicates
	set seen($pid) 1

	# Check that the ps output contains distinguished pid $pid
	# (Note that if $pid was specified on the command line, it will be
	# in $dist_pids, yet need not be an actual pid in the output of
	# the ps command.)
	if {![info exists ::parent($pid)]} { continue }

	# Don't consider ourselves or our child ps distinguished.
	if {$pid == [pid] || $::parent($pid) == [pid]} { continue }

	# If here, it's both distinguished and present.
	set ::dist_and_present($pid) 1

	while {[info exists ::parent($pid)] && $::parent($pid) != $pid} {
	    set ppid $::parent($pid)
	    if {![info exists ::trimmed($ppid)]} {
		# Trim children of parent to just this one pid.
		set ::children($ppid) $pid
		set ::trimmed($ppid) 1
	    } else {
		# If not yet in, append this pid to the
		# previously-trimmed list.
		if {[lsearch $::children($ppid) $pid] == -1} {
		    lappend ::children($ppid) $pid
		}
	    }
	    set pid $ppid
	}
    }
}

#
# Expect lines are 3-elt list of format:
#	pid  ppid {original text}
#
proc sort_ppid_pid {line1 line2} {
    set ppid1 [lindex $line1 1]
    set ppid2 [lindex $line2 1]

    if {$ppid1 == $ppid2} {
	set pid1 [lindex $line1 0]
	set pid2 [lindex $line2 0]
	return [expr $pid1 - $pid2]
    }
    return [expr $ppid1 - $ppid2]
}


proc print_self_then_children {leader amtoplevel thissib totsibs pid} {

    if {![info exists ::psarray($pid)]} {
	# It's a ppid, but it isn't in the actual ps output
	# (for example, we may have inferred the ppid from the ps output).
	# Don't print anything.
    } else {
	set left [string range $::psarray($pid) 0 $::left_end]
	set right [string range $::psarray($pid) $::right_start end]
	if $amtoplevel {
	    set my_ldr ""
	} else {
	    # set my_ldr "$leader \\_ "
	    # set my_ldr "$leader+ "
	    set my_ldr "$leader+-"
	}


	#
	# Don't print unless:
	# o  quiet == 0, or
	# o  quiet == 1 and self is a distinguished pid.
	# In the former case, print with full indenting; in the latter
	# case, ignore leader and indentation of the command.
	#
	if {$::quiet == 0} {
	    # Don't flag ourselves or our child ps
	    if {[info exists ::dist_and_present($pid)]} {
		append right " *"
	    }
	    if {[catch {puts "$left$my_ldr$right"} result]} {
		if {[string compare [lindex $::errorCode 0] POSIX] == 0} {
		    puts stderr [lindex $::errorCode 2]
		} else {
		    puts stderr $result
		}
		exit 1
	    }
	} elseif {$::quiet == 1 && [info exists ::dist_and_present($pid)]} {
	    if {[catch {puts "$left $right"} result]} {
		if {[string compare [lindex $::errorCode 0] POSIX] == 0} {
		    puts stderr [lindex $::errorCode 2]
		} else {
		    puts stderr $result
		}
		exit 1
	    }
	}
    }

    if {[info exists ::children($pid)]} {
	set nchildren [llength $::children($pid)]
	if {$totsibs > 1 && $thissib != $totsibs} {
	    # set child_ldr "$leader |  "
	    set child_ldr "$leader| "
	} elseif {$amtoplevel} {
	    # set child_ldr "$leader"
	    set child_ldr "$leader"
	} else {
	    # set child_ldr "$leader    "
	    set child_ldr "$leader  "
	}
	set ichild 0
	foreach child $::children($pid) {
	    incr ichild
	    # avoid the infinite loop of process 0 being its own parent:
	    if {$child != $pid} {
		print_self_then_children $child_ldr 0 $ichild $nchildren $child
	    }
	}
    }
}

main