feat(lua): add vim.iter (#23029)

vim.iter wraps a table or iterator function into an `Iter` object with
methods such as `filter`, `map`, and `fold` which can be chained to
produce iterator pipelines that do not create new tables at each step.

1 files changed, 105 insertions, 0 deletions

diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua
index f700f4a6b3..1a96ef9bc4 100644
--- a/runtime/lua/vim/shared.lua
+++ b/runtime/lua/vim/shared.lua
@@ -884,4 +884,109 @@ function vim.defaulttable(create)
   })
 end
 
+--- Create an Iter |lua-iter| object from a table or iterator.
+---
+--- The input value can be a table  or a function iterator (see |luaref-in|).
+---
+--- This function wraps the input value into an interface which allows chaining
+--- multiple pipeline stages in an efficient manner. 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 tables pass both the key and value of each element
+--- - Function iterators pass all of the values returned by their respective
+---   function
+---
+--- Examples:
+--- <pre>lua
+--- local it = vim.iter({ 1, 2, 3, 4, 5 })
+--- it:map(function(v)
+---   return v * 3
+--- end)
+--- it:rev()
+--- it:skip(2)
+--- it:totable()
+--- -- { 9, 6, 3 }
+---
+--- vim.iter(ipairs({ 1, 2, 3, 4, 5 })):map(function(i, v)
+---   if i > 2 then return v end
+--- end):totable()
+--- -- { 3, 4, 5 }
+---
+--- local it = vim.iter(vim.gsplit('1,2,3,4,5', ','))
+--- it:map(function(s) return tonumber(s) end)
+--- for i, d in it:enumerate() do
+---   print(string.format("Column %d is %d", i, d))
+--- end
+--- -- Column 1 is 1
+--- -- Column 2 is 2
+--- -- Column 3 is 3
+--- -- Column 4 is 4
+--- -- Column 5 is 5
+---
+--- vim.iter({ a = 1, b = 2, c = 3, z = 26 }):any(function(k, v)
+---   return k == 'z'
+--- end)
+--- -- true
+--- </pre>
+---
+---@see |lua-iter|
+---
+---@param src table|function Table or iterator.
+---@return Iter @|lua-iter|
+function vim.iter(src, ...)
+  local Iter = require('vim.iter')
+  return Iter.new(src, ...)
+end
+
+--- Collect an iterator into a table.
+---
+--- This is a convenience function that performs:
+--- <pre>lua
+--- vim.iter(f):totable()
+--- </pre>
+---
+---@param f function Iterator function
+---@return table
+function vim.totable(f, ...)
+  return vim.iter(f, ...):totable()
+end
+
+--- Filter a table or iterator.
+---
+--- This is a convenience function that performs:
+--- <pre>lua
+--- vim.iter(src):filter(f):totable()
+--- </pre>
+---
+---@see |Iter:filter()|
+---
+---@param f function(...):bool Filter function. Accepts the current iterator or table values as
+---                            arguments and returns true if those values should be kept in the
+---                            final table
+---@param src table|function Table or iterator function to filter
+---@return table
+function vim.filter(f, src, ...)
+  return vim.iter(src, ...):filter(f):totable()
+end
+
+--- Map and filter a table or iterator.
+---
+--- This is a convenience function that performs:
+--- <pre>lua
+--- vim.iter(src):map(f):totable()
+--- </pre>
+---
+---@see |Iter:map()|
+---
+---@param f function(...):?any Map function. Accepts the current iterator or table values as
+---                            arguments and returns one or more new values. Nil values are removed
+---                            from the final table.
+---@param src table|function Table or iterator function to filter
+---@return table
+function vim.map(f, src, ...)
+  return vim.iter(src, ...):map(f):totable()
+end
+
 return vim