Merge pull request #7485 from xmake-io/pairs

improve pairs for lua5.5

6 files changed, 306 insertions, 52 deletions

diff --git a/tests/modules/for_loop/test.lua b/tests/modules/for_loop/test.lua
new file mode 100644
index 000000000..51d5f075a
--- /dev/null
+++ b/tests/modules/for_loop/test.lua
@@ -0,0 +1,189 @@
+-- Regression tests for the Lua 5.4+ for-loop hazard: the parser merges
+-- the generic-for control slot with the user's first loop variable and
+-- marks it `RDKCONST`, so code like
+--
+--     for k, v in pairs(t) do
+--         k = k:gsub("_", "-")
+--         ...
+--     end
+--
+-- normally fails to compile with "attempt to assign to const variable",
+-- and if we only relax the const check the runtime silently corrupts the
+-- iterator state (next iteration gets the modified key and `next` bails
+-- with "invalid key to 'next'" — or worse, silently skips entries).
+--
+-- xmake works around this in two coordinated places:
+--
+--   1. Compile-time: `core/src/lua/xmake.lua` (and xmake.sh) replace
+--      `RDKCONST);` with `VDKREG);` in lparser.c so assignment to the
+--      first loop variable is allowed through.
+--
+--   2. Runtime: the sandbox iterators (`pairs`, `ipairs`, `irpairs`) and
+--      the base containers (`list:items/ritems`, `hashset:items/keys`)
+--      are implemented as stateful closures keeping their cursor in an
+--      upvalue, so user writes to the first loop variable cannot reach
+--      the iterator's internal state.
+--
+-- These tests pin both behaviors: that reassigning the first loop
+-- variable (a) compiles and (b) produces correct, non-lossy iteration.
+
+import("core.base.list")
+import("core.base.hashset")
+
+function test_numeric_for_reassign(t)
+    -- The point of this test is purely compile-time: without the
+    -- RDKCONST->VDKREG replace in core/src/lua/xmake.lua the write to
+    -- `i` would raise "attempt to assign to const variable" and this
+    -- file wouldn't even parse.
+    --
+    -- We do NOT assert on the value of `i` during iteration. Lua 5.4
+    -- merges the numeric-for control slot with the user's first loop
+    -- variable while 5.5 splits them, so the observable sequence of
+    -- `i` values after a reassignment differs between versions. The
+    -- iteration count lives in its own slot and stays stable, so that
+    -- is what we pin.
+    local n = 0
+    for i = 1, 5 do
+        i = i * 10
+        n = n + 1
+    end
+    t:require(n == 5)
+end
+
+function test_generic_for_reassign_key(t)
+    local tbl = {foo_a = 1, foo_b = 2, foo_c = 3}
+    local seen = {}
+    local count = 0
+    for k, v in pairs(tbl) do
+        k = k:gsub("_", "-")
+        seen[k] = v
+        count = count + 1
+    end
+    t:require(count == 3)
+    t:require(seen["foo-a"] == 1)
+    t:require(seen["foo-b"] == 2)
+    t:require(seen["foo-c"] == 3)
+end
+
+function test_generic_for_reassign_many_keys(t)
+    -- Stress the iterator with enough keys that a broken stateless
+    -- iterator would deterministically trip `next` on the second round.
+    local tbl = {}
+    for i = 1, 64 do
+        tbl["key_" .. i] = i
+    end
+    local total = 0
+    for name, value in pairs(tbl) do
+        name = name:gsub("_", "-")  -- would corrupt `next`'s key arg
+        t:require(name:find("^key%-%d+$") ~= nil)
+        total = total + value
+    end
+    t:require(total == (1 + 64) * 64 / 2)
+end
+
+function test_list_items_reassign(t)
+    local l = list.new()
+    for i = 1, 5 do
+        l:push({name = "n" .. i})
+    end
+    for item in l:items() do
+        item = nil  -- would corrupt list:next on the next iteration
+    end
+    local names = {}
+    for item in l:items() do
+        table.insert(names, item.name)
+    end
+    t:require(#names == 5)
+    t:require(names[1] == "n1" and names[5] == "n5")
+end
+
+function test_irpairs_reassign_index(t)
+    local arr = {"a", "b", "c", "d", "e"}
+    local collected = {}
+    for i, v in irpairs(arr) do
+        i = -1  -- would corrupt the index on the next iteration
+        table.insert(collected, v)
+    end
+    t:require(#collected == 5)
+    t:require(collected[1] == "e")
+    t:require(collected[5] == "a")
+end
+
+function test_hashset_items_reassign(t)
+    local set = hashset.from({"key_1", "key_2", "key_3", "key_4", "key_5"})
+    local seen = {}
+    for item in set:items() do
+        item = item:gsub("_", "-")  -- would corrupt `next`'s key arg
+        seen[item] = true
+    end
+    local count = 0
+    for _ in pairs(seen) do count = count + 1 end
+    t:require(count == 5)
+    for i = 1, 5 do
+        t:require(seen["key-" .. i] == true)
+    end
+end
+
+function test_hashset_items_skip_nil_member(t)
+    -- A nil member lives in hashset under the `_NIL` sentinel; iteration
+    -- must skip it rather than terminate, so real entries after it are
+    -- still visited regardless of `next`'s order.
+    local set = hashset.new()
+    set:insert("a")
+    set:insert(nil)
+    set:insert("b")
+    set:insert("c")
+    local seen = {}
+    for item in set:items() do
+        seen[item] = true
+    end
+    local count = 0
+    for _ in pairs(seen) do count = count + 1 end
+    t:require(count == 3)
+    t:require(seen.a and seen.b and seen.c)
+end
+
+function test_hashset_orderitems_skip_nil_member(t)
+    -- use numeric members so the existing orderitems sort (which coerces
+    -- `_NIL` to math.inf) stays within one comparable type.
+    local set = hashset.new()
+    set:insert(1)
+    set:insert(nil)
+    set:insert(2)
+    set:insert(3)
+    local collected = {}
+    for item in set:orderitems() do
+        table.insert(collected, item)
+    end
+    t:require(#collected == 3)
+    t:require(collected[1] == 1 and collected[2] == 2 and collected[3] == 3)
+end
+
+function test_ipairs_reassign_index(t)
+    -- With the stock Lua 5.4 `ipairs`, writing to the first loop variable
+    -- would silently shift the index on the next iteration (no error,
+    -- just wrong results). Sandbox `ipairs` hides the counter in an
+    -- upvalue so the body's write is harmless.
+    local list = {10, 20, 30, 40, 50}
+    local seen = {}
+    for i, v in ipairs(list) do
+        i = -1  -- would corrupt iteration if `i` were the control slot
+        table.insert(seen, v)
+    end
+    t:require(#seen == 5)
+    for idx = 1, 5 do
+        t:require(seen[idx] == list[idx])
+    end
+end
+
+function test_generic_for_reassign_value(t)
+    -- Writing to the second loop variable is always safe (it isn't the
+    -- iterator control), but exercise it anyway to pin the behavior.
+    local tbl = {a = 1, b = 2, c = 3}
+    local total = 0
+    for _, v in pairs(tbl) do
+        v = v * 2
+        total = total + v
+    end
+    t:require(total == 12)
+end
diff --git a/xmake/core/base/hashset.lua b/xmake/core/base/hashset.lua
index c052d36ee..405a23474 100644
--- a/xmake/core/base/hashset.lua
+++ b/xmake/core/base/hashset.lua
@@ -119,15 +119,27 @@ end
 -- end
 -- @endcode
 --
+-- Stateful closure so the loop body can safely reassign the first loop
+-- variable under lua 5.4+ (paired with the RDKCONST->VDKREG compile-time
+-- patch in core/src/lua/xmake.lua).
 function hashset:items()
-    return function (t, item)
-        local k, _ = next(t._DATA, item)
-        if k == hashset._NIL then
-            return nil
-        else
-            return k
-        end
-    end, self, nil
+    -- keep `next`'s key in an upvalue so the loop body can safely reassign
+    -- the first loop variable. In lua 5.4+ the for-in control slot is
+    -- merged with the first user variable; threading the key through the
+    -- loop would otherwise corrupt `next` on the following iteration.
+    --
+    -- nil-as-a-member is stored under the `_NIL` sentinel. For-loop
+    -- semantics don't let us yield nil (it would end the loop), so we
+    -- skip the sentinel and continue to the next real key; the nil
+    -- member is omitted but entries after it are still visited.
+    local data = self._DATA
+    local k = nil
+    return function ()
+        repeat
+            k = next(data, k)
+        until k ~= hashset._NIL
+        return k
+    end
 end
 
 -- iterate order items
@@ -154,16 +166,18 @@ function hashset:orderitems()
         end
         return a < b
     end)
-    local i = 1
-    return function (t, k)
-        k = orderkeys[i]
-        i = i + 1
-        if k == hashset._NIL then
-            return nil
-        else
-            return k
-        end
-    end, self, nil
+    -- see hashset:items() for the `_NIL` handling rationale
+    local n = #orderkeys
+    local i = 0
+    return function ()
+        local k
+        repeat
+            i = i + 1
+            if i > n then return nil end
+            k = orderkeys[i]
+        until k ~= hashset._NIL
+        return k
+    end
 end
 
 -- iterate keys (deprecated, please use items())
@@ -175,14 +189,19 @@ end
 -- @endcode
 --
 function hashset:keys()
-    return function (t, key)
-        local k, _ = next(t._DATA, key)
-        if k == hashset._NIL then
-            return k, nil
-        else
-            return k, k
+    -- see hashset:items() for the stateful-closure rationale and the
+    -- `_NIL` skipping behavior.
+    local data = self._DATA
+    local k = nil
+    return function ()
+        repeat
+            k = next(data, k)
+        until k ~= hashset._NIL
+        if k == nil then
+            return nil
         end
-    end, self, nil
+        return k, k
+    end
 end
 
 -- iterate order keys (deprecated, please use orderitems())
@@ -194,6 +213,8 @@ end
 -- @endcode
 --
 function hashset:orderkeys()
+    -- see hashset:items() for the stateful-closure rationale and the
+    -- `_NIL` skipping behavior.
     local orderkeys = table.keys(self._DATA)
     table.sort(orderkeys, function (a, b)
         if a == hashset._NIL then
@@ -210,16 +231,17 @@ function hashset:orderkeys()
         end
         return a < b
     end)
-    local i = 1
-    return function (t, k)
-        k = orderkeys[i]
-        i = i + 1
-        if k == hashset._NIL then
-            return k, nil
-        else
-            return k, k
-        end
-    end, self, nil
+    local n = #orderkeys
+    local i = 0
+    return function ()
+        local k
+        repeat
+            i = i + 1
+            if i > n then return nil end
+            k = orderkeys[i]
+        until k ~= hashset._NIL
+        return k, k
+    end
 end
 
 -- get size of hashset
diff --git a/xmake/core/base/list.lua b/xmake/core/base/list.lua
index 4caed81bc..222a6e85d 100644
--- a/xmake/core/base/list.lua
+++ b/xmake/core/base/list.lua
@@ -238,22 +238,31 @@ end
 --
 -- @return      the iterator function
 --
+-- Stateful closure so the loop body can safely reassign the first loop
+-- variable under lua 5.4+ (paired with the RDKCONST->VDKREG compile-time
+-- patch in core/src/lua/xmake.lua).
 function list:items()
-    local iter = function (list, item)
-        return list:next(item)
+    -- stateful closure: keep the cursor in an upvalue so the loop body
+    -- can safely reassign the first loop variable (lua 5.4+ merges the
+    -- for-in control slot with the first user variable).
+    local item = nil
+    return function ()
+        item = self:next(item)
+        return item
     end
-    return iter, self, nil
 end
 
 -- iterate elements from back to front
 --
 -- @return      the reverse iterator function
 --
+-- Stateful closure; see `list:items()` for the rationale.
 function list:ritems()
-    local iter = function (list, item)
-        return list:prev(item)
+    local item = nil
+    return function ()
+        item = self:prev(item)
+        return item
     end
-    return iter, self, nil
 end
 
 -- create a new doubly-linked list
diff --git a/xmake/core/sandbox/modules/ipairs.lua b/xmake/core/sandbox/modules/ipairs.lua
index 19591fce5..d910eb2b7 100644
--- a/xmake/core/sandbox/modules/ipairs.lua
+++ b/xmake/core/sandbox/modules/ipairs.lua
@@ -22,6 +22,10 @@
 local table = require("base/table")
 
 -- improve ipairs, wrap nil and single value
+--
+-- Like sandbox `pairs`, this is a stateful closure so the loop body can
+-- safely reassign the first loop variable under lua 5.4+ (paired with
+-- the RDKCONST->VDKREG compile-time patch in core/src/lua/xmake.lua).
 function sandbox_ipairs(t)
 
     -- exists the custom ipairs?
@@ -34,15 +38,21 @@ function sandbox_ipairs(t)
     if not is_table then
         t = t ~= nil and {t} or {}
     end
-    return function (t, i)
+    -- keep the index in an upvalue so the loop body can safely reassign
+    -- the first loop variable. In lua 5.4+ the for-in control slot is
+    -- merged with the first user variable; writes to it would otherwise
+    -- silently skip or repeat entries on the next iteration.
+    local i = 0
+    return function ()
         i = i + 1
         local v = t[i]
         if v ~= nil then
             return i, v
         end
-    end, t, 0
+    end
 end
 
+
 -- load module
 return sandbox_ipairs
 
diff --git a/xmake/core/sandbox/modules/irpairs.lua b/xmake/core/sandbox/modules/irpairs.lua
index 486f89506..c2224b2e1 100644
--- a/xmake/core/sandbox/modules/irpairs.lua
+++ b/xmake/core/sandbox/modules/irpairs.lua
@@ -40,14 +40,22 @@ local table = require("base/table")
 -- end
 --
 -- @endcode
+--
+-- Implemented as a stateful closure so the loop body can safely
+-- reassign the first loop variable under lua 5.4+ (paired with the
+-- RDKCONST->VDKREG compile-time patch in core/src/lua/xmake.lua).
 function sandbox_irpairs(t, filter, ...)
 
     -- has filter?
     local has_filter = type(filter) == "function"
 
-    -- init iterator
+    -- stateful closure: keep the index in an upvalue so the loop body
+    -- can safely reassign the first loop variable (lua 5.4+ merges the
+    -- for-in control slot with the first user variable).
     local args = table.pack(...)
-    local iter = function (t, i)
+    t = table.wrap(t)
+    local i = table.getn(t) + 1
+    return function ()
         i = i - 1
         local v = t[i]
         if v ~= nil then
@@ -57,10 +65,6 @@ function sandbox_irpairs(t, filter, ...)
             return i, v
         end
     end
-
-    -- return iterator and initialized state
-    t = table.wrap(t)
-    return iter, t, table.getn(t) + 1
 end
 
 -- load module
diff --git a/xmake/core/sandbox/modules/pairs.lua b/xmake/core/sandbox/modules/pairs.lua
index 50fa4d10e..dc162bc47 100644
--- a/xmake/core/sandbox/modules/pairs.lua
+++ b/xmake/core/sandbox/modules/pairs.lua
@@ -22,6 +22,19 @@
 local table = require("base/table")
 
 -- improve pairs, wrap nil/single value
+--
+-- Unlike the stock lua `pairs`, this sandbox version tolerates the loop
+-- body reassigning the first loop variable, e.g.:
+--
+--   for k, v in pairs(t) do
+--       k = k:gsub("_", "-")   -- safe here
+--       ...
+--   end
+--
+-- This works together with the `RDKCONST -> VDKREG` compile-time patch
+-- applied to lparser.c in core/src/lua/xmake.lua (and xmake.sh): that
+-- patch lifts lua 5.4+'s ban on writing to the for-in control variable,
+-- and this stateful closure makes such writes harmless at runtime.
 function sandbox_pairs(t)
 
     -- exists the custom ipairs?
@@ -34,9 +47,16 @@ function sandbox_pairs(t)
     if not is_table then
         t = t ~= nil and {t} or {}
     end
-    return function (t, i)
-        return next(t, i)
-    end, t, nil
+    -- keep `next`'s key in an upvalue so the loop body can safely reassign
+    -- the first loop variable. In lua 5.4+ the for-in control slot is
+    -- merged with the first user variable; if we threaded the key through
+    -- the loop it would corrupt `next` on the following iteration.
+    local k = nil
+    return function ()
+        local nk, nv = next(t, k)
+        k = nk
+        return nk, nv
+    end
 end
 
 -- load module