luamem.c

#include <stdio.h>
#include <string.h>
#include <inttypes.h>
#include <stdlib.h>

#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>

#ifdef __APPLE__
#include <mach/mach_time.h>
#else
#include <time.h>
#endif

typedef struct {
    uint64_t time;
    uint64_t bytes;
}
tracking_t;

typedef struct {
    union {
        uint64_t size;
        char padding[16];
    }
    info;

    char block[0];
}
header_t;

static tracking_t s_tracking[128 * 1024 * 1024];
static int s_count;
static uint64_t s_bytes;

#ifdef __APPLE__
// Apple's high-performance counter needs this
static mach_timebase_info_data_t s_tb = {0, 0};
#endif

uint64_t get_time() {
#ifdef __APPLE__
    uint64_t time = mach_absolute_time();
    return time * s_tb.numer / s_tb.denom;
#else
    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);
    return (uint64_t)now.tv_sec * 1000000000ULL + now.tv_nsec;
#endif
}

static void* customAlloc(void* ud, void* ptr, size_t osize, size_t nsize) {
    tracking_t* t;
    header_t* h;
    header_t* p;

    if (s_count == sizeof(s_tracking) / sizeof(s_tracking[0])) {
        fprintf(stderr, "s_tracking is too small\n");
        exit(1);
    }

    t = s_tracking + s_count++;
    t->time = get_time();

    if (nsize == 0) {
        if (ptr == NULL) {
            /* Uhhh... */
        }
        else {
            h = (header_t*)ptr - 1;
            s_bytes -= h->info.size;
            free(h);
            ptr = NULL;
        }
    }
    else {
        if (ptr == NULL) {
            h = (header_t*)malloc(nsize + sizeof(header_t));
            h->info.size = nsize;
            s_bytes += nsize;
            ptr = h->block;
        }
        else {
            h = (header_t*)malloc(nsize + sizeof(header_t));
            h->info.size = nsize;
            s_bytes += nsize;

            p = (header_t*)ptr - 1;
            s_bytes -= p->info.size;

            memmove(h->block, p->block, osize < nsize ? osize : nsize);
            free(p);
            ptr = h->block;
        }
    }

    t->bytes = s_bytes;
    return ptr;
}

static void dumpTracking(void) {
    tracking_t* t = s_tracking;
    FILE* log;
    int i;

    log = fopen("lua.log", "w");

    for (i = 0; i < s_count; i++, t++) {
        fprintf(log, "%f\t%f\n", (t->time - s_tracking[0].time) / 1000000.0, t->bytes / (1024.0 * 1024.0));
    }

    fclose(log);
}

int main(int argc, const char* argv[]) {
    lua_State* L;
    int top;

    if (argc != 2) {
        fprintf(stderr, "USAGE: luamem filename.lua");
        return 1;
    }

    s_count = 0;
    s_bytes = 0;

#ifdef __APPLE__
    mach_timebase_info(&s_tb);
#endif

    L = lua_newstate(customAlloc, NULL);

    top = lua_gettop(L);
    luaL_openlibs(L);
    lua_settop(L, top);

    if (luaL_loadfilex(L, argv[1], "t") != LUA_OK) {
        fprintf(stderr, "Error loading file: %s\n", lua_tostring(L, -1));
        lua_close(L);
        return 1;
    }

    if (lua_pcall(L, 0, 0, 0) != LUA_OK) {
        fprintf(stderr, "Error running file: %s\n", lua_tostring(L, -1));
        lua_close(L);
        return 1;
    }

    lua_close(L);
    dumpTracking();
    return 0;
}