(Simple) Lua Allocation Tracker

ijlat.h

/* ijlat.h - v0.1a - (Simple) Lua Allocation Tracker -- public domain -- @incrediblejr, May 2016

ABOUT

A simple Lua allocation tracker.

USAGE

The memory tracker is implemented as a [stb-style header-file library][1] which means that
in *ONE* source file, put:

#define IJLAT_IMPLEMENTATION

// check out @FUNC_CUSTOMIZATION for providing own functions(if dependency on stdlib
// is undesirable for example)

// check @CONFIG and define one of :
// IJLAT_COMMIT_IN_HOOK / IJLAT_COMMIT_IN_ALLOC / IJLAT_IGNORE_MIKE_PALL
// like :
// #define IJLAT_COMMIT_IN_HOOK 1
#include "ijlat.h"

EXAMPLE
#include <stdlib.h>

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

#define IJLAT_STATIC
//#define IJLAT_COMMIT_IN_ALLOC 1
#define IJLAT_IGNORE_MIKE_PALL 1
//#define IJLAT_COMMIT_IN_HOOK 1
#define IJLAT_IMPLEMENTATION
#include "ijlat.h"

#include <assert.h>

static int qsort_ijlat_alloc_trace_comp(const void *a, const void *b)
{
	const ijlat_alloc_trace *ta = a;
	const ijlat_alloc_trace *tb = b;
	if (ta->tallocated == tb->tallocated) {
		if (ta->nallocations == tb->nallocations)
			return 0;
		if (ta->nallocations > tb->nallocations)
			return 1;
		return -1;

	}
	else if (ta->tallocated > tb->tallocated)
		return 1;
	return -1;
}

static ijlat gtraces_module;
static int gtraces_module_sorted = 0;

static int lua_ijlat_init(lua_State* L)
{
	ijlat_init(&gtraces_module, L);
	gtraces_module_sorted = 0;
	return 0;
}

static int lua_ijlat_report(lua_State* L)
{
	int back_is_empty = ijlat_is_backindex_and_empty(&gtraces_module, gtraces_module.pending_index);

	(void)L;
	if ((back_is_empty && gtraces_module.tused == 1) || !gtraces_module.tused) {
		printf("_no_ allocations\n");
	}
	else {
		unsigned i, n = gtraces_module.tused - (back_is_empty ? 1 : 0);
		char *sb = gtraces_module.s;
		ijlat_alloc_trace *traces = gtraces_module.t;
		if (!gtraces_module_sorted) {
			qsort(traces, n, sizeof *traces, qsort_ijlat_alloc_trace_comp);
			gtraces_module_sorted = 1;
		}
		for (i = 0; i != n; ++i) {
			ijlat_alloc_trace *current = traces + i;
			printf("allocation #%u. num allocations %u. bytes allocated %u. callstack :\n%s",
				i, current->nallocations, current->tallocated, sb + current->soffset);
		}
	}
	return 0;
}

static int lua_ijlat_shutdown(lua_State* L)
{
	(void)L;
	ijlat_shutdown(&gtraces_module);
	return 0;
}

LICENSE
This software is dual-licensed to the public domain and under the following
license: you are granted a perpetual, irrevocable license to copy, modify,
publish, and distribute this file as you see fit.

[1]: https://github.com/nothings/stb

*/

#ifndef INCLUDE_IJLAT_H
#define INCLUDE_IJLAT_H

#ifdef IJLAT_STATIC
#define IJLAT_DEF static
#else
#define IJLAT_DEF extern
#endif

#ifdef __cplusplus
extern "C" {
#endif

	typedef struct ijlat_alloc_trace {
		unsigned hash;
		unsigned nallocations;
		unsigned tallocated;
		unsigned soffset;
	} ijlat_alloc_trace;

	typedef struct lua_State lua_State;

	typedef struct ijlat {
		lua_State *L;
		unsigned talloc;
		unsigned tused;
		ijlat_alloc_trace *t;

		unsigned salloc;
		unsigned sused;
		char *s;

		unsigned pending_index;
	} ijlat;

	IJLAT_DEF void ijlat_init(ijlat *self, lua_State *L);
	IJLAT_DEF void ijlat_shutdown(ijlat *self);

	/* reset all buffers (_not_ free:ing them) */
	#define ijlat_reset(self) ((self)->tused = 0, (self)->sused = 0, (self)->pending_index = (unsigned)-1)

	#define ijlat_is_backindex(self, index) ((self)->tused && (index) == (self)->tused - 1)
	#define ijlat_is_backindex_and_empty(self, index) (ijlat_is_backindex((self), (index)) ? !(self)->t[(index)].nallocations : 0)

#ifdef __cplusplus
}
#endif

#endif

#ifdef IJLAT_IMPLEMENTATION

/* I _know_ I could make this a bit cuter by storing original hook+alloc+L and anchor
self in registry to avoid globals. */

#ifdef IJLAT_USER_CONFIG
#include IJLAT_USER_CONFIG
#endif

/* @CONFIG */

/* IJLAT_COMMIT_IN_HOOK
allocation size is recorded in wrapped allocation function and stored in next
hook callback. this way we only have to 'serialize' the callstack when a allocation
has occurred but gives very wrong results.
*/


/* IJLAT_COMMIT_IN_ALLOC
callstack is _always_ 'serialized' in the hook callback and the latest callstack
is stored/committed when a allocation occurs. gives the correct result but is very
slow/wasteful.
*/

/* IJLAT_IGNORE_MIKE_PALL
callstack is 'serialized' and stored in the wrapped allocation function, ignoring
Mike Pall's advice. use at own risk :)
http://www.freelists.org/post/luajit/Calling-lua-getinfo-in-allocation-callback-occasionally-crashes,5
*/


#if IJLAT_IGNORE_MIKE_PALL
#undef IJLAT_COMMIT_IN_HOOK
#undef IJLAT_COMMIT_IN_ALLOC
#define IJLAT_COMMIT_IN_ALLOC 1
#endif

#if IJLAT_COMMIT_IN_HOOK
#undef IJLAT_COMMIT_IN_ALLOC
#undef IJLAT_IGNORE_MIKE_PALL
#endif

#if !IJLAT_COMMIT_IN_ALLOC && !IJLAT_COMMIT_IN_HOOK
#error wrong configuration
#endif

/* @FUNC_CUSTOMIZATION BEGIN */
#ifndef IJLAT_hash
/* MurmurHash2, by Austin Appleby https://sites.google.com/site/murmurhash/ */
static unsigned int ijlat__murmur_hash(const void *key, int len, unsigned int seed)
{
	/* 'm' and 'r' are mixing constants generated offline.
	They're not really 'magic', they just happen to work well. */
	const unsigned int m = 0x5bd1e995;
	const int r = 24;

	/* Initialize the hash to a 'random' value */
	unsigned int h = seed ^ len;

	/* Mix 4 bytes at a time into the hash */
	const unsigned char * data = (const unsigned char *)key;

	while (len >= 4) {
		unsigned int k = *(unsigned int *)data;

		k *= m;
		k ^= k >> r;
		k *= m;

		h *= m;
		h ^= k;

		data += 4;
		len -= 4;
	}

	/* Handle the last few bytes of the input array */
	switch (len)
	{
	case 3: h ^= data[2] << 16;
	case 2: h ^= data[1] << 8;
	case 1: h ^= data[0];
		h *= m;
	};

	/* Do a few final mixes of the hash to ensure the last few
	bytes are well-incorporated. */
	h ^= h >> 13;
	h *= m;
	h ^= h >> 15;

	return h;
}

#define IJLAT_hash ijlat__murmur_hash

#endif

#ifndef IJLAT_assert
#include <assert.h>
#define IJLAT_assert assert
#endif

#ifndef IJLAT_memcpy
#include <string.h>
#define IJLAT_memcpy memcpy
#endif

#ifndef IJLAT_memset
#include <string.h>
#define IJLAT_memset memset
#endif

#ifndef IJLAT_strlen
#include <string.h>
#define IJLAT_strlen strlen
#endif

#ifndef IJLAT_int2str

static unsigned ijlat__ndigits(int value)
{
	unsigned r = (value < 0) ? 2 : 1;
	unsigned v = (value < 0) ? -value : value;
	for (;;) {
		if (v < 10) return r;
		if (v < 100) return r + 1;
		if (v < 1000) return r + 2;
		if (v < 10000) return r + 3;

		v /= 10000u;
		r += 4;
	}
}

static int ijlat__int2str(char *s, int value)
{
	unsigned r = ijlat__ndigits(value);
	unsigned v = (value < 0) ? -value : value;

	unsigned pos = r - 1;
	s[r] = 0;
	while (v >= 10) {
		char c = (v % 10);
		s[pos--] = '0' + c;
		v /= 10;
	}
	s[pos--] = (char)v + '0';
	if (value < 0) {
		IJLAT_assert(pos == 0);
		s[0] = '-';
	}
	return r;
}

#define IJLAT_int2str ijlat__int2str

#endif

/* @FUNC_CUSTOMIZATION END */

#define IJLAT_INVALID_INDEX (unsigned)-1

static ijlat *ijlat__gtraces;

static lua_Hook ijlat__ohook;
static int ijlat__omask = 0;

#if IJLAT_IGNORE_MIKE_PALL
static lua_State *ijlat__globall;
#endif
static lua_Alloc ijlat__oalloc;
static void *ijlat__oallocud;
static size_t ijlat__last_alloc_size;

static void *ijlat__alloc_internal(void *ptr, size_t osize, size_t nsize)
{
	return ijlat__oalloc(ijlat__oallocud, ptr, osize, nsize);
}

static int ijlat__stack_info(lua_State *L, lua_Debug *s, unsigned n, const char *what)
{
	int level = 0;

	while (n-- && lua_getstack(L, level, &s[level])) {
		int r = lua_getinfo(L, what, &s[level]);
		IJLAT_assert(r);
		++level;
	}

	return level;
}

static unsigned ijlat__stack_hash(lua_Debug *s, unsigned n)
{
	unsigned i, final_hash = 0;
	for (i = 0; i != n; ++i) {
		lua_Debug *current = s + i;
		/* this assumes that the source points to interned strings */
		final_hash = IJLAT_hash(current->source, sizeof current->source, final_hash);
		final_hash = IJLAT_hash(&current->currentline, sizeof current->currentline, final_hash);
	}
	return final_hash;
}

#define ijroundup32(x) (--(x), (x)|=(x)>>1, (x)|=(x)>>2, (x)|=(x)>>4, (x)|=(x)>>8, (x)|=(x)>>16, ++(x))

static void ijlat__trace_ensure_can_fit(ijlat *self, unsigned n)
{
	unsigned newallocn;
	if (self->talloc >= self->tused + n)
		return;

	newallocn = self->tused + n;
	ijroundup32(newallocn);

	/* how to handle memory failure ? */
	self->t = (ijlat_alloc_trace*)ijlat__alloc_internal(self->t, sizeof *self->t*self->talloc, sizeof *self->t*newallocn);
	self->talloc = newallocn;
}

static void ijlat__string_ensure_can_fit(ijlat *self, unsigned n)
{
	unsigned newallocn;
	if (self->salloc >= self->sused + n)
		return;

	newallocn = self->sused + n;
	ijroundup32(newallocn);

	/* how to handle memory failure ? (no, I did not just copy/paste above function :) ) */
	self->s = (char*)ijlat__alloc_internal(self->s, sizeof *self->s*self->salloc, sizeof *self->s*newallocn);
	self->salloc = newallocn;
}

static ijlat_alloc_trace *ijlat__trace_extend(ijlat *self)
{
	ijlat__trace_ensure_can_fit(self, 1);
	return &self->t[self->tused++];
}

static ijlat_alloc_trace *ijlat__trace_find(ijlat_alloc_trace *traces, unsigned ntraces, unsigned hash)
{
	unsigned i;
	for (i = 0; i != ntraces; ++i) {
		ijlat_alloc_trace *c = traces + i;
		if (c->hash == hash)
			return c;
	}
	return 0;
}

/* debug _remove_ */
static ijlat_alloc_trace *ijlat__trace_find_empty(ijlat_alloc_trace *traces, unsigned ntraces)
{
	unsigned i;
	for (i = 0; i != ntraces; ++i) {
		ijlat_alloc_trace *c = traces + i;
		if (!c->nallocations)
			return c;
	}
	return 0;
}

#define IJLAT_NL "\n"
#define IJLAT_SEP " : "

#define IJLAT_SEP_NS " ["
#define IJLAT_SEP_NE "]"

#define IJLAT_LL(s) (sizeof((s)) / sizeof(char)-1)

/* returns length of string representation (assumes fit if buffer is supplied)
if allocation size/length is needed:
- call with NULL buffer
- allocate len+1 bytes
- call again with allocated buffer
*/

static unsigned ijlat__stack_to_string(lua_Debug *ars, unsigned n, char *b)
{
	char temp[33];

	unsigned i, cursor = 0;

	for (i = 0; i != n; ++i) {
		lua_Debug *current = ars + i;
		int sres;
		unsigned len;
		char *linestring;
		IJLAT_assert(current->source);
		len = (unsigned)IJLAT_strlen(current->source);
		if (b) IJLAT_memcpy(b + cursor, current->source, len);

		cursor += len;

		if (b) IJLAT_memcpy(b + cursor, IJLAT_SEP, IJLAT_LL(IJLAT_SEP));

		cursor += IJLAT_LL(IJLAT_SEP);

		linestring = b ? (b + cursor) : temp;

		sres = IJLAT_int2str(linestring, current->currentline);
		IJLAT_assert(sres > 0);
		cursor += (unsigned)sres;

		if (current->name) {
			if (b) IJLAT_memcpy(b + cursor, IJLAT_SEP_NS, IJLAT_LL(IJLAT_SEP_NS));

			cursor += IJLAT_LL(IJLAT_SEP_NS);

			len = (unsigned)IJLAT_strlen(current->name);
			if (b) IJLAT_memcpy(b + cursor, current->name, len);
			cursor += len;

			if (b) IJLAT_memcpy(b + cursor, IJLAT_SEP_NE, IJLAT_LL(IJLAT_SEP_NE));

			cursor += IJLAT_LL(IJLAT_SEP_NE);
		}

		if (b) IJLAT_memcpy(b + cursor, IJLAT_NL, IJLAT_LL(IJLAT_NL));

		cursor += IJLAT_LL(IJLAT_NL);
	}

	return cursor;
}

static void ijlat__record_allocation(ijlat *self, lua_State *L, size_t allocation_size)
{
	if (!L || !allocation_size)
		return;
	else {
		lua_Debug debug_infos[8];
		unsigned num_filled = ijlat__stack_info(L, debug_infos, sizeof debug_infos / sizeof *debug_infos, "Sl");
		unsigned chash = ijlat__stack_hash(debug_infos, num_filled);

		ijlat_alloc_trace *current_trace = ijlat__trace_find(self->t, self->tused, chash);
		if (!current_trace) {
			unsigned len_needed, num_filled_trace = ijlat__stack_info(L, debug_infos, sizeof debug_infos / sizeof *debug_infos, "n");

			IJLAT_assert(num_filled_trace == num_filled);
			current_trace = ijlat__trace_extend(self);
			current_trace->hash = chash;
			current_trace->nallocations = 0;
			current_trace->tallocated = 0;
			current_trace->soffset = self->sused;

			len_needed = ijlat__stack_to_string(debug_infos, num_filled_trace, 0);
			ijlat__string_ensure_can_fit(self, len_needed + 1);
			ijlat__stack_to_string(debug_infos, num_filled_trace, self->s + self->sused);

			self->sused += len_needed + 1;
			self->s[self->sused - 1] = 0; /* null-terminate */
		}

		current_trace->nallocations += 1;
		current_trace->tallocated += (unsigned)allocation_size;
	}
}

static void *ijlat__alloc_wrap(void *ud, void *ptr, size_t osize, size_t nsize)
{
	ijlat__last_alloc_size = (nsize > osize) ? (nsize - osize) : 0;

#if IJLAT_COMMIT_IN_ALLOC
#if IJLAT_IGNORE_MIKE_PALL
	ijlat__record_allocation(ijlat__gtraces, ijlat__globall, ijlat__last_alloc_size);
#else
	if (ijlat__last_alloc_size && ijlat__gtraces->pending_index != IJLAT_INVALID_INDEX) {
		ijlat_alloc_trace *current_trace = ijlat__gtraces->t + ijlat__gtraces->pending_index;
		current_trace->nallocations += 1;
		current_trace->tallocated += ijlat__last_alloc_size;
	}
#endif
#endif
	return ijlat__oalloc(ud, ptr, osize, nsize);
}

static void ijlat__hook(lua_State *L, lua_Debug *ar)
{
	if (ar->event == LUA_HOOKLINE) {
#if IJLAT_COMMIT_IN_ALLOC && !IJLAT_IGNORE_MIKE_PALL
		/* record the latest callstack that will be committed in the allocation
		callback */
		lua_Debug debug_infos[8];
		unsigned num_filled = ijlat__stack_info(L, debug_infos, sizeof debug_infos / sizeof *debug_infos, "Sl");
		unsigned chash = ijlat__stack_hash(debug_infos, num_filled);

		ijlat_alloc_trace *current_trace = ijlat__trace_find(ijlat__gtraces->t, ijlat__gtraces->tused, chash);

		if (!current_trace) {
			unsigned len_needed, num_filled_trace = ijlat__stack_info(L, debug_infos, sizeof debug_infos / sizeof *debug_infos, "n");
			IJLAT_assert(num_filled_trace == num_filled);
			if (ijlat_is_backindex_and_empty(ijlat__gtraces, ijlat__gtraces->pending_index)) {
				/* replace back */
				current_trace = &ijlat__gtraces->t[ijlat__gtraces->pending_index];
				ijlat__gtraces->sused = current_trace->soffset; /* rewind */
			}
			else {
				IJLAT_assert(!ijlat__trace_find_empty(ijlat__gtraces->t, ijlat__gtraces->tused));
				ijlat__gtraces->pending_index = ijlat__gtraces->tused;
				current_trace = ijlat__trace_extend(ijlat__gtraces);
			}
			current_trace->hash = chash;
			current_trace->nallocations = 0;
			current_trace->tallocated = 0;
			current_trace->soffset = ijlat__gtraces->sused;

			len_needed = ijlat__stack_to_string(debug_infos, num_filled_trace, 0);
			ijlat__string_ensure_can_fit(ijlat__gtraces, len_needed + 1);
			ijlat__stack_to_string(debug_infos, num_filled_trace, ijlat__gtraces->s + ijlat__gtraces->sused);

			ijlat__gtraces->sused += len_needed + 1;
			ijlat__gtraces->s[ijlat__gtraces->sused - 1] = 0; /* null-terminate */

		}
		else {
			unsigned newpending = current_trace - ijlat__gtraces->t;
			unsigned oldpending = ijlat__gtraces->pending_index;

			ijlat__gtraces->pending_index = newpending;

			if (newpending != oldpending &&
				ijlat_is_backindex_and_empty(ijlat__gtraces, oldpending)) {
				ijlat__gtraces->sused = ijlat__gtraces->t[oldpending].soffset; /* rewind */
				--ijlat__gtraces->tused; /* pop back */
			}
		}
#elif IJLAT_COMMIT_IN_HOOK
		ijlat__record_allocation(ijlat__gtraces, L, ijlat__last_alloc_size);
		ijlat__last_alloc_size = 0;
#endif

	}

	if (ijlat__ohook)
		ijlat__ohook(L, ar);
}

IJLAT_DEF void ijlat_init(ijlat *self, lua_State *L)
{
	IJLAT_assert(!ijlat__oalloc);
#if IJLAT_IGNORE_MIKE_PALL
	ijlat__globall = L;
#endif
	{
		ijlat__oalloc = lua_getallocf(L, &ijlat__oallocud);
		lua_setallocf(L, ijlat__alloc_wrap, ijlat__oallocud);
		IJLAT_assert(ijlat__oalloc);
	}

	{
		ijlat__omask = lua_gethookmask(L);
		ijlat__ohook = lua_gethook(L);

		lua_sethook(L, ijlat__hook, ijlat__omask | LUA_MASKLINE, lua_gethookcount(L));
	}

	IJLAT_memset(self, 0, sizeof *self);
	self->L = L;
	self->pending_index = IJLAT_INVALID_INDEX;

	ijlat__trace_ensure_can_fit(self, 64);
	ijlat__string_ensure_can_fit(self, 1024);
	ijlat__gtraces = self;
}

IJLAT_DEF void ijlat_shutdown(ijlat *self)
{
	void *ud;
	lua_State *L = self->L;

	IJLAT_assert(ijlat__oalloc);
	lua_getallocf(L, &ud);
	lua_setallocf(L, ijlat__oalloc, ud);

	lua_sethook(L, ijlat__ohook, ijlat__omask, lua_gethookcount(L));

	ijlat__oalloc(ud, self->t, sizeof *self->t*self->talloc, 0);
	ijlat__oalloc(ud, self->s, sizeof *self->s*self->salloc, 0);
	IJLAT_memset(self, 0, sizeof *self);
	self->pending_index = IJLAT_INVALID_INDEX;

	ijlat__oalloc = 0;
	ijlat__ohook = 0;

	ijlat__gtraces = 0;

#if IJLAT_IGNORE_MIKE_PALL
	ijlat__globall = 0;
#endif
}
#endif