A stupid sampling profiler thing for Lua

heatmap.c

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

#include <string.h>

static int table_ref = LUA_NOREF;

static void
hook(lua_State *L, lua_Debug *ar)
{
    int count;

    // load the source filename and current line number into the activation record
    lua_getinfo(L, "Sl", ar);

    // ignore C functions
    if(!strcmp(ar->what, "C")) {
        return;
    }

    lua_rawgeti(L, LUA_REGISTRYINDEX, table_ref);
    
    // build a key for the table: filename .. UNIT_SEPARATOR .. line_number
    lua_pushstring(L, ar->source + 1); // Lua filename sources have a leading "@" character - strip it
    lua_pushliteral(L, "\x1f");
    lua_pushinteger(L, ar->currentline);
    lua_concat(L, 3);
    
    lua_pushvalue(L, -1); // duplicate the key so we can use it for lua_settable later
    
    // look up the current value for the key - defaulting to 0
    lua_gettable(L, -3);
    if(lua_isnil(L, -1)) {
        count = 0;
    } else {
        count = lua_tointeger(L, -1);
    }
    lua_pop(L, 1);
    
    count++;
    lua_pushinteger(L, count);
    lua_settable(L, -3);
    lua_pop(L, 1); // cleanup
}

static int
stophook(lua_State *L)
{
    lua_sethook(L, NULL, 0, 0);
    lua_rawgeti(L, LUA_REGISTRYINDEX, table_ref);
    luaL_unref(L, LUA_REGISTRYINDEX, table_ref);
    table_ref = LUA_NOREF;

    return 1;
}

static int
heatmap(lua_State *L)
{
    lua_newtable(L);
    table_ref = luaL_ref(L, LUA_REGISTRYINDEX);
    lua_sethook(L, hook, LUA_MASKCOUNT, luaL_checkinteger(L, 1));
    lua_pushcfunction(L, stophook);

    return 1;
}

int
luaopen_heatmap(lua_State *L)
{
    lua_pushcfunction(L, heatmap);

    return 1;
}

Makefile

CFLAGS += -fPIC

all: heatmap.so

%.so: %.o
	gcc -shared -o $@ $^

clean:
	rm -f *.o *.so