#define _PSTAT64
#include <sys/param.h>
#include <sys/pstat.h>
#include <sys/unistd.h>
#include <malloc.h>
#include <stdio.h>
#include <stdlib.h>
#include <pwd.h>

typedef struct shared_segment_struct {
    long long pst_space;
    long long pst_vaddr;
    int refs;
    struct shared_segment_struct *next;
} segment;

static segment *shared_segs = NULL;

void pstatvm(uid_t uid, int pid)
{
    struct pst_vm_status pst;
    int idx, count;
    long long shared_vm = 0;
    long long shared_ram = 0;
    long long private_vm = 0;
    long long private_ram = 0;

    idx=0;
    count = pstat_getprocvm(&pst, sizeof(pst), (size_t)pid, idx);
    while (count > 0) {
        switch ((long)pst.pst_type) {
            case PS_IO: break;
                /* Don't count IO space.  It really is not RAM or swap. */
            default: 
                if (pst.pst_flags & PS_SHARED) {
                    segment *s;
                    int refs = 1;
                    for (s=shared_segs; s; s=s->next) {
                        if (s->pst_space == pst.pst_space
                        && s->pst_vaddr == pst.pst_vaddr) {
                            refs = s->refs;
                            break;
                        }
                    }
                    shared_vm += (long long) pst.pst_length;
                    shared_ram += (long long)pst.pst_phys_pages/refs;
                } else {
                    private_vm += (long long) pst.pst_length;
                    private_ram += (long long)pst.pst_phys_pages;
                }
                break;
        }

        idx++;
        count = pstat_getprocvm(&pst, sizeof(pst), (size_t)pid, idx);
    }
    printf("%6d ", uid);
    printf("%6d ", pid);
    printf("%11lldK ", shared_vm*4);
    printf("%11lldK ", shared_ram*4);
    printf("%11lldK ", private_vm*4);
    printf("%11lldK ", private_ram*4);
    printf("%11.1fM\n", (shared_ram+private_ram)/256.0);
}

void pstatvm_uid(uid_t uid)
{
#define BURST ((size_t)10)

    struct pst_status pst[BURST];
    int i, count;
    int idx = 0; /* index within the context */

    /* loop until count == 0, will occur when all have been returned */
    while ((count = pstat_getproc(pst, sizeof(pst[0]), BURST, idx)) > 0)
    {
        /* got count (max of BURST) this time.  process them */
        for (i = 0; i < count; i++) {
            if (pst[i].pst_pid==0) continue; /* Can't getprocvm real pid 0 */
            if (pst[i].pst_uid==uid) {
                pstatvm(uid, pst[i].pst_pid);
            }
        }

        /*
         * now go back and do it again, using the next index after
         * the current 'burst'
         */
        idx = pst[count-1].pst_idx + 1;
    }

    if (count == -1)
        perror("pstat_getproc()");

#undef BURST
}

void pstat_refcount_all(void)
{
#define BURST ((size_t)10)

    struct pst_status pst[BURST];
    int i, count;
    int idx = 0; /* index within the context */

    /* loop until count == 0, will occur when all have been returned */
    while ((count = pstat_getproc(pst, sizeof(pst[0]), BURST, idx)) > 0)
    {
        /* got count (max of BURST) this time.  process them */
        for (i = 0; i < count; i++) {
            struct pst_vm_status vm;
            int reg, count;
            if (pst[i].pst_pid==0) continue; /* Can't getprocvm real pid 0 */
            reg=0;
            while (pstat_getprocvm(&vm, sizeof(vm), pst[i].pst_pid, reg++)) {
                segment *s;
                for (s=shared_segs; s; s=s->next) {
                    if (s->pst_space == vm.pst_space
                    && s->pst_vaddr == vm.pst_vaddr) {
                        s->refs += 1;
                        break;
                    }
                }
                if (!s) {
                    s = (segment *) malloc(sizeof(segment));
                    s->pst_space = vm.pst_space;
                    s->pst_vaddr = vm.pst_vaddr;
                    s->refs = 1;
                    s->next = shared_segs;
                    shared_segs = s;
                }
                reg++;
            }
        }
        idx = pst[count-1].pst_idx + 1;
    }

    if (count == -1)
        perror("pstat_getproc()");

#undef BURST
}

int main(int argc, char *argv[])
{
    int i;
    pstat_refcount_all();
    printf("   uid    pid    shared_vm   shared_ram   private_vm  private_ram      Res_Mem\n
");
    if (argc > 1) {
        for (i=1; i<argc; i++) {
            struct passwd *p = getpwnam(argv[i]);
            if (p) {
                pstatvm_uid(p->pw_uid);
            } else {
                pstatvm_uid(atoi(argv[i]));
            }
        }
    } else {
        pstatvm_uid(getuid());
    }
    return 0;
}