commit 8a0cbf04d6a60f91c69c707789b18986a6921f8f
parent 179e7fccd74b02da06664de7103ba4ccee40778d
Author: Elijah Koulaxis <90087463+kx0101@users.noreply.github.com>
Date: Thu, 12 Feb 2026 18:55:16 +0200
feat(iter): peek(), skip(predicate) for non-list iterators #37604
Problem:
Iter:peek() only works if the iterator is a |list-iterator| (internally, an `ArrayIter`).
However, it is possible to implement :peek() support for any iterator.
Solution:
- add `_peeked` buffer for lookahead without actually consuming values
- `peek()` now works for function, pairs(), and array iterators
- `skip(predicate)` stops at the first non matching element without consuming it
- keep existing optimized behavior for `ArrayIter` to maintain backward compatibility
- use `pack`/`unpack` to support iterators that return multiple values
Diffstat:
4 files changed, 166 insertions(+), 23 deletions(-)
diff --git 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
@@ -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
@@ -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
+
+ self.next = function()
+ while true do
+ local peeked = self._peeked or pack(next(self))
+
+ if not peeked then
+ return nil
+ end
- for _ = 1, n do
- local _ = self:next()
+ 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
@@ -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))
