Merge pull request #23198 from gpanders/iter-fix

fix(iter): allow table values in iterator pipelines

3 files changed, 97 insertions, 56 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index 60067f1bc3..8e68e9a792 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -3037,6 +3037,19 @@ Iter:find({self}, {f})                                           *Iter:find()*
 Iter:fold({self}, {init}, {f})                                   *Iter:fold()*
     Fold an iterator or table into a single value.
 
+    Examples: >
+
+     -- Create a new table with only even values
+     local t = { a = 1, b = 2, c = 3, d = 4 }
+     local it = vim.iter(t)
+     it:filter(function(k, v) return v % 2 == 0 end)
+     it:fold({}, function(t, k, v)
+       t[k] = v
+       return t
+     end)
+     -- { b = 2, d = 4 }
+<
+
     Parameters: ~
       • {init}  any Initial value of the accumulator.
       • {f}     function(acc:any, ...):A Accumulation function.
@@ -3297,9 +3310,7 @@ Iter:totable({self})                                          *Iter:totable()*
     The resulting table depends on the initial source in the iterator
     pipeline. List-like tables and function iterators will be collected into a
     list-like table. If multiple values are returned from the final stage in
-    the iterator pipeline, each value will be included in a table. If a
-    map-like table was used as the initial source, then a map-like table is
-    returned.
+    the iterator pipeline, each value will be included in a table.
 
     Examples: >lua
 
@@ -3310,9 +3321,13 @@ Iter:totable({self})                                          *Iter:totable()*
      -- { { 1, 2 }, { 2, 4 }, { 3, 6 } }
 
      vim.iter({ a = 1, b = 2, c = 3 }):filter(function(k, v) return v % 2 ~= 0 end):totable()
-     -- { a = 1, c = 3 }
+     -- { { 'a', 1 }, { 'c', 3 } }
 <
 
+    The generated table is a list-like table with consecutive, numeric
+    indices. To create a map-like table with arbitrary keys, use
+    |Iter:fold()|.
+
     Return: ~
         (table)
 
diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua
index b73c03ba9a..c5d5ef835b 100644
--- a/runtime/lua/vim/iter.lua
+++ b/runtime/lua/vim/iter.lua
@@ -18,26 +18,19 @@ ListIter.__call = function(self)
   return self:next()
 end
 
---- Special case implementations for iterators on non-list tables.
----@class TableIter : Iter
-local TableIter = {}
-TableIter.__index = setmetatable(TableIter, Iter)
-TableIter.__call = function(self)
-  return self:next()
-end
-
 ---@private
 local function unpack(t)
-  if type(t) == 'table' then
-    return _G.unpack(t)
+  if type(t) == 'table' and t.__n ~= nil then
+    return _G.unpack(t, 1, t.__n)
   end
   return t
 end
 
 ---@private
 local function pack(...)
-  if select('#', ...) > 1 then
-    return { ... }
+  local n = select('#', ...)
+  if n > 1 then
+    return { __n = n, ... }
   end
   return ...
 end
@@ -184,10 +177,10 @@ end
 
 --- Collect the iterator into a table.
 ---
---- The resulting table depends on the initial source in the iterator pipeline. List-like tables
---- and function iterators will be collected into a list-like table. If multiple values are returned
---- from the final stage in the iterator pipeline, each value will be included in a table. If a
---- map-like table was used as the initial source, then a map-like table is returned.
+--- The resulting table depends on the initial source in the iterator pipeline.
+--- List-like tables and function iterators will be collected into a list-like
+--- table. If multiple values are returned from the final stage in the iterator
+--- pipeline, each value will be included in a table.
 ---
 --- Examples:
 --- <pre>lua
@@ -198,9 +191,13 @@ end
 --- -- { { 1, 2 }, { 2, 4 }, { 3, 6 } }
 ---
 --- vim.iter({ a = 1, b = 2, c = 3 }):filter(function(k, v) return v % 2 ~= 0 end):totable()
---- -- { a = 1, c = 3 }
+--- -- { { 'a', 1 }, { 'c', 3 } }
 --- </pre>
 ---
+--- The generated table is a list-like table with consecutive, numeric indices.
+--- To create a map-like table with arbitrary keys, use |Iter:fold()|.
+---
+---
 ---@return table
 function Iter.totable(self)
   local t = {}
@@ -210,6 +207,12 @@ function Iter.totable(self)
     if args == nil then
       break
     end
+
+    if type(args) == 'table' then
+      -- Removed packed table tag if it exists
+      args.__n = nil
+    end
+
     t[#t + 1] = args
   end
   return t
@@ -218,23 +221,35 @@ end
 ---@private
 function ListIter.totable(self)
   if self._head == 1 and self._tail == #self._table + 1 and self.next == ListIter.next then
+    -- Remove any packed table tags
+    for i = 1, #self._table do
+      local v = self._table[i]
+      if type(v) == 'table' then
+        v.__n = nil
+        self._table[i] = v
+      end
+    end
     return self._table
   end
 
   return Iter.totable(self)
 end
 
----@private
-function TableIter.totable(self)
-  local t = {}
-  for k, v in self do
-    t[k] = v
-  end
-  return t
-end
-
 --- Fold an iterator or table into a single value.
 ---
+--- Examples:
+--- <pre>
+--- -- Create a new table with only even values
+--- local t = { a = 1, b = 2, c = 3, d = 4 }
+--- local it = vim.iter(t)
+--- it:filter(function(k, v) return v % 2 == 0 end)
+--- it:fold({}, function(t, k, v)
+---   t[k] = v
+---   return t
+--- end)
+--- -- { b = 2, d = 4 }
+--- </pre>
+---
 ---@generic A
 ---
 ---@param init A Initial value of the accumulator.
@@ -747,7 +762,7 @@ function ListIter.enumerate(self)
   local inc = self._head < self._tail and 1 or -1
   for i = self._head, self._tail - inc, inc do
     local v = self._table[i]
-    self._table[i] = { i, v }
+    self._table[i] = pack(i, v)
   end
   return self
 end
@@ -768,7 +783,7 @@ function Iter.new(src, ...)
       count = count + 1
       local v = src[count]
       if v == nil then
-        return TableIter.new(src)
+        return Iter.new(pairs(src))
       end
       t[count] = v
     end
@@ -812,25 +827,4 @@ function ListIter.new(t)
   return it
 end
 
---- Create a new TableIter
----
----@param t table Table to iterate over. For list-like tables, use ListIter.new instead.
----@return Iter
----@private
-function TableIter.new(t)
-  local it = {}
-
-  local index = nil
-  function it.next()
-    local k, v = next(t, index)
-    if k ~= nil then
-      index = k
-      return k, v
-    end
-  end
-
-  setmetatable(it, TableIter)
-  return it
-end
-
 return Iter
diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index e5caf6f6f7..e37d477376 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -3374,13 +3374,45 @@ describe('lua stdlib', function()
     end)
 
     it('handles map-like tables', function()
-      local t = { a = 1, b = 2, c = 3 }
-      local it = vim.iter(t):map(function(k, v)
+      local it = vim.iter({ a = 1, b = 2, c = 3 }):map(function(k, v)
         if v % 2 ~= 0 then
           return k:upper(), v * 2
         end
       end)
-      eq({ A = 2, C = 6 }, it:totable())
+
+      local t = it:fold({}, function(t, k, v)
+        t[k] = v
+        return t
+      end)
+      eq({ A = 2, C = 6 }, t)
+    end)
+
+    it('handles table values mid-pipeline', function()
+      local map = {
+        item = {
+          file = 'test',
+        },
+        item_2 = {
+          file = 'test',
+        },
+        item_3 = {
+          file = 'test',
+        },
+      }
+
+      local output = vim.iter(map):map(function(key, value)
+        return { [key] = value.file }
+      end):totable()
+
+      table.sort(output, function(a, b)
+        return next(a) < next(b)
+      end)
+
+      eq({
+        { item = 'test' },
+        { item_2 = 'test' },
+        { item_3 = 'test' },
+      }, output)
     end)
   end)
 end)