fix drop_while work with stateful iterators

drewdzzz · Totktonada · commit 90967c98abaf · 2024-10-09T18:30:48.000+03:00
Currently, method `drop_while` rewinds the iterator once it yields values not satisfying the condition. Such approach doesn't work with stateful iterators since they cannot be rewound by simply copying the state. The commit brings a new implementation that works correctly with all types of iterators. Performance comparison. In the first scenario, functor from `drop_while` returns `true` on the first million elements, then it returns `false` on the other million elements. With LuaJIT enabled, the new implementation is about 7% faster. Without LuaJIT, results are identical. In the second scenario, we leave only the first half of elements so that functor always returns `true`, which results in an empty iterator since all the elements are dropped. In this case, the new implementation is about 10-15% faster with JIT enabled. In the third scenario, we leave the second half of elements so that the functor returns `false` right from the start. In this case, the new implementation is about 4% slower than the old one with JIT enabled. Part of #83
diff --git a/fun.lua b/fun.lua
@@ -411,24 +411,41 @@ end
 methods.drop_n = method1(drop_n)
 exports.drop_n = export1(drop_n)
 
+-- Unpack values from param[3] on the first iteration, then return
+-- values from the provided iterator.
+--
+-- A generator function for drop_while().
+local drop_while_gen = function(param, state)
+    local results = param[3]
+    if not results then
+        return param[1](param[2], state)
+    else
+        param[3] = nil
+        return state, unpack(results, 1, table.maxn(results))
+    end
+end
+
+-- Checks if drop_while should continue skipping. If iterator is not exhausted
+-- and skipping is over, elements returned by iterator are wrapped into a table
+-- and returned as the second return value. Note that a table is created only
+-- once, on the last iteration, for the sake of performance.
 local drop_while_x = function(fun, state_x, ...)
-    if state_x == nil or not fun(...) then
-        return state_x, false
+    if state_x ~= nil and not fun(...) then
+        return state_x, {...}
     end
-    return state_x, true, ...
+    return state_x
 end
 
 local drop_while = function(fun, gen_x, param_x, state_x)
     assert(type(fun) == "function", "invalid first argument to drop_while")
-    local cont, state_x_prev
-    repeat
-        state_x_prev = deepcopy(state_x)
-        state_x, cont = drop_while_x(fun, gen_x(param_x, state_x))
-    until not cont
+    local pivot = nil
+    while state_x ~= nil and pivot == nil do
+        state_x, pivot = drop_while_x(fun, gen_x(param_x, state_x))
+    end
     if state_x == nil then
         return wrap(nil_gen, nil, nil)
     end
-    return wrap(gen_x, param_x, state_x_prev)
+    return wrap(drop_while_gen, {gen_x, param_x, pivot}, state_x)
 end
 methods.drop_while = method1(drop_while)
 exports.drop_while = export1(drop_while)
diff --git a/tests/stateful.lua b/tests/stateful.lua
@@ -0,0 +1,71 @@
+-- compatibility with Lua 5.1/5.2
+local unpack = rawget(table, "unpack") or unpack
+
+function gen_stateful_iter(values)
+    local i = 1
+    local gen = function(_param, values)
+        if i > #values then
+            return nil
+        end
+        local t = values[i]
+        i = i + 1
+        if type(t) == 'table' then
+            return values, unpack(t, 1, table.maxn(t))
+        else
+            return values, t
+        end
+    end
+    return iter(gen, nil, values)
+end
+
+--------------------------------------------------------------------------------
+-- drop_while
+--------------------------------------------------------------------------------
+
+-- Simple test
+dump(gen_stateful_iter({1, 2, 3, 4, 5, 6, 5, 4, 3}):drop_while(function(x) return x < 5 end))
+--[[test
+5
+6
+5
+4
+3
+--test]]
+
+-- Multireturn
+dump(gen_stateful_iter({{1, 10}, {3, 30}, {5, 50}, {6, 60}, {3, 20}}):drop_while(function(x) return x < 5 end))
+--[[test
+5 50
+6 60
+3 20
+--test]]
+
+-- Multireturn with nil
+dump(gen_stateful_iter({{1, nil, 10}, {3, nil, 30}, {5, nil, 50}, {6, nil, 60}, {3, nil, 20}}):drop_while(function(x) return x < 5 end))
+--[[test
+5 nil 50
+6 nil 60
+3 nil 20
+--test]]
+
+-- Multireturn with condition on second returned value
+dump(gen_stateful_iter({{0, 1}, {0, 3}, {0, 5}, {0, 4}}):drop_while(function(x, y) return y < 5 end))
+--[[test
+0 5
+0 4
+--test]]
+
+-- Empty iterator
+dump(gen_stateful_iter({}):drop_while(function(x) return x < 5 end))
+--[[test
+--test]]
+
+-- Always false
+dump(gen_stateful_iter({1, 2, 3, 4, 5}):drop_while(function(x) return x > 100 end))
+--[[test
+1
+2
+3
+4
+5
+--test]]
