/*
* immutable.c
* An immutable table library by Bruce Hill. This library returns a single function
* that can be used to declare immutable classes, like so:
*
*     immutable = require 'immutable'
*     local Foo = immutable({"baz","qux"})
*     local foo = Foo("hello", 99)
*     assert(not pcall(function() foo.x = 'mutable' end))
*     local t = {[foo]="it works"}
*     assert(t[Foo("hello", 99)] == "it works")
* 
* Instances *are* garbage collected.
*/

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

// 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

#if !defined(LUAI_HASHLIMIT)
#define LUAI_HASHLIMIT		5
#endif

// This is used to create a unique light userdata to store in the registry to allow all
// hash buckets to share the same {__mode='k'} metatable
static int SHARED_BUCKET_METATABLE;
// This is used to create a unique light userdata to store in the registry to allow all
// constructed classes to have the same metatable:
//  {__new=Lcreate_instance, __tostring=function(cls) return cls.name or 'immutable(...)' end}
static int SHARED_CLASS_METATABLE;

static int Lcreate_instance(lua_State *L)
{
    // arg 1: class table, ...
    lua_getfield(L, 1, "__fields");
    size_t n = lua_objlen(L, -1);
    lua_pop(L, 1);
    
    // Compute the hash:
    lua_Integer hash = (lua_Integer)lua_topointer(L, 1); // Hash depends on the metatable used in creation
    for (lua_Integer i=1; i <=(lua_Integer)n; i++) {
        lua_Integer item_hash;
        switch (lua_type(L, 1+i)) {
            case LUA_TNIL: case LUA_TNONE:
                // Arbitrarily chosen value
                item_hash = 0x97167da9;
                break;
            case LUA_TNUMBER:
                {
                    // Cast float bits to integer
                    lua_Number n = lua_tonumber(L, 1+i);
                    item_hash = *((lua_Integer*)&n);
                    break;
                }
            case LUA_TBOOLEAN:
                // Arbitrarily chosen values
                item_hash = lua_toboolean(L, 1+i)? 0x82684f71 : 0x88d66f2a;
                break;
            case LUA_TTABLE:
            case LUA_TFUNCTION:
            case LUA_TUSERDATA:
            case LUA_TTHREAD:
            case LUA_TLIGHTUSERDATA:
                item_hash = (lua_Integer)lua_topointer(L, 1+i);
                break;
            case LUA_TSTRING:
                {
                    // Algorithm taken from Lua 5.3's implementation
                    size_t len;
                    const char *str = lua_tolstring(L, 1+i, &len);
                    item_hash = len ^ 0xd2e9e9ac; // Arbitrary seed
                    size_t step = (len >> LUAI_HASHLIMIT) + 1;
                    for (; len >= step; len -= step)
                        item_hash ^= ((item_hash<<5) + (item_hash>>2) + (unsigned char)(str[len - 1]));
                    break;
                }
            default:
                item_hash = 0;
        }
        hash = (1000003 * hash) ^ item_hash;
    }

    lua_getfield(L, 1, "__instances");

    // Stack: [buckets]
    // Find bucket
    lua_rawgeti(L, -1, hash);
    // Stack: [buckets, bucket]
    if (lua_isnil(L, -1)) {
        // Make a new bucket
        // Stack: [buckets, nil]
        lua_pop(L, 1);
        // Stack: [buckets]
        lua_createtable(L, 1, 0);
        // Stack: [buckets, bucket]
        lua_pushlightuserdata(L, (void*)&SHARED_BUCKET_METATABLE);
        lua_gettable(L, LUA_REGISTRYINDEX);
        // Stack: [buckets, bucket, {'__mode'='k'}]
        lua_setmetatable(L, -2);
        // Stack: [buckets, bucket]
        lua_pushvalue(L, -1);
        // Stack: [buckets, bucket, bucket]
        lua_rawseti(L, -3, hash);
        // Stack: [buckets, bucket]
    }
    // Stack: [buckets, bucket]
    // scan bucket
    lua_pushnil(L);
    while (lua_next(L, -2) != 0) { // for hash_collider_inst, hash_collider in pairs(bucket) do
        // Stack: [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 i, collider_value in pairs(hash_collider) do
            // Stack: [buckets, bucket, hash_collider_inst, hash_collider, i, value]
            if (!lua_rawequal(L, -1, 1+lua_tonumber(L, -2))) { // If the i'th entry doesn't match the i'th arg
                bucket_item_matches = 0;
                // Stack: [buckets, bucket, hash_collider_inst, hash_collider, i, value]
                lua_pop(L, 3);
                // Stack: [buckets, bucket, hash_collider_inst]
                break; // go to next item in the bucket
            } else {
                // Stack: [buckets, bucket, hash_collider_inst, hash_collider, i, value]
                lua_pop(L, 1);
                // Stack: [buckets, bucket, hash_collider_inst, hash_collider, i]
            }
        }
        if (bucket_item_matches) {
            // Stack: [buckets, bucket, hash_collider_inst, hash_collider]
            lua_pop(L, 1);
            // Found matching singleton
            return 1;
        }
    }

    // Failed to find an existing instance, so create a new one
    // Stack: [buckets, bucket]
    lua_Integer* userdata = (lua_Integer*)lua_newuserdata(L, sizeof(lua_Integer));
    *userdata = hash;
    
    // Stack [buckets, bucket, inst_userdata]
    lua_pushvalue(L, 1);
    // Stack [buckets, bucket, inst_userdata, metatable]
    lua_setmetatable(L, -2);
    // Stack [buckets, bucket, inst_userdata]
    lua_pushvalue(L, -1);
    // Stack [buckets, bucket, inst_userdata, inst_userdata]
    lua_createtable(L, n, 0); // Create the table to store the instance's data
    // Stack [buckets, bucket, inst_userdata, inst_userdata, inst_table]
    for (lua_Integer i=1; i <=(lua_Integer)n; i++) {
        lua_pushvalue(L, i+1);
        // Stack [buckets, bucket, inst_userdata, inst_userdata, inst_table, arg #1+i]
        lua_rawseti(L, -2, i);
    }
    // Stack [buckets, bucket, inst_userdata, inst_userdata, inst_table]
    lua_settable(L, -4); // buckets[inst_userdata] = inst_table
    // Stack [buckets, bucket, inst_userdata]
    return 1;
}

static int Lfrom_table(lua_State *L)
{
    lua_pushvalue(L, 1);
    // Stack: [mt]
    lua_getfield(L, -1, "__fields");
    // Stack: [mt, fields]
    lua_pushnil(L);
    int num_args = 0;
    while (lua_next(L, -2) != 0) {
        // Stack: [mt, fields, i, field_i]
        lua_gettable(L, 2);
        // Stack: [mt, fields, i, table[field_i]]
        lua_insert(L, -3);
        // Stack: [mt, table[field_i], fields, i]
        num_args++;
    }
    // Stack: [mt, table[field], ..., fields]
    lua_pop(L, 1);
    // Stack: [mt, table[field], ...]
    lua_pushcfunction(L, Lcreate_instance);
    // Stack: [mt, table[field], ..., create]
    lua_insert(L, -(num_args+2));
    // Stack: [create, mt, table[field_1], ...]
    lua_call(L, num_args+1, 1);
    return 1;
}

static int Lis_instance(lua_State *L)
{
    if (lua_type(L, 2) != LUA_TUSERDATA) {
        lua_pushboolean(L, 0);
        return 1;
    }
    lua_getmetatable(L, 2);
    lua_pushboolean(L, lua_rawequal(L, -1, 1));
    return 1;
}

static int Llen(lua_State *L)
{
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    // Stack: [mt]
    lua_getfield(L, -1, "__fields");
    // Stack: [mt, fields]
    lua_pushinteger(L, lua_objlen(L, -1));
    return 1;
}

static int Lindex(lua_State *L)
{
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    // Stack: [mt]
    lua_getfield(L, -1, "__indices");
    // Stack: [mt, indices]
    lua_pushvalue(L, 2);
    // Stack: [mt, indices, k]
    lua_gettable(L, -2);
    // Stack: [mt, indices, i]
    if (! lua_isnil(L, -1)) { // Found the field name
        // Stack: [mt, indices, i]
        lua_getfield(L, -3, "__instances");
        // Stack: [mt, indices, i, buckets]
        lua_Integer* hash_address = (lua_Integer*)lua_touserdata(L, 1);
        if (! hash_address) {
            luaL_error(L, "invalid type");
        }
        lua_rawgeti(L, -1, *hash_address);
        // Stack: [mt, indices, i, buckets, bucket]
        lua_pushvalue(L, 1);
        // Stack: [mt, indices, i, buckets, bucket, inst_udata]
        lua_rawget(L, -2);
        // Stack: [mt, indices, i, buckets, bucket, inst_table]
        lua_rawgeti(L, -1, lua_tointeger(L, -4));
        return 1;
    } else if (lua_type(L, 2) == LUA_TNUMBER) {
        lua_pop(L, 2);
        // Stack: [mt]
        lua_getfield(L, -1, "__instances");
        // Stack: [mt, buckets]
        lua_Integer* hash_address = (lua_Integer*)lua_touserdata(L, 1);
        if (! hash_address) {
            luaL_error(L, "invalid type");
        }
        lua_rawgeti(L, -1, *hash_address);
        // Stack: [mt, buckets, bucket]
        lua_pushvalue(L, 1);
        // Stack: [mt, buckets, bucket, inst_udata]
        lua_rawget(L, -2);
        // Stack: [mt, buckets, bucket, inst_table]
        lua_rawgeti(L, -1, lua_tointeger(L, 2));
        if (! lua_isnil(L, -1)) {
            // Found numeric index
            return 1;
        } else {
            // Fall back to class
            // Stack: [mt, buckets, bucket, inst_table, v, k]
            lua_pushvalue(L, 2);
            // Stack: [mt, key]
            lua_gettable(L, -6);
            return 1;
        }
    } else {
        // Fall back to class:
        // Stack: [mt, indices, i]
        lua_pop(L, 2);
        // Stack: [mt]
        lua_pushvalue(L, 2);
        // Stack: [mt, key]
        lua_gettable(L, -2);
        return 1;
    }
}

static int Ltostring(lua_State *L)
{
    luaL_Buffer b;
    luaL_buffinit(L, &b);

    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    // Stack: [mt]
    
    lua_getfield(L, -1, "name");
    if (!lua_isnil(L, -1)) {
        luaL_addvalue(&b);
    } else {
        lua_pop(L, 1);
    }
    luaL_addstring(&b, "(");
    
    lua_getfield(L, -1, "__instances");
    // Stack: [mt, buckets]
    lua_Integer* hash_address = (lua_Integer*)lua_touserdata(L, 1);
    if (! hash_address) {
        luaL_error(L, "invalid type");
    }
    lua_rawgeti(L, -1, *hash_address);
    // Stack: [mt, buckets, bucket]
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    // Stack: [mt, buckets, bucket, inst_table]
    lua_getglobal(L, "tostring");
    // Stack: [mt, buckets, bucket, inst_table, tostring]
    lua_getfield(L, -5, "__fields");
    // Stack: [mt, buckets, bucket, inst_table, tostring, fields]
    
    lua_pushnil(L);
    int needs_comma = 0;
    int numeric_index = 1;
    while (lua_next(L, -2) != 0) {
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, fieldname]
        if (needs_comma) {
            luaL_addstring(&b, ", ");
        } else {
            needs_comma = 1;
        }
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, fieldname]
        if (lua_type(L, -1) == LUA_TNUMBER && lua_tointeger(L, -1) == numeric_index) {
            lua_pop(L, 1);
        } else {
            lua_pushvalue(L, -4);
            // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, fieldname, tostring]
            lua_insert(L, -2);
            // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, tostring, fieldname]
            lua_call(L, 1, 1);
            // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, field string]
            luaL_addvalue(&b);
            // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i]
            luaL_addstring(&b, "=");
        }
        lua_rawgeti(L, -4, lua_tonumber(L, -1));
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, value]
        lua_pushvalue(L, -4);
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, value, tostring]
        lua_insert(L, -2);
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, tostring, value]
        lua_call(L, 1, 1);
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i, value string]
        luaL_addvalue(&b);
        // Stack: [mt, buckets, bucket, inst_table, tostring, fields, i]
        numeric_index++;
    }
    luaL_addstring(&b, ")");
    luaL_pushresult(&b);
    return 1;
}

static int Lnexti(lua_State *L)
{
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    // Stack: [mt]
    lua_getfield(L, -1, "__instances");
    // Stack: [mt, buckets]
    lua_Integer* hash_address = (lua_Integer*)lua_touserdata(L, 1);
    if (! hash_address) {
        luaL_error(L, "invalid type");
    }
    lua_rawgeti(L, -1, *hash_address);
    // Stack: [mt, buckets, bucket]
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    // Stack: [mt, buckets, bucket, inst_table]
    lua_getfield(L, -4, "__fields");
    // Stack: [mt, buckets, bucket, inst_table, fields]
    lua_pushvalue(L, 2);
    // Stack: [mt, buckets, bucket, inst_table, fields, i]
    if (lua_next(L, -2) == 0) {
        return 0;
    }
    // Stack: [mt, buckets, bucket, inst_table, fields, i2, next_fieldname]
    lua_pop(L, 1);
    // Stack: [mt, buckets, bucket, inst_table, fields, i2]
    lua_rawgeti(L, -3, lua_tonumber(L, -1));
    // Stack: [mt, buckets, bucket, inst_table, fields, i2, value]
    return 2;
}

static int Lipairs(lua_State *L)
{
    lua_pushcfunction(L, Lnexti);
    // Stack: [Lnexti]
    lua_pushvalue(L, 1);
    // Stack: [Lnexti, inst_udata]
    lua_pushnil(L);
    // Stack: [Lnexti, inst_udata, nil]
    return 3;
}

static int Lnext(lua_State *L)
{
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    // Stack: [mt]
    lua_getfield(L, -1, "__instances");
    // Stack: [mt, buckets]
    lua_Integer* hash_address = (lua_Integer*)lua_touserdata(L, 1);
    if (! hash_address) {
        luaL_error(L, "invalid type");
    }
    lua_rawgeti(L, -1, *hash_address);
    // Stack: [mt, buckets, bucket]
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    // Stack: [mt, buckets, bucket, inst_table]
    lua_getfield(L, -4, "__indices");
    // Stack: [mt, buckets, bucket, inst_table, fields]
    lua_pushvalue(L, 2);
    // Stack: [mt, buckets, bucket, inst_table, fields, k]
    if (lua_next(L, -2) == 0) {
        return 0;
    }
    // Stack: [mt, buckets, bucket, inst_table, fields, k2, next_i]
    lua_gettable(L, -4);
    // Stack: [mt, buckets, bucket, inst_table, fields, k2, value]
    return 2;
}

static int Lpairs(lua_State *L)
{
    lua_pushcfunction(L, Lnext);
    // Stack: [Lnexti]
    lua_pushvalue(L, 1);
    // Stack: [Lnexti, inst_udata]
    lua_pushnil(L);
    // Stack: [Lnexti, inst_udata, nil]
    return 3;
}

static int Lgc(lua_State *L)
{
    // Unfortunately, a __gc metamethod is needed here, because the data storage format
    // is cls.__instances[hash][udata] = data, and even though cls.__instances[hash]
    // is a weak-keyed table, cls.__instances will leak hash key -> empty table mappings
    // over time.
    lua_getmetatable(L, 1);
    // Stack: [mt]
    lua_getfield(L, -1, "__instances");
    // Stack: [mt, buckets]
    lua_Integer* hash_address = (lua_Integer*)lua_touserdata(L, 1);
    if (! hash_address) {
        return 0;
    }
    lua_rawgeti(L, -1, *hash_address);
    // Stack: [mt, indices, i, buckets, bucket]
    if (lua_isnil(L, -1)) {
        return 0;
    }
    // Set cls.__instances[hash][udata] = nil
    lua_pushvalue(L, 1);
    // Stack: [mt, indices, i, buckets, bucket, inst_udata]
    lua_pushnil(L);
    // Stack: [mt, indices, i, buckets, bucket, inst_udata, nil]
    lua_settable(L, -3);
    // Stack: [mt, indices, i, buckets, bucket]
    lua_pushnil(L);
    // If next(cls.__instances[hash]) == nil, then set cls.__instances[hash] = nil
    // Stack: [mt, indices, i, buckets, bucket, nil]
    if (lua_next(L, -2) == 0) {
        lua_pop(L, 1);
        // Stack: [mt, indices, i, buckets]
        lua_pushnil(L);
        // Stack: [mt, indices, i, buckets, nil]
        lua_rawseti(L, -2, *hash_address);
    }
    return 1;
}

static const luaL_Reg Rinstance_metamethods[] =
{
    { "__len", Llen},
    { "__index", Lindex},
    { "__tostring", Ltostring},
    { "__ipairs", Lipairs},
    { "__pairs", Lpairs},
    { "__gc", Lgc},
    { "from_table", Lfrom_table},
    { "is_instance", Lis_instance},
    { NULL, NULL}
};

static int Lmake_class(lua_State *L)
{
    // immutable([fields], [methods/metamethods])
    
    lua_newtable(L);
    // Stack: [CLS]
    // Populate CLS.__len, CLS.__index, CLS.__pairs, etc.
    luaL_register(L,NULL,Rinstance_metamethods);

    // If methods were passed in, copy them over, overwriting defaults if desired
    if (lua_type(L, 2) == LUA_TTABLE) {
        // Stack: [CLS]
        lua_pushnil(L);
        // Stack: [CLS, nil]
        while (lua_next(L, 2) != 0) {
            // Stack: [CLS, method_name, method_value]
            lua_pushvalue(L, -2);
            // Stack: [CLS, method_name, method_value, method_name]
            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]
        }
        // Stack: [CLS]
    }

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

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

            size_t n = lua_objlen(L, 1);
            lua_createtable(L, 0, n);
            // Stack: [CLS, __indices]
            lua_pushnil(L);
            while (lua_next(L, 1) != 0) {
                // Stack: [CLS, __indices, i, fieldname]
                lua_pushvalue(L, -2);
                // Stack: [CLS, __indices, i, fieldname, i]
                lua_settable(L, -4);
                // Stack: [CLS, __indices, i]
            }
            lua_setfield(L, -2, "__indices");
            break;
        }
        case LUA_TNUMBER: {
            // If no fields were passed in, make them empty (i.e. a singleton)
            lua_Integer n = lua_tointeger(L, 1);
            if (n < 0) {
                luaL_error(L, "immutable table size must be positive");
            }
            lua_createtable(L, n, 0);
            lua_createtable(L, n, 0);
            // Stack: [CLS, __fields, __indices]
            for (lua_Integer i = 1; i <= n; i++) {
                lua_pushinteger(L, i);
                // Stack: [CLS, __fields, __indices, i]
                lua_rawseti(L, -2, i);
                // Stack: [CLS, __fields, __indices]
                lua_pushinteger(L, i);
                // Stack: [CLS, __fields, __indices, i]
                lua_rawseti(L, -3, i);
                // Stack: [CLS, __fields, __indices]
            }
            // Stack: [CLS, __fields, __indices]
            lua_setfield(L, -3, "__indices");
            // Stack: [CLS, __fields]
            lua_setfield(L, -2, "__fields");
            break;
        }
        case LUA_TNIL: case LUA_TNONE: {
            // If no fields were passed in, make them empty (i.e. a singleton)
            lua_createtable(L, 0, 0);
            lua_setfield(L, -2, "__fields");
            lua_createtable(L, 0, 0);
            lua_setfield(L, -2, "__indices");
            break;
        }
        default: {
            luaL_error(L, "expected number, table, or nil");
        }
    }
    // Stack: [CLS]
    lua_pushlightuserdata(L, (void*)&SHARED_CLASS_METATABLE);
    lua_gettable(L, LUA_REGISTRYINDEX);
    lua_setmetatable(L, -2);
    // Stack [CLS]

    return 1;
}

static int Lclass_tostring(lua_State *L)
{
    lua_getfield(L, 1, "name");
    if (lua_isnil(L, -1)) {
        lua_pushfstring(L, "immutable type: %p", lua_topointer(L, 1));
    }
    return 1;
}

static const luaL_Reg Rclass_metamethods[] =
{
    { "__call", Lcreate_instance},
    { "__tostring", Lclass_tostring},
    { NULL, NULL}
};

LUALIB_API int luaopen_immutable(lua_State *L)
{
    lua_pushlightuserdata(L, (void*)&SHARED_BUCKET_METATABLE);
    lua_createtable(L, 0, 1);
    lua_pushstring(L, "k");
    lua_setfield(L, -2, "__mode");
    lua_settable(L, LUA_REGISTRYINDEX);

    lua_pushlightuserdata(L, (void*)&SHARED_CLASS_METATABLE);
    lua_createtable(L, 0, 2);
    luaL_register(L,NULL,Rclass_metamethods);
    lua_settable(L, LUA_REGISTRYINDEX);

    lua_pushcfunction(L, Lmake_class);
    return 1;
}