diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt index 993480283e..a0daea9bc7 100644 --- a/runtime/doc/lua.txt +++ b/runtime/doc/lua.txt @@ -3185,7 +3185,10 @@ Iter:nth({n}) *Iter:nth()* (`any`) Iter:peek() *Iter:peek()* - Gets the next value in a |list-iterator| without consuming it. + Gets the next value from the iterator without consuming it. + + The value returned by |Iter:peek()| will be returned again by the next + call to |Iter:next()|. Example: >lua @@ -3291,8 +3294,12 @@ Iter:rskip({n}) *Iter:rskip()* (`Iter`) Iter:skip({n}) *Iter:skip()* - Skips `n` values of an iterator pipeline, or all values satisfying a - predicate of a |list-iterator|. + Skips `n` values of an iterator pipeline, or skips values while a + predicate returns |lua-truthy|. + + When a predicate is used, skipping stops at the first value for which the + predicate returns non-truthy. That value is not consumed and will be + returned by the next call to |Iter:next()| Example: >lua diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt index 39f2eaa3a4..3ae225dc66 100644 --- a/runtime/doc/news.txt +++ b/runtime/doc/news.txt @@ -317,6 +317,7 @@ LUA • |vim.version.range()| output can be converted to a human-readable string with |tostring()|. • |vim.version.intersect()| computes intersection of two version ranges. • |Iter:take()| and |Iter:skip()| now optionally accept predicates. +• |Iter:peek()| now works for all iterator types, not just |list-iterator|. • Built-in plugin manager |vim.pack| • |vim.list.unique()| and |Iter:unique()| to deduplicate lists and iterators, respectively. diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua index 758ef9b5dc..654ab9bf12 100644 --- a/runtime/lua/vim/iter.lua +++ b/runtime/lua/vim/iter.lua @@ -75,6 +75,8 @@ local M = {} ---@nodoc ---@class Iter +---@field _peeked any +---@field _next fun():... The underlying function that returns the next value(s) from the source. local Iter = {} Iter.__index = Iter Iter.__call = function(self) @@ -572,8 +574,14 @@ end --- ---@return any function Iter:next() - -- This function is provided by the source iterator in Iter.new. This definition exists only for - -- the docstring + if self._peeked then + local v = self._peeked + self._peeked = nil + + return unpack(v) + end + + return self._next() end ---@private @@ -610,7 +618,10 @@ function ArrayIter:rev() return self end ---- Gets the next value in a |list-iterator| without consuming it. +--- Gets the next value from the iterator without consuming it. +--- +--- The value returned by |Iter:peek()| will be returned again by the next call +--- to |Iter:next()|. --- --- Example: --- @@ -628,7 +639,11 @@ end --- ---@return any function Iter:peek() - error('peek() requires an array-like table') + if not self._peeked then + self._peeked = pack(self:next()) + end + + return unpack(self._peeked) end ---@private @@ -856,8 +871,11 @@ function ArrayIter:rpeek() end end ---- Skips `n` values of an iterator pipeline, or all values satisfying a ---- predicate of a |list-iterator|. +--- Skips `n` values of an iterator pipeline, or skips values while a predicate returns |lua-truthy|. +--- +--- When a predicate is used, skipping stops at the first value for which the +--- predicate returns non-truthy. That value is not consumed and will be returned +--- by the next call to |Iter:next()| --- --- Example: --- @@ -876,13 +894,30 @@ end ---@param n integer|fun(...):boolean Number of values to skip or a predicate. ---@return Iter function Iter:skip(n) - if type(n) == 'function' then - -- We would need to evaluate the perdicate without advancing iterator - error('skip() with predicate requires an array-like table') - end + if type(n) == 'number' then + for _ = 1, n do + self._peeked = nil + local _ = self:next() + end + elseif type(n) == 'function' then + local next = self.next - for _ = 1, n do - local _ = self:next() + self.next = function() + while true do + local peeked = self._peeked or pack(next(self)) + + if not peeked then + return nil + end + + if not n(unpack(peeked)) then + self._peeked = nil + return unpack(peeked) + end + + self._peeked = nil + end + end end return self end @@ -890,11 +925,13 @@ end ---@private function ArrayIter:skip(n) if type(n) == 'function' then - local inc = self._head < self._tail and 1 or -1 - local i = self._head - while n(unpack(self:peek())) and i ~= self._tail do - self:next() - i = i + inc + while self._head ~= self._tail do + local v = self._table[self._head] + if not n(unpack(v)) then + break + end + + self._head = self._head + (self._head < self._tail and 1 or -1) end return self end @@ -1128,7 +1165,7 @@ function Iter.new(src, ...) local mt = getmetatable(src) if mt and type(mt.__call) == 'function' then ---@private - function it.next() + it._next = function() return src() end @@ -1162,7 +1199,7 @@ function Iter.new(src, ...) end ---@private - function it.next() + it._next = function() return fn(src(s, var)) end diff --git a/test/functional/lua/iter_spec.lua b/test/functional/lua/iter_spec.lua index 76495c4d5f..59bf4be38f 100644 --- a/test/functional/lua/iter_spec.lua +++ b/test/functional/lua/iter_spec.lua @@ -196,6 +196,26 @@ describe('vim.iter', function() end end) + it('skip(predicate) preserves first non-matching element', function() + local it = vim.iter(vim.gsplit('1|2|3|4', '|')) + + it:skip(function(x) + return tonumber(x) < 3 + end) + + eq('3', it:next()) + eq('4', it:next()) + end) + + it('skip() followed by peek() works correctly', function() + local it = vim.iter(vim.gsplit('a|b|c|d', '|')) + + it:skip(2) + + eq('c', it:peek()) + eq('c', it:next()) + end) + it('rskip()', function() do local q = { 4, 3, 2, 1 } @@ -434,10 +454,88 @@ describe('vim.iter', function() do local it = vim.iter(vim.gsplit('hi', '')) - matches('peek%(%) requires an array%-like table', pcall_err(it.peek, it)) + eq('h', it:peek()) + eq('h', it:peek()) + eq('h', it:next()) + eq('i', it:peek()) + eq('i', it:next()) end end) + it('peek() does not consume on function iterators', function() + local it = vim.iter(vim.gsplit('a|b|c', '|')) + + eq('a', it:peek()) + eq('a', it:peek()) + eq('a', it:next()) + eq('b', it:next()) + end) + + it('peek() before skip(predicate) does not break iteration', function() + local it = vim.iter(vim.gsplit('1|2|3|4', '|')) + + eq('1', it:peek()) + + it:skip(function(x) + return tonumber(x) < 3 + end) + + eq('3', it:next()) + end) + + it('multiple peek() calls after next()', function() + local it = vim.iter(vim.gsplit('a|b|c', '|')) + + eq('a', it:next()) + eq('b', it:peek()) + eq('b', it:peek()) + eq('b', it:next()) + eq('c', it:next()) + end) + + describe('peek() with multi-value returns', function() + it('peek() preserves multiple return values from ipairs()', function() + local it = vim.iter(ipairs({ 'a', 'b', 'c' })) + local i1, v1 = it:peek() + + eq(1, i1) + eq('a', v1) + + local i2, v2 = it:next() + + eq(1, i2) + eq('a', v2) + end) + + it('peek() works with pairs() returning multiple values', function() + local tbl = { x = 10, y = 20 } + local it = vim.iter(pairs(tbl)) + local k1, v1 = it:peek() + local k2, v2 = it:peek() + + eq(k1, k2) + eq(v1, v2) + end) + end) + + describe('peek() after transformations', function() + it('peek() works after map() on function iterator', function() + local it = vim.iter(vim.gsplit('1|2|3', '|')):map(tonumber) + + eq(1, it:peek()) + eq(1, it:next()) + eq(2, it:peek()) + end) + + it('peek() at end of iterator returns nil', function() + local it = vim.iter({ 1 }) + + eq(1, it:next()) + eq(nil, it:peek()) + eq(nil, it:next()) + end) + end) + it('find()', function() local q = { 3, 6, 9, 12 } eq(12, vim.iter(q):find(12))