/* memtrck.c: implementation of a simple memory tracker wrapper */
/*
 * [2000.07.17] Memory Tracker v1.0 by Jonathan Mayo <jmayo@telocity.net>
 * 
 * THIS FILE IS PUBLIC DOMAIN. YOU ARE FREE TO USE IT FOR ANY PURPOSE.
 *
 */

/* Updated 2007.07.02 : Jonathan Mayo <jon.mayo@gmail.com>
 *
 *
 *
 */

#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>
#include "memtrck.h"

#ifndef NR
#define NR(x) (sizeof(x)/sizeof*(x))
#endif

struct meminfo {
    unsigned signature;
    size_t allocated_size;
    unsigned long generation;
    const char *origin; /* calling functions name */
    struct meminfo *next, **prev;
};

static struct meminfo *head, **last=&head;
static unsigned long generation=1;
static unsigned long histogram[sizeof(size_t) * CHAR_BIT];
/* TODO: graph of size by log2 of generation */

/* calculate the integer log2 of size_t */
static unsigned ilog2(size_t n) {
    unsigned y;
    int i, c;
    i=sizeof n * CHAR_BIT;
    c=i/2;
    do {
        y=n>>c;
        if(y) {
            i=i-c;
            n=y;
        }
        c>>=1;
    } while(c);
    return (sizeof n * CHAR_BIT) - i;
}

static void list_add(struct meminfo *meta, size_t size) {
    /* add this entry to end of linked list */
    meta->signature=0x77779999;
    meta->allocated_size=size;
    meta->next=0;
    meta->prev=last;
    meta->generation=generation++;
    meta->origin=0;
    *last=meta;
    last=&meta->next;
    histogram[ilog2(size)]++; /* add new allocation to the list */
}

static void list_delete(struct meminfo *meta) {
    /* unhitch this entry from the list */
    if(last==&meta->next) {
        last=meta->prev;    
    }
    *meta->prev=meta->next;
    memset(meta, 0, sizeof *meta);
}

void *dbgMalloc(size_t size)
{
    void *ret;
    assert(size+sizeof(struct meminfo) > size); /* check for overflow */
    ret=malloc(size + sizeof(struct meminfo));
    if(!ret) return 0;
    list_add(ret, size);
    return ret + sizeof(struct meminfo);
}

void *dbgCalloc(size_t nmemb, size_t size)
{
    void *ret;
    size*=nmemb;
    ret=dbgMalloc(size);
    if(ret) {
        memset(ret, 0, size);
    }
    return ret;
}

void dbgFree(void *ptr)
{
    if(ptr) {
        ptr=ptr-sizeof(struct meminfo); /* point to meminfo header */
        assert(((struct meminfo*)ptr)->signature==0x77779999);
        list_delete(ptr);
        free(ptr);
    }
}

void *dbgRealloc(void *ptr, size_t size)
{
    void *ret;

    assert(size+sizeof(struct meminfo) > size); /* check for overflow */

    if(ptr) {
        ptr=ptr-sizeof(struct meminfo); /* point to meminfo header */
        assert(((struct meminfo*)ptr)->signature==0x77779999);
    }
    ret=realloc(ptr, size + sizeof(struct meminfo));
    if(!ret) return 0;
    if(!ptr) { /* behaves like a malloc */
        list_add(ret, size);
    } else {
        struct meminfo *meta = ret;
#if 0
        /* enable this if realloc should be tracked in the histogram as a
         * single item. disable if realloc is added to the histogram each time
         * it is modified.
         */
        histogram[ilog2(meta->allocated_size)]--; /* remove old count */
#endif
        meta->allocated_size=size; /* updated size */
        meta->origin=0; /* updated origin */
        histogram[ilog2(size)]++; /* update histogram */
    }
    return ret + sizeof(struct meminfo);
}

/**
 * hardcoded to draw no more than 24 rows */
void dbgStatisticsShow(void) {
    unsigned largest;
    int i, y, rows, scale;

    /* scale the graph */
    largest=0;
    for(i=0;i<NR(histogram);i++) {
        if(largest<histogram[i]) largest=histogram[i];
    }
    scale=largest/(24-6);
    rows=largest/scale;

    /* draw the graph */
    for(y=0;y<rows;y++) {
        for(i=0;i<NR(histogram);i++) {
            putchar(histogram[i]<=(rows-y)*scale ? ' ' : '#');
        }
        putchar('\n');
    }

    /* label the entries */
    for(scale=10000;scale;scale/=10) {
        for(i=0;i<NR(histogram);i++) {
            char ch;
            if(10000*10<=histogram[i]) {
                ch='-';
            } else if(scale>histogram[i]) {
                ch=' ';
            } else {
                ch="0123456789"[(histogram[i]/scale)%10];
            }
            putchar(ch);
        }
        putchar('\n');
    }

    /* TODO: count the leaked memory */
}

#if 1
#include <stdio.h>
/** TEST **/
#define NR_TEST 32768
int main() {
    void *test[NR_TEST];
    unsigned i, curr;
    
    /* allocate the test data and reallocate it */
    for(i=0;i<NR_TEST;i++) {
        test[i]=dbgMalloc(rand()%4096);
        test[i]=dbgRealloc(test[i], rand()%65536);
    }
    
    /* free entries randomly, shuffling in the end slot to filling gaps */
    for(i=NR_TEST;i!=0;i--) {
        curr=rand()%i;
/*        fprintf(stderr, "curr=%u i=%u\n", curr, i); */
        assert(test[curr] != NULL);
        dbgFree(test[curr]);
        test[curr]=test[i-1];
        test[i-1]=0;
    }
    
    /* verify the test */
    for(i=0;i<NR_TEST;i++) {
        assert(test[i] == NULL);
    }

    dbgStatisticsShow();
    return 0;
}
#endif