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: ---
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 } }
---
---
+--- 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:
+---
+--- -- 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 }
+---
+---
---@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)