#!/usr/bin/env bash

# Memstat.sh is a shell script that calculates linux memory usage
# for each program or arguments passed script only for that pid list.

# Original author Bobbin Zachariah
# Full rewroted by Vladimir Oleynik <dzo@simtreas.ru> (c) 2022
#  ftp://ftp.simtreas.ru/pub/my/memstat.sh
#
#  - do not use temporary files
#  - do not use cat/awk/paste/bc/cut/rm/ls/sort/uniq/wc/id utilities
#  - allow pid list
#  - allow non-root usage
#  - more fast and clear
#  If like my qsort-bash, do not destroy my (c) while coping, please!

# Script outputs shared and private memory for each program running in linux.
# Since memory calculation is bit complex,
# this shell script tries best to find more accurate results.
# Script use 2 files ie /proc/PID/status (to get name of process)
# and /proc/PID/smaps for memory statistic of process.

# Used external `grep' utility.
# Yes, I tested with "internally" `grep', but it very very slow!


# Global variables
PRG=${0##*/} Err=0 Err_Perm=1 Err_pid_NN=2 Err_pid_NF=4 Err_Proc=8
declare -Ai A_Name
declare -a  Name PIDs
declare -ai Private RAM Count
declare -i  Total i
declare     nn a p

echoe() {
	echo "$PRG: $1" >&2
	(( Err |= $2 ))
}

# Ohh, /proc/$PID/smaps have r--r--r-- mode but cat not read may be for root too
# next two functions call by the pointer 'fp' for fast run with different:
# call for all process or specified pid list.

declare -f fp

# for all process: disallow zero PSS, run without errors
dzPSS() {
	[[ s -ne 0 ]]
}

# for pid list: allow zero PSS (kernel process) but if PSS is not set then show error
azPSS_E() {
  if [[ -z $s ]]; then
	echoe "read /proc/$1/smaps permission denied, program name: '$name'" Err_Perm
	return 1
  fi
  return 0
}

# This function will count memory statistic for passed PID=$1
get_process_mem ()
{
  local name s r=0

  # get the process name
  read nn name < <(grep -m1 -s "^Name:" /proc/$1/status)
  [[ -z $name ]] && return 1

  # here we count memory usage with paste uniq process name

  # summarize the Pss and Private* values from /proc/$1/smaps
  # The trick for out variable name: ${"Pss":1:1} is 's' and ${"Private":1:1} is 'r'
  while read nn a p; do
	(( ${nn:1:1}+=a ))
  done < <(grep -s -e "^Pss:" -e "^Private" /proc/$1/smaps)
  $fp $1 || return 1

  # make a uniq process name
  if [[ -z ${A_Name["$name"]} ]]; then
	  A_Name["$name"]=$((i=${#Name[*]}))
	  Name+=("$name")
	  PIDs+=($1)
  else
	  i=${A_Name["$name"]}
	  PIDs[i]+=",$1"
  fi
  Count[i]+=1
  RAM[i]+=s
  Total+=s
  Private[i]+=r
}


# compare function.
q_cmp() {
	[[ ${RAM[$1]} -gt ${RAM[$2]} ]]
}

# internal function for qsort, sort left part of array
_qsortl() {
   local -i n k t=l[0] x=()

   for i in ${l[@]:1:$1-1}; do
	if q_cmp $i $t; then
		x[n++]=i
	else
		l[k++]=i
	fi
   done
   case ${n:-1} in
    1) ;;
    2) q_cmp ${x[1]} ${x[0]} && { n=x[0]; x[0]=x[1]; x[1]=n; };;
    *) _qsort $n;;
   esac
   case $k in
   '') l=(${x[@]} t);;
    1) l=(${x[@]} t ${l[0]});;
    2) q_cmp ${l[1]} ${l[0]}; l=(${x[@]} t ${l[$?^1]} ${l[$?]});;
    *) _qsortl $k; l=(${x[@]} t ${l[@]});;
   esac
}

# internal function for qsort
_qsort() {
   local -i n k t=x[0] l=()

   for i in ${x[@]:1:$1-1}; do
	if q_cmp $i $t; then
		x[n++]=i
	else
		l[k++]=i
	fi
   done
   case ${k:-1} in
    1) ;;
    2) q_cmp ${l[1]} ${l[0]} && { k=l[0]; l[0]=l[1]; l[1]=k; };;
    *) _qsortl $k;;
   esac
   case $n in
   '') x=(t ${l[@]});;
    1) x=(${x[0]} t ${l[@]});;
    2) q_cmp ${x[1]} ${x[0]}; x=(${x[$?^1]} ${x[$?]} t ${l[@]});;
    *) _qsort $n; x+=(t ${l[@]});;
   esac
}

# outname: x - is integer indexs of sorted input array values
qsort() {
	declare -gai x=(${!RAM[@]})
	i=${#x[@]}
	[[ i -gt 1 ]] && _qsort $i
}

# this function make conversion integers from kilobytes to Mb or Gb
# do integer manipulations for make the float rezult with a 3-digits factor
make_out()
{
  local out

  for a; do
      if [[ $a != *[^0-9]* ]]; then
	if (( a > 1024*1024 )); then
		# bc <<< "scale=3;${a}/1024.0/1024.0"
		(( nn = a / 1024 / 1024 ))
		printf -v a "%4d.%03d GB" $nn $(( (( a - nn * 1024 * 1024 ) * 1000 ) / 1024 / 1024 ))
	elif [[ a -gt 1024 ]]; then
		# bc <<< "scale=3;${a}/1024.0"
		(( nn = a / 1024 ))
		printf -v a "%4d.%03d MB" $nn $(( (( a - nn * 1024 ) * 1000 ) / 1024 ))
	else
		printf -v a "%4d.000 kB" $a
	fi
      fi
      out+=$a
  done
  printf '%s\n' "$out"
}


# main()

if [[ ! -d /proc ]]; then
	echoe "This script require Linux with mounted /proc" Err_Proc
	exit $Err
fi

# if arguments passed script will show statistic only for that pid list,
# or not we list all processes in /proc/ and get statistic for all of them

if [[ $# -eq 0 ]]; then
	fp=dzPSS
	for p in /proc/[1-9]*; do
		[[ ${p#/proc/} != *[^0-9]* ]] && get_process_mem ${p#/proc/}
	done
else
	fp=azPSS_E
	for p; do
		if [[ $p == *[^0-9]* ]]; then
			echoe "'$p' is not PID number" Err_pid_NN
			continue
		fi
		if [[ ! -d /proc/$p ]]; then
			echoe "PID=$p do not found" Err_pid_NF
			continue
		fi
		get_process_mem $p
	done
fi

# now we print result, first header
[[ ${#Name[@]} -gt 0 ]] && echo -e "   Private    +     Shared     =    RAM used \t Program \t\tPID"

# sort by RAM and do output
qsort
for i in ${x[@]}; do
	printf -v p "%-16s" "${Name[i]}"
	[[ Count[i] -gt 1 ]] && p+=" (${Count[i]})"
	make_out ${Private[i]} $' \t ' \
			$((RAM[i]-Private[i])) $' \t ' \
			${RAM[i]} $' \t '"$p" $' \t'${PIDs[i]}
done

[[ ${#Name[@]} -gt 1 ]] && \
make_out $'--------------------------------------------\n\t\t\t\t ' \
	$Total $'\n'"============================================"
exit $Err