Reimplement the 'pairs' iterator and 'next' iterator

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)) {

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<typename RT>
+    void rawgeti(RT &target, LuaSlot tab, lua_Integer key) const {
+        lua_rawgeti(L_, tab, key);
+        pop_any_value(target);
+    }
+
     template<typename KT, typename VT>
     void rawset(LuaSlot tab, KT key, VT value) const {
         push_any_value(key);
@@ -447,6 +456,12 @@ public:
         lua_rawset(L_, tab);
     }
     
+    template<typename VT>
+    void rawseti(LuaSlot tab, lua_Integer key, VT value) const {
+        push_any_value(value);
+        lua_rawseti(L_, tab, key);
+    }
+
     // template<typename VT>
     // void rawset(LuaSlot tab, const char *field, VT value) const {
     //     push_any_value(value);

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();
 }

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