Snapshot.cs

/*
    The C# version of https://github.com/cloudwu/lua-snapshot on https://github.com/topameng/tolua_runtime
*/
using System;
using System.Text;
using LuaInterface;

public static class Snapshot
{
    private readonly static int TABLE = 1;
    private readonly static int FUNCTION = 2;
    private readonly static int SOURCE = 3;
    private readonly static int THREAD = 4;
    private readonly static int USERDATA = 5;
    private readonly static int MARK = 6;
    private readonly static int LUA_MINSTACK = 20;

    private readonly static byte NEWLINE = Convert.ToByte('\n');
    private readonly static string MODE = "__mode";
    private readonly static byte[] MODE_BYTES = System.Text.UTF8Encoding.Default.GetBytes("__mode");

    private static void luaL_checkversion(IntPtr L)
    {
        if (LuaDLL.lua_pushthread(L) == 0)
        {
            LuaDLL.tolua_error(L, "Must require in main thread");
        }
        LuaDLL.lua_setfield(L, LuaIndexes.LUA_REGISTRYINDEX, "mainthread");
    }

    private static void lua_rawsetp(IntPtr L, int idx, IntPtr lightuserdata)
    {
        if (idx < 0)
            idx += LuaDLL.lua_gettop(L) + 1;

        LuaDLL.lua_pushlightuserdata(L, lightuserdata);
        LuaDLL.lua_insert(L, -2);
        LuaDLL.lua_rawset(L, idx);
    }

    private static void lua_rawgetp(IntPtr L, int idx, IntPtr lightuserdata)
    {
        if (idx < 0)
            idx += LuaDLL.lua_gettop(L) + 1;

        LuaDLL.lua_pushlightuserdata(L, lightuserdata);
        LuaDLL.lua_rawget(L, idx);
    }

    private static void lua_getuservalue(IntPtr L, int idx)
    {
        LuaDLL.lua_getfenv(L, idx);
    }

    private static void mark_function_env(IntPtr L, IntPtr dL, IntPtr t)
    {
        LuaDLL.lua_getfenv(L, -1);
        if (LuaDLL.lua_istable(L, -1))
        {
            mark_object(L, dL, t, "[environment]");
        }
        else
        {
            LuaDLL.lua_pop(L, 1);
        }
    }


    private static bool ismarked(IntPtr dL, IntPtr p)
    {
        lua_rawgetp(dL, MARK, p);
        if (LuaDLL.lua_isnil(dL, -1))
        {
            LuaDLL.lua_pop(dL, 1);
            LuaDLL.lua_pushboolean(dL, 1);
            lua_rawsetp(dL, MARK, p);
            return false;
        }

        LuaDLL.lua_pop(dL, 1);
        return true;
    }

    private static IntPtr readobject(IntPtr L, IntPtr dL, IntPtr parent, string desc)
    {
        LuaTypes t = LuaDLL.lua_type(L, -1);
        int tidx = 0;

        switch (t)
        {
            case LuaTypes.LUA_TTABLE:
                tidx = TABLE;
                break;
            case LuaTypes.LUA_TFUNCTION:
                tidx = FUNCTION;
                break;
            case LuaTypes.LUA_TUSERDATA:
                tidx = USERDATA;
                break;
            case LuaTypes.LUA_TTHREAD:
                tidx = THREAD;
                break;
            default:
                return IntPtr.Zero;
        }

        IntPtr p = LuaDLL.lua_topointer(L, -1);
        if (ismarked(dL, p))
        {
            lua_rawgetp(dL, tidx, p);
            if (!LuaDLL.lua_isnil(dL, -1))
            {
                LuaDLL.lua_pushstring(dL, desc);
                lua_rawsetp(dL, -2, parent);
            }

            LuaDLL.lua_pop(dL, 1);
            LuaDLL.lua_pop(L, 1);
            return IntPtr.Zero;
        }

        LuaDLL.lua_newtable(dL);
        LuaDLL.lua_pushstring(dL, desc);
        lua_rawsetp(dL, -2, parent);
        lua_rawsetp(dL, tidx, p);

        return p;
    }

    private static string keystring(IntPtr L, int index)
    {
        var sb = new StringBuilder();
        LuaTypes t = LuaDLL.lua_type(L, index);
        switch (t)
        {
            case LuaTypes.LUA_TNIL:
                sb.Append("[nil]");
                break;
            case LuaTypes.LUA_TBOOLEAN:
                sb.AppendFormat("[{0}]", LuaDLL.lua_toboolean(L, index));
                break;
            case LuaTypes.LUA_TNUMBER:
                sb.AppendFormat("[{0}]", LuaDLL.lua_tonumber(L, index));
                break;
            case LuaTypes.LUA_TSTRING:
                return LuaDLL.lua_tostring(L, index);
            default:
                sb.AppendFormat("[{0}:{1}]", LuaDLL.lua_typename(L, t), LuaDLL.lua_topointer(L, index));
                break;
        }

        return sb.ToString();
    }

    private static void mark_table(IntPtr L, IntPtr dL, IntPtr parent, string desc)
    {
        IntPtr t = readobject(L, dL, parent, desc);
        if (t == IntPtr.Zero)
            return;

        bool weak_k = false;
        bool weak_v = false;

        if (LuaDLL.lua_getmetatable(L, -1) != 0)
        {
            LuaDLL.lua_pushstring(L, MODE);
            LuaDLL.lua_rawget(L, -2);
            if (LuaDLL.lua_isstring(L, -1) == 1)
            {
                string mode = LuaDLL.lua_tostring(L, -1);
                if (mode.Contains("k"))
                    weak_k = true;

                if (mode.Contains("v"))
                    weak_v = true;
            }

            LuaDLL.lua_pop(L, 1);

            LuaDLL.lua_checkstack(L, LUA_MINSTACK);
            mark_table(L, dL, t, "[metatable]");
        }

        LuaDLL.lua_pushnil(L);

        while (LuaDLL.lua_next(L, -2) != 0)
        {
            if (weak_v)
            {
                LuaDLL.lua_pop(L, 1);
            }
            else
            {
                string desc2 = keystring(L, -2);
                mark_object(L, dL, t, desc2);
            }

            if (!weak_k)
            {
                LuaDLL.lua_pushvalue(L, -1);
                mark_object(L, dL, t, "[key]");
            }
        }

        LuaDLL.lua_pop(L, 1);
    }

    private static void mark_userdata(IntPtr L, IntPtr dL, IntPtr parent, string desc)
    {
        IntPtr t = readobject(L, dL, parent, desc);
        if (t == IntPtr.Zero)
            return;

        if (LuaDLL.lua_getmetatable(L, -1) == 1)
        {
            mark_table(L, dL, t, "[metatable]");
        }

        lua_getuservalue(L, -1);
        if (LuaDLL.lua_isnil(L, -1))
        {
            LuaDLL.lua_pop(L, 2);
        }
        else
        {
            mark_table(L, dL, t, "[uservalue]");
            LuaDLL.lua_pop(L, 1);
        }
    }

    private static void mark_function(IntPtr L, IntPtr dL, IntPtr parent, string desc)
    {
        IntPtr t = readobject(L, dL, parent, desc);
        if (t == IntPtr.Zero)
            return;

        mark_function_env(L, dL, t);

        int i = 0;
        for (; ; i++)
        {
            string name = LuaDLL.lua_getupvalue(L, -1, i);
            if (name == null || name == string.Empty)
                break;

            mark_object(L, dL, t, name.Length != 0 ? name : "[upvalue]");
        }

        if (LuaDLL.lua_iscfunction(L, -1) == 1)
        {
            if (i == 1)
            {
                LuaDLL.lua_pushnil(dL);
                lua_rawsetp(dL, FUNCTION, t);
            }

            LuaDLL.lua_pop(L, 1);
        }
        else
        {
            Lua_Debug ar = new Lua_Debug();
            LuaDLL.lua_getinfo(L, ">S", ref ar);
            IntPtr b = LuaDLL.tolua_buffinit(dL);
            LuaDLL.tolua_addstring(b, ar.short_src);
            LuaDLL.tolua_addstring(b, string.Format(":{0}", ar.linedefined));
            LuaDLL.tolua_pushresult(b);
            lua_rawsetp(dL, SOURCE, t);
        }
    }

    private static void mark_thread(IntPtr L, IntPtr dL, IntPtr parent, string desc)
    {
        IntPtr t = readobject(L, dL, parent, desc);
        if (t == IntPtr.Zero)
            return;

        int level = 0;
        IntPtr cL = LuaDLL.lua_tothread(L, -1);
        if (cL.Equals(L))
        {
            level = 1;
        }
        else
        {
            int top = LuaDLL.lua_gettop(cL);
            LuaDLL.lua_checkstack(cL, 1);

            for (int i = 0; i < top; i++)
            {
                LuaDLL.lua_pushvalue(cL, i + 1);
                mark_object(cL, dL, cL, string.Format("[{0}]", i + 1));
            }
        }

        Lua_Debug ar = new Lua_Debug();
        IntPtr b = LuaDLL.tolua_buffinit(dL);
        while (LuaDLL.lua_getstack(cL, level, ref ar) == 1)
        {
            LuaDLL.lua_getinfo(cL, "Sl", ref ar);
            LuaDLL.tolua_addstring(b, ar.short_src);
            if (ar.currentline >= 0)
            {
                LuaDLL.tolua_addstring(b, string.Format(":{0}", ar.currentline));
            }

            for (int j = 1; j > -1; j -= 2)
            {
                for (int i = j; ; i += j)
                {
                    string name = LuaDLL.lua_getlocal(cL, ref ar, i);
                    if (name == null)
                        break;

                    mark_object(cL, dL, t,
                        string.Format("{0} : {1}:{2}", name, ar.short_src, ar.currentline));
                }
            }

            ++level;
        }

        LuaDLL.tolua_pushresult(b);
        lua_rawsetp(dL, SOURCE, t);
        LuaDLL.lua_pop(L, 1);
    }

    private static void mark_object(IntPtr L, IntPtr dL, IntPtr parent, string desc)
    {
        LuaDLL.lua_checkstack(L, LUA_MINSTACK);
        var t = LuaDLL.lua_type(L, -1);
        switch (t)
        {
            case LuaTypes.LUA_TTABLE:
                mark_table(L, dL, parent, desc);
                break;
            case LuaTypes.LUA_TFUNCTION:
                mark_function(L, dL, parent, desc);
                break;
            case LuaTypes.LUA_TUSERDATA:
                mark_userdata(L, dL, parent, desc);
                break;
            case LuaTypes.LUA_TTHREAD:
                mark_thread(L, dL, parent, desc);
                break;
            default:
                LuaDLL.lua_pop(L, 1);
                break;
        }
    }

    private static int count_table(IntPtr L, int idx)
    {
        int n = 0;
        LuaDLL.lua_pushnil(L);
        while (LuaDLL.lua_next(L, idx) != 0)
        {
            ++n;
            LuaDLL.lua_pop(L, 1);
        }

        return n;
    }

    // b -> luaL_Buffer
    private static void gen_table_desc(IntPtr b, IntPtr parent, string desc)
    {
        string addr = string.Format("{0} : ", parent.ToString());
        LuaDLL.tolua_addlstring(b, addr, addr.Length);
        LuaDLL.tolua_addstring(b, desc);
        LuaDLL.tolua_addchar(b, NEWLINE);
    }

    private static void pdesc(IntPtr L, IntPtr dL, int idx, string typename)
    {
        LuaDLL.lua_pushnil(dL);
        while (LuaDLL.lua_next(dL, idx) != 0)
        {
            IntPtr b = LuaDLL.tolua_buffinit(L);
            IntPtr key = LuaDLL.lua_touserdata(dL, -2);
            if (idx == FUNCTION)
            {
                lua_rawgetp(dL, SOURCE, key);
                if (LuaDLL.lua_isnil(dL, -1))
                {
                    LuaDLL.tolua_addstring(b, "cfunction\n");
                }
                else
                {
                    string s = System.Text.UTF8Encoding.Default.GetString(ToLua.ToByteBuffer(dL, -1));
                    LuaDLL.tolua_addlstring(b, s, s.Length);
                    LuaDLL.tolua_addchar(b, NEWLINE);
                }

                LuaDLL.lua_pop(dL, 1);
            }
            else if (idx == THREAD)
            {
                lua_rawgetp(dL, SOURCE, key);
                string s = System.Text.UTF8Encoding.Default.GetString(ToLua.ToByteBuffer(dL, -1));
                LuaDLL.tolua_addlstring(b, s, s.Length);
                LuaDLL.tolua_addchar(b, NEWLINE);
                LuaDLL.lua_pop(dL, 1);
            }
            else
            {
                LuaDLL.tolua_addstring(b, typename);
                LuaDLL.tolua_addchar(b, NEWLINE);
            }

            LuaDLL.lua_pushnil(dL);

            while (LuaDLL.lua_next(dL, -2) != 0)
            {
                int l = 0;
                IntPtr parent = LuaDLL.lua_touserdata(dL, -2);
                string desc = LuaDLL.luaL_checklstring(dL, -1, out l);
                gen_table_desc(b, parent, desc);
                LuaDLL.lua_pop(dL, 1);
            }

            LuaDLL.tolua_pushresult(b);
            lua_rawsetp(L, -2, key);
            LuaDLL.lua_pop(dL, 1);
        }
    }

    private static void gen_result(IntPtr L, IntPtr dL)
    {
        int count = 0;
        count += count_table(dL, TABLE);
        count += count_table(dL, FUNCTION);
        count += count_table(dL, USERDATA);
        count += count_table(dL, THREAD);

        LuaDLL.lua_createtable(L, 0, count);
        pdesc(L, dL, TABLE,     "table");
        pdesc(L, dL, FUNCTION,  "function");
        pdesc(L, dL, USERDATA,  "userdata");
        pdesc(L, dL, THREAD,    "thread");
    }

    private static int _snapshot(IntPtr L)
    {
        int i;
        IntPtr dL = LuaDLL.luaL_newstate();
        for (i = 0; i < MARK; i++)
        {
            LuaDLL.lua_newtable(dL);
        }

        LuaDLL.lua_pushvalue(L, LuaIndexes.LUA_REGISTRYINDEX);
        mark_table(L, dL, IntPtr.Zero, "[registry]");
        gen_result(L, dL);
        LuaDLL.lua_close(dL);
        return 1;
    }

    public static int snapshot()
    {
        IntPtr currentL = LuaClient.GetMainState().GetL();
        luaL_checkversion(currentL);
        return _snapshot(currentL);
    }
}