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

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
 

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.

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
 

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