fix(iter): make pipeline termination conditions consistent (#24614)

If an iterator pipeline stage returns nil as its first return value, the
other return values are ignored and it is treated as if that stage
returned only nil (the semantics of returning nil are different between
different stages). This is consistent with how for loops work in Lua
more generally, where the for loop breaks when the first return value
from the function iterator is nil (see :h for-in for details).

3 files changed, 28 insertions, 42 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index fb92ad411e..08022c3ed6 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -3207,13 +3207,18 @@ receives as input the output values from the prior stage. The values used
 in the first stage of the pipeline depend on the type passed to this
 function:
 
-• List tables pass only the value of each element
-• Non-list (dict) tables pass both the key and value of each element
+• List tables (arrays) pass only the value of each element
+• Non-list tables (dictionaries) pass both the key and value of each
+  element
 • Function |iterator|s pass all of the values returned by their respective
   function
 • Tables with a metatable implementing |__call()| are treated as function
   iterators
 
+The iterator pipeline terminates when the original table or function
+iterator runs out of values (for function iterators, this means that the
+first value returned by the function is nil).
+
 Examples: >lua
 
    local it = vim.iter({ 1, 2, 3, 4, 5 })
@@ -3225,6 +3230,7 @@ Examples: >lua
    it:totable()
    -- { 9, 6, 3 }
 
+   -- ipairs() is a function iterator which returns both the index (i) and the value (v)
    vim.iter(ipairs({ 1, 2, 3, 4, 5 })):map(function(i, v)
      if i > 2 then return v end
    end):totable()
diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua
index a278e96a5f..56c130dd0c 100644
--- a/runtime/lua/vim/iter.lua
+++ b/runtime/lua/vim/iter.lua
@@ -6,11 +6,14 @@
 --- Each pipeline stage receives as input the output values from the prior stage. The values used in
 --- the first stage of the pipeline depend on the type passed to this function:
 ---
---- - List tables pass only the value of each element
---- - Non-list (dict) tables pass both the key and value of each element
+--- - List tables (arrays) pass only the value of each element
+--- - Non-list tables (dictionaries) pass both the key and value of each element
 --- - Function |iterator|s pass all of the values returned by their respective function
 --- - Tables with a metatable implementing |__call()| are treated as function iterators
 ---
+--- The iterator pipeline terminates when the original table or function iterator runs out of values
+--- (for function iterators, this means that the first value returned by the function is nil).
+---
 --- Examples:
 --- <pre>lua
 ---   local it = vim.iter({ 1, 2, 3, 4, 5 })
@@ -22,6 +25,7 @@
 ---   it:totable()
 ---   -- { 9, 6, 3 }
 ---
+---   -- ipairs() is a function iterator which returns both the index (i) and the value (v)
 ---   vim.iter(ipairs({ 1, 2, 3, 4, 5 })):map(function(i, v)
 ---     if i > 2 then return v end
 ---   end):totable()
@@ -111,7 +115,7 @@ end
 ---@return boolean True if the iterator stage should continue, false otherwise
 ---@return any Function arguments.
 local function continue(...)
-  if select('#', ...) > 0 then
+  if select(1, ...) ~= nil then
     return false, ...
   end
   return true
@@ -127,7 +131,7 @@ end
 ---@return boolean True if the iterator pipeline should continue, false otherwise
 ---@return any Return values of f
 local function apply(f, ...)
-  if select('#', ...) > 0 then
+  if select(1, ...) ~= nil then
     return continue(f(...))
   end
   return false
@@ -258,7 +262,7 @@ end
 ---                       in the pipeline as arguments.
 function Iter.each(self, f)
   local function fn(...)
-    if select('#', ...) > 0 then
+    if select(1, ...) ~= nil then
       f(...)
       return true
     end
@@ -740,6 +744,12 @@ end
 ---@param last number
 ---@return Iter
 function Iter.slice(self, first, last) -- luacheck: no unused args
+  error('slice() requires a list-like table')
+  return self
+end
+
+---@private
+function ListIter.slice(self, first, last)
   return self:skip(math.max(0, first - 1)):skipback(math.max(0, self._tail - last - 1))
 end
 
@@ -912,6 +922,8 @@ function Iter.new(src, ...)
 
     --- Use a closure to handle var args returned from iterator
     local function fn(...)
+      -- Per the Lua 5.1 reference manual, an iterator is complete when the first returned value is
+      -- nil (even if there are other, non-nil return values). See |for-in|.
       if select(1, ...) ~= nil then
         var = select(1, ...)
         return ...
diff --git a/test/functional/lua/iter_spec.lua b/test/functional/lua/iter_spec.lua
index 3b603c9911..ffa28e7b11 100644
--- a/test/functional/lua/iter_spec.lua
+++ b/test/functional/lua/iter_spec.lua
@@ -154,6 +154,9 @@ describe('vim.iter', function()
     eq({1, 2}, vim.iter(t):slice(1, 2):totable())
     eq({10}, vim.iter(t):slice(10, 10):totable())
     eq({8, 9, 10}, vim.iter(t):slice(8, 11):totable())
+
+    local it = vim.iter(vim.gsplit('a|b|c|d', '|'))
+    matches('slice%(%) requires a list%-like table', pcall_err(it.slice, it, 1, 3))
   end)
 
   it('nth()', function()
@@ -396,39 +399,4 @@ describe('vim.iter', function()
       { item_3 = 'test' },
     }, output)
   end)
-
-  it('handles nil values', function()
-    local t = {1, 2, 3, 4, 5}
-    do
-      local it = vim.iter(t):enumerate():map(function(i, v)
-        if i % 2 == 0 then
-          return nil, v*v
-        end
-        return v, nil
-      end)
-      eq({
-        { [1] = 1 },
-        { [2] = 4 },
-        { [1] = 3 },
-        { [2] = 16 },
-        { [1] = 5 },
-      }, it:totable())
-    end
-
-    do
-      local it = vim.iter(ipairs(t)):map(function(i, v)
-        if i % 2 == 0 then
-          return nil, v*v
-        end
-        return v, nil
-      end)
-      eq({
-        { [1] = 1 },
-        { [2] = 4 },
-        { [1] = 3 },
-        { [2] = 16 },
-        { [1] = 5 },
-      }, it:totable())
-    end
-  end)
 end)