/*
* immutable.c
* An immutable table library.
*/

#include "lua.h"
#include "lauxlib.h"

#define MYNAME "immutable"
#define MYTYPE MYNAME

// The C API changed from 5.1 to 5.2, so these shims help the code compile on >=5.2
#if LUA_VERSION_NUM >= 502
#define lua_objlen(L, i) lua_rawlen(L, i)
#define lua_equal(L, i, j) lua_compare(L, i, j, LUA_OPEQ)
#define luaH_getnum(t, k) luaH_getint(t, k)
#define luaL_register(L, _, R) luaL_setfuncs(L, R, 0)
#endif

static int Lcreate_instance(lua_State *L)
{
    int n_args = lua_gettop(L);
    // arg 1: class table, ...

    lua_getfield(L, 1, "__fields");
    // Stack: [__fields]
    size_t n = lua_objlen(L,-1);
    if ((size_t)n_args-1 != n) {
        lua_pushstring(L, "incorrect number of arguments");
        lua_error(L);
    }
    lua_pop(L,1);
    // Stack: []

    lua_createtable(L, n, 0);
    // Stack [inst]
    
    // Copy in all the values, and simultaneously compute the hash:
    lua_Integer hash = 0;
    for (lua_Integer i=1; i <=(lua_Integer)n; i++) {
        lua_pushvalue(L, i+1);
        // Stack [inst, args[i+1]]
        lua_Integer item_hash;
        switch (lua_type(L, -1)) {
            case LUA_TNIL:
                item_hash = 0;
                break;
            case LUA_TNUMBER:
                item_hash = (lua_Integer)lua_tonumber(L, -1);
                break;
            case LUA_TBOOLEAN:
                item_hash = (lua_Integer)lua_toboolean(L, -1);
                break;
            case LUA_TTABLE:
            case LUA_TFUNCTION:
            case LUA_TUSERDATA:
            case LUA_TTHREAD:
            case LUA_TLIGHTUSERDATA:
                item_hash = (lua_Integer)lua_topointer(L, -1);
                break;
            case LUA_TSTRING:
                {
                    size_t strlen;
                    const char *str = lua_tolstring(L, -1, &strlen);
                    item_hash = *str << 7;
                    for (const char *end = &str[strlen]; str < end; str++)
                        item_hash = (1000003*item_hash) ^ *str;
                    item_hash ^= strlen;
                    break;
                }
            default:
                item_hash = 0;
        }
        hash = (1000003 * hash) ^ item_hash;
        lua_rawseti(L, -2, i);
        // Stack [inst]
    }

    lua_getfield(L, 1, "__instances");

    // Stack: [inst, buckets]
    // Find bucket
    lua_rawgeti(L, -1, hash);
    // Stack: [inst, buckets, bucket]
    if (lua_isnil(L, -1)) {
        // Make a new bucket
        // Stack: [inst, buckets, nil]
        lua_pop(L, 1);
        // Stack: [inst, buckets]
        lua_createtable(L, 1, 0);
        // Stack: [inst, buckets, bucket]
        // TODO: share bucket metatables
        lua_createtable(L, 0, 1);
        // Stack: [inst, buckets, bucket, bucket_mt]
        lua_pushstring(L, "k");
        lua_setfield(L, -2, "__mode");
        // Stack: [inst, buckets, bucket, {'__mode'='k'}]
        lua_setmetatable(L, -2);
        // Stack: [inst, buckets, bucket]

        lua_pushvalue(L, -1);
        // Stack: [inst, buckets, bucket, bucket]
        lua_rawseti(L, -3, hash);
        // Stack: [inst, buckets, bucket]
    }
    // Stack: [inst, buckets, bucket]
    // scan bucket
    lua_pushnil(L);
    while (lua_next(L, -2) != 0) { // for hash_collider_inst, hash_collider in pairs(bucket) do
        // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider]
        int bucket_item_matches = 1;

        // Perform a full equality check
        lua_pushnil(L);
        while (lua_next(L, -2) != 0) { // for collider_key, collider_value in pairs(hash_collider) do
            // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider, collider_key, collider_value]
            lua_pushvalue(L, -2);
            // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider, collider_key, collider_value, collider_key]
            lua_gettable(L, -8); // inst[collider_key]
            // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider, collider_key, collider_value, inst_value]
            if (!lua_rawequal(L, -1, -2)) {
                // go to next item in the bucket
                bucket_item_matches = 0;
                // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider, collider_key, collider_value, inst_value]
                lua_pop(L, 4);
                // Stack: [inst, buckets, bucket, hash_collider_inst]
                break;
            } else {
                // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider, collider_key, collider_value, inst_value]
                lua_pop(L, 2);
                // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider, collider_key]
            }
        }
        if (bucket_item_matches) {
            // Stack: [inst, buckets, bucket, hash_collider_inst, hash_collider]
            lua_pop(L, 1);
            // Found matching singleton
            return 1;
        }
    }

    // failed to find a singleton
    // Stack: [inst, buckets, bucket]

    lua_Integer* userdata = (lua_Integer*)lua_newuserdata(L, sizeof(lua_Integer));
    *userdata = hash;
    
    // Stack [inst, buckets, bucket, inst_userdata]
    lua_pushvalue(L, 1);
    // Stack [inst, buckets, bucket, inst_userdata, metatable]
    lua_setmetatable(L, -2);
    // Stack [inst, buckets, bucket, inst_userdata]
    
    lua_pushvalue(L, -1);
    // Stack [inst, buckets, bucket, inst_userdata, inst_userdata]
    lua_pushvalue(L, -5);
    // Stack [inst, buckets, bucket, inst_userdata, inst_userdata, inst]
    lua_settable(L, -4); // buckets[inst_userdata] = inst
    // Stack [inst, buckets, bucket, inst_userdata]
    return 1;
}


static int Llen(lua_State *L)
{
    lua_pushvalue(L, 1);
    luaL_getmetatable(L, MYTYPE);
    lua_getfield(L, -1, "fields");
    return 1;
}

static int Lindex(lua_State *L)
{
    lua_getmetatable(L, 1);
    // Stack: [mt]
    lua_getfield(L, -1, "__fields");
    // Stack: [mt, fields]
    
    lua_pushnil(L);
    while (lua_next(L, -2) != 0) {
        // Stack: [mt, fields, i, field]
        if (lua_rawequal(L, -1, 2)) {
            lua_pop(L, 1);
            // All the slowness is concentrated here
            // Stack: [mt, fields, i]
            lua_getfield(L, -3, "__instances");
            // Stack: [mt, fields, i, buckets]
            lua_rawgeti(L, -1, *((lua_Integer*)lua_touserdata(L, 1)));
            // Stack: [mt, fields, i, buckets, bucket]
            lua_pushvalue(L, 1);
            // Stack: [mt, fields, i, buckets, bucket, inst_udata]
            lua_rawget(L, -2);
            // Stack: [mt, fields, i, buckets, bucket, inst_table]
            lua_rawgeti(L, -1, lua_tonumber(L, -4));
            // Stack: [mt, fields, i, buckets, bucket, inst_table, result]
            return 1;
        }
        lua_pop(L, 1);
    }
    lua_pop(L, 1);
    // Stack: [mt]
    lua_getfield(L, -1, "__methods");
    // Stack: [mt, __methods]
    lua_pushvalue(L, 2);
    // Stack: [mt, __methods, key]
    lua_gettable(L, -2);
    return 1;
}

static const luaL_Reg R[] =
{
    { "__len", Llen},
    { "__index", Lindex},
    { "new", Lcreate_instance},
    { NULL, NULL}
};

static int Lmake_class(lua_State *L)
{
    // args: fields, [methods, [metamethods]]
    int n_args = lua_gettop(L);

    // CLS = {}
    lua_newtable(L);
    // Stack: [CLS]

    // If metamethods were passed in, copy them over
    if (n_args > 2) {
        lua_pushnil(L);
        // Stack: [CLS, nil]
        while (lua_next(L, 3) != 0) {
            lua_pushvalue(L, -2);
            lua_pushvalue(L, -2);
            // Stack: [CLS, method_name, method_value, method_name, method_value]
            lua_settable(L, -5);
            // Stack: [CLS, method_name, method_value]
            lua_pop(L, 1);
            // Stack: [CLS, method_name]
        }
    }

    // CLS.buckets = {}
    lua_newtable(L);
    // Stack: [CLS, CLS.buckets]
    lua_setfield(L, -2, "__instances");

    // Stack: [CLS]

    // CLS.__fields = arg1
    lua_pushvalue(L, 1);
    // Stack: [CLS, __fields]
    lua_setfield(L, -2, "__fields");
    // Stack: [CLS]

    lua_newtable(L);
    // Stack: [CLS, __methods]
    // If methods were passed in, copy them over
    if (n_args > 1) {
        // Stack: [
        lua_pushnil(L);
        // Stack: [CLS, __methods, nil]
        while (lua_next(L, 2) != 0) {
            lua_pushvalue(L, -2);
            lua_pushvalue(L, -2);
            // Stack: [CLS, __methods, method_name, method_value, method_name, method_value]
            lua_settable(L, -5);
            // Stack: [CLS, __methods, method_name, method_value]
            lua_pop(L, 1);
            // Stack: [CLS, __methods, method_name]
        }
        // Stack: [CLS, __methods]
    }
    lua_setfield(L, -2, "__methods");
    // Stack: [CLS]

    // Populate __len, __eq, __index, new, etc.
    luaL_register(L,NULL,R);

    // setmetatable(CLS, {__new=CLS.new})
    lua_createtable(L, 0, 1);
    lua_pushcfunction(L, Lcreate_instance);
    lua_setfield(L, -2, "__call");
    lua_setmetatable(L, -2);
    // Stack [CLS]

    return 1;
}

LUALIB_API int luaopen_immutable(lua_State *L)
{
    lua_pushcfunction(L, Lmake_class);
    return 1;
}