refactor(iter): move helper functions under vim.iter

vim.iter is now both a function and a module (similar to vim.version).

1 files changed, 109 insertions, 3 deletions

diff --git a/runtime/lua/vim/iter.lua b/runtime/lua/vim/iter.lua
index 2545853b41..c2e2c5bd9f 100644
--- a/runtime/lua/vim/iter.lua
+++ b/runtime/lua/vim/iter.lua
@@ -1,4 +1,56 @@
---- Iterator implementation.
+---@defgroup lua-iter
+---
+--- The \*vim.iter\* module provides a generic "iterator" interface over tables and iterator
+--- functions.
+---
+--- \*vim.iter()\* wraps its table or function argument into an \*Iter\* object with methods (such
+--- as |Iter:filter()| and |Iter:map()|) that transform the underlying source data. These methods
+--- can be chained together to create iterator "pipelines". 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>
+---
+--- In addition to the |vim.iter()| function, the |vim.iter| module provides convenience functions
+--- like |vim.iter.filter()| and |vim.iter.totable()|.
+
+local M = {}
 
 ---@class Iter
 local Iter = {}
@@ -733,7 +785,6 @@ end
 --- </pre>
 ---
 --- over
----
 --- <pre>lua
 --- vim.iter(t):enumerate()
 --- </pre>
@@ -776,6 +827,7 @@ end
 ---
 ---@param src table|function Table or iterator to drain values from
 ---@return Iter
+---@private
 function Iter.new(src, ...)
   local it = {}
   if type(src) == 'table' then
@@ -807,6 +859,7 @@ function Iter.new(src, ...)
       end
     end
 
+    ---@private
     function it.next()
       return fn(src(s, var))
     end
@@ -832,4 +885,57 @@ function ListIter.new(t)
   return it
 end
 
-return Iter
+--- 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 M.totable(f, ...)
+  return Iter.new(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 M.filter(f, src, ...)
+  return Iter.new(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 M.map(f, src, ...)
+  return Iter.new(src, ...):map(f):totable()
+end
+
+return setmetatable(M, {
+  __call = function(_, ...)
+    return Iter.new(...)
+  end,
+})