diff --git a/luprex/core/cpp/luastack.cpp b/luprex/core/cpp/luastack.cpp index 266cfca9..9c6e3529 100644 --- a/luprex/core/cpp/luastack.cpp +++ b/luprex/core/cpp/luastack.cpp @@ -151,6 +151,11 @@ void LuaStack::newtable(LuaSlot target) const { lua_replace(L_, target); } +void LuaStack::createtable(LuaSlot target, int narr, int nrec) const { + lua_createtable(L_, narr, nrec); + lua_replace(L_, target); +} + lua_State *LuaStack::newthread(LuaSlot target) const { lua_State *result = lua_newthread(L_); lua_replace(L_, target); @@ -252,6 +257,10 @@ void LuaStack::cleartable(LuaSlot tab) const { } } +int LuaStack::rawlen(LuaSlot obj) const { + return lua_rawlen(L_, obj.index()); +} + void LuaStack::check_nret(int xnret, int otop, int nret) const { int ntop = lua_gettop(L_); if ((nret != xnret)||(ntop != otop + xnret)) { diff --git a/luprex/core/cpp/luastack.hpp b/luprex/core/cpp/luastack.hpp index f90abe08..4e3912a7 100644 --- a/luprex/core/cpp/luastack.hpp +++ b/luprex/core/cpp/luastack.hpp @@ -392,11 +392,14 @@ public: void checknometa(LuaSlot index) const; void newtable(LuaSlot target) const; + void createtable(LuaSlot target, int narr, int nrec) const; lua_State *newthread(LuaSlot target) const; void getglobaltable(LuaSlot gltab) const; void makesubtable(LuaSlot sub, LuaSlot tab, const char *name) const; void cleartable(LuaSlot tab) const; + int rawlen(LuaSlot val) const; + int next(LuaSlot tab, LuaSlot key, LuaSlot value) const; void makeclass(LuaSlot tab, LuaSlot name) const; @@ -440,6 +443,12 @@ public: pop_any_value(target); } + template + void rawgeti(RT &target, LuaSlot tab, lua_Integer key) const { + lua_rawgeti(L_, tab, key); + pop_any_value(target); + } + template void rawset(LuaSlot tab, KT key, VT value) const { push_any_value(key); @@ -447,6 +456,12 @@ public: lua_rawset(L_, tab); } + template + void rawseti(LuaSlot tab, lua_Integer key, VT value) const { + push_any_value(value); + lua_rawseti(L_, tab, key); + } + // template // void rawset(LuaSlot tab, const char *field, VT value) const { // push_any_value(value); diff --git a/luprex/core/cpp/table.cpp b/luprex/core/cpp/table.cpp index 4534afce..2703c3ce 100644 --- a/luprex/core/cpp/table.cpp +++ b/luprex/core/cpp/table.cpp @@ -1,6 +1,14 @@ #include "table.hpp" #include "source.hpp" +LuaDefine(table_getregistry, "f") { + LuaArg key; + LuaRet result; + LuaStack LS(L, key, result); + LS.rawget(result, LuaRegistry, key); + return LS.result(); +} + LuaDefine(table_equal, "c") { LuaArg t1, t2; LuaRet eql; @@ -196,10 +204,215 @@ LuaDefine(queue_nth, "c") { return LS.result(); } -LuaDefine(table_getregistry, "f") { - LuaArg key; - LuaRet result; - LuaStack LS(L, key, result); - LS.rawget(result, LuaRegistry, key); +///////////////////////////////////////////////////////////// +// +// table_sortedpairs +// +// Given a table, return all the pairs in the table as a sequence. +// The resulting sequence contains a 1 in the first position, +// followed by the pairs, like so: +// +// sortedpairs({a=1,b=2,c=3}) => {1,"a",1,"b",2,"c",3} +// +// Note that this function is not directly exposed to lua. +// +///////////////////////////////////////////////////////////// + +static void pushkey (lua_State *L, int tab, int i) { + lua_rawgeti(L, tab, i*2); +} + +static void push2keys (lua_State *L, int tab, int i, int j) { + lua_rawgeti(L, tab, i*2); + lua_rawgeti(L, tab, j*2); +} + +static void pop2keys (lua_State *L, int tab, int i, int j) { + lua_rawseti(L, tab, i*2); + lua_rawseti(L, tab, j*2); +} + +static void swap2values (lua_State *L, int tab, int i, int j) { + lua_rawgeti(L, tab, i*2+1); + lua_rawgeti(L, tab, j*2+1); + lua_rawseti(L, tab, i*2+1); + lua_rawseti(L, tab, j*2+1); +} + +static void auxsort (lua_State *L, int tab, int l, int u) { + while (l < u) { /* for tail recursion */ + int i, j; + /* sort elements a[l], a[(l+u)/2] and a[u] */ + push2keys(L, tab, l, u); + if (lua_genlt(L, -1, -2)) { + pop2keys(L, tab, l, u); + swap2values(L, tab, l, u); + } else { + lua_pop(L, 2); + } + if (u - l == 1) break; /* only 2 elements */ + i = (l + u) / 2; + push2keys(L, tab, i, l); + if (lua_genlt(L, -2, -1)) { + pop2keys(L, tab, i, l); + swap2values(L, tab, i, l); + } else { + lua_pop(L, 1); /* remove a[l] */ + pushkey(L, tab, u); + if (lua_genlt(L, -1, -2)) { + pop2keys(L, tab, i, u); + swap2values(L, tab, i, u); + } else { + lua_pop(L, 2); + } + } + if (u - l == 2) break; /* only 3 elements */ + /* put the pivot value on top of the stack and keep it there */ + pushkey(L, tab, i); + /* move the pivot from i to u-1 */ + lua_pushvalue(L, -1); + pushkey(L, tab, u-1); + pop2keys(L, tab, i, u-1); + swap2values(L, tab, i, u-1); + /* a[l] <= P == a[u-1] <= a[u], only need to sort from l+1 to u-2 */ + i = l; + j = u - 1; + for (;;) { /* invariant: a[l..i] <= P <= a[j..u] */ + /* repeat ++i until a[i] >= P */ + while (pushkey(L, tab, ++i), lua_genlt(L, -1, -2)) { + if (i >= u) luaL_error(L, "invalid order function for sorting"); + lua_pop(L, 1); /* remove a[i] */ + } + /* repeat --j until a[j] <= P */ + while (pushkey(L, tab, --j), lua_genlt(L, -3, -1)) { + if (j <= l) luaL_error(L, "invalid order function for sorting"); + lua_pop(L, 1); /* remove a[j] */ + } + if (j < i) { + lua_pop(L, 3); /* pop pivot, a[i], a[j] */ + break; + } + pop2keys(L, tab, i, j); + swap2values(L, tab, i, j); + } + push2keys(L, tab, u-1, i); + pop2keys(L, tab, u-1, i); + swap2values(L, tab, u-1, i); + /* a[l..i-1] <= a[i] == P <= a[i+1..u] */ + /* adjust so that smaller half is in [j..i] and larger one in [l..u] */ + if (i - l < u - i) { + j = l; + i = i - 1; + l = i + 2; + } else { + j = i + 1; + i = u; + u = j - 2; + } + auxsort(L, tab, j, i); /* call recursively the smaller one */ + } /* repeat the routine for the larger one */ +} + +int table_sortedpairs(lua_State *L, bool *unsortable) { + lua_checkstack(L, 40); + LuaArg tab; + LuaVar key, value; + LuaRet pairs; + LuaStack LS(L, tab, key, value, pairs); + if (unsortable != nullptr) { + *unsortable = false; + } + LS.checktable(tab); + // Count the total number of pairs. + // TODO: add lua_npairs to make this faster. + lua_pushnil(L); + int total = 0; + while (lua_next(L, tab.index()) != 0) { + total += 1; + lua_pop(L, 1); + } + // Create the table, store the initial 1. + LS.createtable(pairs, total * 2 + 1, 0); + LS.rawseti(pairs, 1, 1); + // Transfer the pairs into the sequence. + lua_pushnil(L); + int offset = 2; + while (lua_next(L, tab.index()) != 0) { + int ktype = lua_type(L, -2); + if (ktype == LUA_TNUMBER || ktype == LUA_TSTRING || ktype == LUA_TBOOLEAN) { + lua_pushvalue(L, -2); // K V K + lua_rawseti(L, pairs.index(), offset++); + lua_rawseti(L, pairs.index(), offset++); + } else { + if (unsortable != nullptr) { + *unsortable = true; + } + lua_pop(L, 1); + } + } + auxsort(L, pairs.index(), 1, total); + return LS.result(); +} + +///////////////////////////////////////////////////////////// +// +// Given a sortedpairs vector, return the (key, value) pairs +// one by one. The first element of the vector is used as a +// counter to keep track of our position. +// +///////////////////////////////////////////////////////////// + +static int nextsortedpair(lua_State *L) { + if (lua_gettop(L) < 2) { + luaL_error(L, "Not enough arguments to nextpair"); + } + luaL_checktype(L, 1, LUA_TTABLE); + lua_rawgeti(L, 1, 1); + lua_Integer i = luaL_checkinteger(L, -1); + lua_pop(L, 1); + lua_pushinteger(L, i+1); + lua_rawseti(L, 1, 1); + lua_rawgeti(L, 1, i*2); + if (lua_isnil(L, -1)) { + return 1; + } else { + lua_rawgeti(L, 1, i*2+1); + return 2; + } +} + +///////////////////////////////////////////////////////////// +// +// Replace the lua 'pairs' and 'next' iterators. +// +///////////////////////////////////////////////////////////// + +LuaDefine(table_pairs, "f") { + LuaArg tab; + LuaRet closure, rtab, key; + LuaStack LS(L, tab, closure, rtab, key); + lua_pushvalue(L, tab.index()); + bool unsortable; + table_sortedpairs(L, &unsortable); + if (unsortable) { + luaL_error(L, "Cannot iterate over a table with unsortable keys"); + } + lua_replace(L, rtab.index()); + LS.set(closure, nextsortedpair); + LS.set(key, LuaNil); + return LS.result(); +} + +LuaDefine(table_next, "f") { + luaL_error(L, "The 'next' iterator is not allowed in this version of lua"); + return 0; +} + +LuaDefine(table_genlt, "f") { + LuaArg o1,o2; + LuaRet lt; + LuaStack LS(L, o1, o2, lt); + int ltf = lua_genlt(L, o1.index(), o2.index()); + LS.set(lt, ltf ? true:false); return LS.result(); } diff --git a/luprex/core/cpp/table.hpp b/luprex/core/cpp/table.hpp index 84b784c4..f5a23b7e 100644 --- a/luprex/core/cpp/table.hpp +++ b/luprex/core/cpp/table.hpp @@ -72,6 +72,17 @@ int table_empty(lua_State *L); // int table_count(lua_State *L); +// table_sortedpairs +// +// Return a vector containing pairs from the table. +// This is not exposed directly to lua, but it is used +// in the new version of the 'pairs' iterator. +// +// The boolean 'unsortable' is set to indicate whether +// or not the table contains unsortable values. +// +int table_sortedpairs(lua_State *L, bool *unsortable); + // table_clear // // Remove all key/value pairs from the table. Does not