/*
* 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"}, {class_var=5})
*     local foo = Foo("hello", 99)
*     assert(foo.baz == "hello")
*     assert(not pcall(function() foo.x = 'mutable' end))
*     local t = {[foo]="it works"}
*     assert(t[Foo("hello", 99)] == "it works")
*     assert(foo.class_var == 5)
* 
* Instances *are* garbage collected.
*
* Class layout:
*     __instances: weak-valued map from hash -> hash bucket
*         hash buckets: weak-keyed map from instance userdata -> table of values
*     __buckets: weak-keyed map from instance userdata -> hash bucket (used to manage hash bucket lifetimes)
*     __fields: list of named fields
*     metamethods, methods, class variables, etc.
*/

#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 luaL_register(L, _, R) luaL_setfuncs(L, R, 0)
#endif
// Lua 5.3 introduced lua_isinteger, fall back to lua_isnumber
#if LUA_VERSION_NUM < 503
#define lua_isinteger(L, i) lua_isnumber(L, i)
#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 WEAK_KEY_METATABLE;
// This is used to create a unique light userdata to store in the registry to allow all
// __instances tables to share the same {__mode='v'} metatable
static int WEAK_VALUE_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 SUPER_METHODS, NIL_SENTINEL;

typedef struct {
    int hash;
    size_t array_len;
} immutable_info_t;

#define IMMUTABLE_GETVALUE(i) if (from_table) {\
    if (i <= num_fields) {\
        lua_rawgeti(L, fields_index, i);\
    } else {\
        lua_pushinteger(L, i-num_fields);\
    }\
    lua_gettable(L, 2);\
} else {\
    if (i <= num_args) {\
        lua_pushvalue(L, 1+i);\
    } else {\
        lua_pushlightuserdata(L, &NIL_SENTINEL);\
    }\
}\
if (lua_isnil(L, -1)) {\
    lua_pop(L, 1);\
    lua_pushlightuserdata(L, &NIL_SENTINEL);\
}

static inline int _create_instance(lua_State *L, int from_table)
{
    size_t num_args = lua_gettop(L)-1;
    lua_getfield(L, 1, "__fields");
    int fields_index = lua_gettop(L);
    size_t num_fields = lua_objlen(L, -1);
    size_t num_values = from_table ? (num_fields + lua_objlen(L, 2)) : (num_args >= num_fields ? num_args : num_fields);
    size_t array_len = num_values - num_fields;
    
    // Compute the hash and populate the values table:
    // Stack: [fields]
    unsigned long long ull_hash = 0x9a937c4d; // Seed
    for (size_t i=1; i <= num_values; i++) {
        unsigned long long item_hash;
        IMMUTABLE_GETVALUE(i);
        // Stack: [fields, value[i]]
        int type = lua_type(L, -1);
        switch (type) {
            case LUA_TNIL: case LUA_TNONE:
                // Arbitrarily chosen value
                item_hash = 0x97167da9;
                break;
            case LUA_TNUMBER:
                {
                    // Cast float bits to integer
                    lua_Number num = lua_tonumber(L, -1);
                    item_hash = *((lua_Integer*)&num);
                    if (item_hash == 0) {
                        item_hash = 0x2887c992;
                    }
                    break;
                }
            case LUA_TBOOLEAN:
                // Arbitrarily chosen values
                item_hash = lua_toboolean(L, -1)? 0x82684f71 : 0x88d66f2a;
                break;
            case LUA_TLIGHTUSERDATA:
                if (lua_touserdata(L, -1) == &NIL_SENTINEL) {
                    item_hash = (1000003 * type) ^ (lua_Integer)lua_topointer(L, -1);
                    break;
                }
                // fallthrough
            case LUA_TTABLE:
            case LUA_TFUNCTION:
            case LUA_TUSERDATA:
            case LUA_TTHREAD:
                item_hash = (1000003 * type) ^ (lua_Integer)lua_topointer(L, -1);
                break;
            case LUA_TSTRING:
                {
                    // Algorithm taken from Lua 5.3's implementation
                    size_t len;
                    const char *str = lua_tolstring(L, -1, &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;
        }
        ull_hash = (1000003 * ull_hash) ^ item_hash;
        lua_pop(L, 1);
        // Stack: [fields]
    }
    int hash = (int)ull_hash;

    lua_getfield(L, 1, "__instances");
    // Stack: [fields, __instances]
    // Find bucket
    lua_rawgeti(L, -1, hash);
    // Stack: [fields, __instances, bucket]
    if (lua_isnil(L, -1)) {
        // Make a new bucket
        // Stack: [fields, __instances, nil]
        lua_pop(L, 1);
        // Stack: [fields, __instances]
        lua_createtable(L, 0, 1);
        // Stack: [fields, __instances, bucket]
        lua_pushlightuserdata(L, (void*)&WEAK_KEY_METATABLE);
        lua_gettable(L, LUA_REGISTRYINDEX);
        // Stack: [fields, __instances, bucket, {__mode='k'}]
        lua_setmetatable(L, -2);
        // Stack: [fields, __instances, bucket]
        lua_pushvalue(L, -1);
        // Stack: [fields, __instances, bucket, bucket]
        lua_rawseti(L, -3, hash);
        // Stack: [fields, __instances, bucket]
    } else {
        // Stack: [fields, __instances, bucket]
        // scan bucket
        lua_pushnil(L);
        while (lua_next(L, -2) != 0) { // for hash_collider_inst, collider_table in pairs(bucket) do
            // Stack: [fields, __instances, bucket, hash_collider_inst, collider_table]
            // Shallow equality check:
            immutable_info_t *collider_info = (immutable_info_t*)lua_touserdata(L, -2);
            if (collider_info->array_len != array_len) {
                lua_pop(L, 1);
                goto next_bucket_item;
            }
            // Stack: [fields, __instances, bucket, hash_collider_inst, collider_table]
            for (size_t i = 1; i <= num_fields; i++) {
                lua_rawgeti(L, fields_index, i);
                lua_gettable(L, -2);
                IMMUTABLE_GETVALUE(i);
                // Stack: [fields, __instances, bucket, hash_collider_inst, collider_table, collider_val, inst_val]
                if (! lua_rawequal(L, -1, -2)) { // If the i'th entry doesn't match the i'th arg
                    lua_pop(L, 3);
                    goto next_bucket_item;
                } else {
                    lua_pop(L, 2);
                }
            }
            for (size_t i = 1; i <= array_len; i++) {
                lua_rawgeti(L, -1, i);
                IMMUTABLE_GETVALUE(i+num_fields);
                // Stack: [fields, __instances, bucket, hash_collider_inst, collider_table, collider_val, inst_val]
                if (! lua_rawequal(L, -1, -2)) { // If the i'th entry doesn't match the i'th arg
                    lua_pop(L, 3);
                    goto next_bucket_item;
                } else {
                    lua_pop(L, 2);
                }
            }
            // bucket item matches
            // Stack: [fields, __instances, bucket, hash_collider_inst, collider_table]
            lua_pop(L, 1);
            return 1;

            next_bucket_item: ;
        }
    }
    // Stack: [fields, __instances, bucket]
    int bucket_index = lua_gettop(L);

    // Failed to find an existing instance, so create a new one
    immutable_info_t *userdata = (immutable_info_t*)lua_newuserdata(L, sizeof(immutable_info_t));
    // Stack: [fields, __instances, bucket, inst_userdata]
    int userdata_index = lua_gettop(L);
    userdata->hash = hash;
    userdata->array_len = array_len;
    
    lua_pushvalue(L, 1);
    // Stack: [fields, __instances, bucket, inst_userdata, cls]
    lua_setmetatable(L, -2);
    // Stack: [fields, __instances, bucket, inst_userdata]

    // Set up a ref to the bucket so its lifetime is tied to inst_userdata
    lua_getfield(L, 1, "__buckets");
    // Stack: [fields, __instances, bucket, inst_userdata, __buckets]
    lua_pushvalue(L, userdata_index);
    // Stack: [fields, __instances, bucket, inst_userdata, __buckets, inst_userdata]
    lua_pushvalue(L, bucket_index);
    // Stack: [fields, __instances, bucket, inst_userdata, __buckets, inst_userdata, bucket]
    lua_settable(L, -3);
    lua_pop(L, 1);
    // Stack: [fields, __instances, bucket, inst_userdata]
    
    lua_pushvalue(L, userdata_index);
    // Stack: [fields, __instances, bucket, inst_userdata, inst_userdata]
    lua_createtable(L, num_values-num_fields, num_fields);
    for (size_t i=1; i <= num_fields; i++) {
        lua_rawgeti(L, fields_index, i);
        IMMUTABLE_GETVALUE(i);
        lua_settable(L, -3);
    }
    for (size_t i=1; i <= array_len; i++) {
        IMMUTABLE_GETVALUE(i+num_fields);
        lua_rawseti(L, -2, i);
    }
    // Stack: [fields, __instances, bucket, inst_userdata, inst_userdata, inst_table]
    lua_settable(L, -4); // buckets[inst_userdata] = inst_table
    // Stack: [fields, __instances, bucket, inst_userdata]
    lua_gc(L, LUA_GCSTEP, 3);
    return 1;
}

static int Lcreate_instance(lua_State *L)
{
    int num_args = lua_gettop(L)-1;
    lua_getfield(L, 1, "__new");
    if (! lua_isnil(L, -1)) {
        lua_insert(L, 2); // move __new
        lua_pushvalue(L, 1);
        lua_insert(L, 3); // make cls the first argument to __new
        lua_call(L, num_args+1, LUA_MULTRET);
    } else {
        lua_pop(L, 1);
    }
    return _create_instance(L, 0);
}

static int Lfrom_table(lua_State *L)
{
    return _create_instance(L, 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)
{
    luaL_checktype(L, 1, LUA_TUSERDATA);
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    immutable_info_t *info = (immutable_info_t *)lua_touserdata(L, 1);
    lua_pushinteger(L, info->array_len);
    return 1;
}

static int Lindex(lua_State *L)
{
    // Return inst[key], <whether inst has key>
    luaL_checktype(L, 1, LUA_TUSERDATA);
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    int class_index = lua_gettop(L);
    // Stack: [mt]
    immutable_info_t *info = (immutable_info_t*)lua_touserdata(L, 1);
    if (! info) {
        luaL_error(L, "invalid type");
    }

    // Stack: [mt]
    lua_getfield(L, -1, "__instances");
    // Stack: [mt, buckets]
    lua_rawgeti(L, -1, info->hash);
    // Stack: [mt, buckets, bucket]
    if (lua_isnil(L, -1)) {
        luaL_error(L, "Failed to find hash bucket");
    }
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    // Stack: [mt, buckets, bucket, inst_table]
    if (lua_isnil(L, -1)) {
        luaL_error(L, "Failed to find instance");
    }
    lua_pushvalue(L, 2);
    lua_gettable(L, -2);
    if (lua_isnil(L, -1)) {
        // Fall back to class:
        lua_pushvalue(L, 2);
        lua_gettable(L, class_index);
        return 1;
    }
    if (lua_islightuserdata(L, -1) && lua_touserdata(L, -1) == &NIL_SENTINEL) {
        lua_pop(L, 1);
        lua_pushnil(L);
    }
    lua_pushboolean(L, 1);
    return 2;
}

static int Ltostring(lua_State *L)
{
    luaL_Buffer b;
    luaL_buffinit(L, &b);
    if (! lua_getmetatable(L, 1)) {
        luaL_error(L, "invalid type");
    }
    int cls_index = lua_gettop(L);
    // Stack: [mt]
    
    lua_getfield(L, cls_index, "name");
    if (!lua_isnil(L, -1)) {
        luaL_addvalue(&b);
    } else {
        lua_pop(L, 1);
    }
    luaL_addstring(&b, "(");
    
    lua_getfield(L, cls_index, "__instances");
    // Stack: [mt, buckets]
    immutable_info_t *info = (immutable_info_t*)lua_touserdata(L, 1);
    if (! info) {
        luaL_error(L, "invalid type");
    }
    lua_rawgeti(L, -1, info->hash);
    // Stack: [mt, buckets, bucket]
    if (lua_isnil(L, -1)) {
        luaL_error(L, "Failed to find hash bucket");
    }
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    if (lua_isnil(L, -1)) {
        luaL_error(L, "Failed to find instance table");
    }
    int inst_table_index = lua_gettop(L);
    // Stack: [mt, buckets, bucket, inst_table]
    lua_getfield(L, cls_index, "__fields");
    int fields_index = lua_gettop(L);
    lua_getglobal(L, "tostring");
    int tostring_index = lua_gettop(L);
    // Stack: [mt, buckets, bucket, inst_table, __fields, tostring]
    int num_fields = lua_objlen(L, fields_index);
    int n = num_fields + info->array_len, i = 1;
    if (i <= n) {
        goto first_list_item;
    }
    while (++i <= n) {
        // Stack: [mt, buckets, bucket, inst_table, tostring, ???]
        luaL_addstring(&b, ", ");
      first_list_item:
        lua_pushvalue(L, tostring_index);
        // Stack: [mt, buckets, bucket, inst_table, tostring, ???, tostring]
        if (i <= num_fields) {
            lua_rawgeti(L, fields_index, i);
            lua_gettable(L, inst_table_index);
        } else {
            lua_rawgeti(L, inst_table_index, i - num_fields);
        }
        if (lua_islightuserdata(L, -1) && lua_touserdata(L, -1) == &NIL_SENTINEL) {
            lua_pop(L, 1);
            lua_pushnil(L);
        }
        // Stack: [mt, buckets, bucket, inst_table, tostring, ???, tostring, value]
        int quotes = lua_type(L, -1) == LUA_TSTRING;
        lua_call(L, 1, 1);
        if (quotes) {
            luaL_addchar(&b, '"');
        }
        // Stack: [mt, buckets, bucket, inst_table, tostring, ???, value string]
        // TODO: properly escape strings? str:gsub('["\\]',"\\%1"):gsub("\n","\\n")
        luaL_addvalue(&b);
        // Stack: [mt, buckets, bucket, inst_table, tostring, ???]
        if (quotes) {
            luaL_addchar(&b, '"');
        }
    }
    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]
    immutable_info_t *info = (immutable_info_t*)lua_touserdata(L, 1);
    if (! info) {
        luaL_error(L, "invalid type");
    }
    int i = lua_tointeger(L, 2)+1;
    if (i > (int)info->array_len) {
        return 0;
    }
    lua_rawgeti(L, -1, info->hash);
    // Stack: [mt, buckets, bucket]
    if (lua_isnil(L, -1)) {
        luaL_error(L, "Failed to find hash bucket");
    }
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    // Stack: [mt, buckets, bucket, inst_table]
    lua_pushinteger(L, i);
    // Stack: [mt, buckets, bucket, inst_table, i]
    lua_rawgeti(L, -2, i);
    // Stack: [mt, buckets, bucket, inst_table, i, table[i]]
    if (lua_islightuserdata(L, -1) && lua_touserdata(L, -1) == &NIL_SENTINEL) {
        lua_pop(L, 1);
        lua_pushnil(L);
    }
    return 2;
}

static int Lipairs(lua_State *L)
{
    lua_pushcfunction(L, Lnexti);
    // Stack: [Lnexti]
    lua_pushvalue(L, 1);
    // Stack: [Lnexti, inst_udata]
    lua_pushinteger(L, 0);
    // Stack: [Lnexti, inst_udata, 0]
    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]
    immutable_info_t *info = (immutable_info_t*)lua_touserdata(L, 1);
    if (! info) {
        luaL_error(L, "invalid type");
    }
    lua_rawgeti(L, -1, info->hash);
    // Stack: [mt, buckets, bucket]
    if (lua_isnil(L, -1)) {
        luaL_error(L, "Failed to find hash bucket");
    }
    lua_pushvalue(L, 1);
    // Stack: [mt, buckets, bucket, inst_udata]
    lua_rawget(L, -2);
    // Stack: [mt, buckets, bucket, inst_table]
    lua_pushvalue(L, 2);
    // TODO: this is in a random order, and it might be good to force it to be in the same order as __fields
    if (lua_next(L, -2) == 0) {
        return 0;
    }
    if (lua_islightuserdata(L, -1) && lua_touserdata(L, -1) == &NIL_SENTINEL) {
        lua_pop(L, 1);
        lua_pushnil(L);
    }
    return 2;
}

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

static int Lhash(lua_State *L)
{
    luaL_checktype(L, 1, LUA_TUSERDATA);
    immutable_info_t *info = (immutable_info_t *)lua_touserdata(L, 1);
    lua_pushinteger(L, info->hash);
    return 1;
}

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

static int Lmake_class(lua_State *L)
{
    size_t num_args = lua_gettop(L);
    // immutable([fields], [methods/metamethods])
    lua_createtable(L, 0, 24); // Rough guess of number of fields from Rinstance_metamethods + __fields, etc.
    // 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 (num_args >= 2 && 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_insert(L, -2);
            // Stack: [CLS, method_name, method_name, method_value]
            lua_settable(L, -4);
            // Stack: [CLS, method_name]
        }
        // Stack: [CLS]
    }

    lua_pushlightuserdata(L, (void*)&SUPER_METHODS);
    lua_gettable(L, LUA_REGISTRYINDEX);
    lua_setfield(L, -2, "__super");

    // Stack: [CLS]
    lua_createtable(L, 0, 32); // Rough guess: at least 32 instances concurrently
    // Stack: [CLS, __instances]
    lua_pushlightuserdata(L, (void*)&WEAK_VALUE_METATABLE);
    lua_gettable(L, LUA_REGISTRYINDEX);
    // Stack: [CLS, __instances, {__mode='v'}]
    lua_setmetatable(L, -2);
    // Stack: [CLS, __instances]
    lua_setfield(L, -2, "__instances");

    // Stack: [CLS]
    lua_createtable(L, 0, 32); // Rough guess: at least 32 instances concurrently
    // Stack: [CLS, __buckets]
    lua_pushlightuserdata(L, (void*)&WEAK_KEY_METATABLE);
    lua_gettable(L, LUA_REGISTRYINDEX);
    // Stack: [CLS, __buckets, {__mode='k'}]
    lua_setmetatable(L, -2);
    // Stack: [CLS, __buckets]
    lua_setfield(L, -2, "__buckets");

    // Stack: [CLS]
    switch (num_args == 0 ? LUA_TNIL : lua_type(L, 1)) {
        case LUA_TTABLE:
            lua_pushvalue(L, 1);
            break;
        case LUA_TNIL: case LUA_TNONE:
            // If no fields were passed in, set __fields to empty table
            lua_createtable(L, 0, 0);
            break;
        default: {
            luaL_error(L, "immutable expected the fields to be either table or nil");
        }
    }
    lua_setfield(L, -2, "__fields");
    // 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*)&SUPER_METHODS);
    lua_createtable(L, 0, 8);
    luaL_register(L,NULL,Rinstance_metamethods);
    lua_settable(L, LUA_REGISTRYINDEX);

    lua_pushlightuserdata(L, (void*)&WEAK_VALUE_METATABLE);
    lua_createtable(L, 0, 1);
    lua_pushstring(L, "v");
    lua_setfield(L, -2, "__mode");
    lua_settable(L, LUA_REGISTRYINDEX);

    lua_pushlightuserdata(L, (void*)&WEAK_KEY_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;
}