feat(vim.iter): add Iter:flatten (#26786)

Co-authored-by: Gregory Anders <greg@gpanders.com>
Co-authored-by: Jongwook Choi <wookayin@gmail.com>

runtime/doc/lua.txt

@@ -3431,6 +3431,28 @@ Iter:find({f})                                                   *Iter:find()*
     Return: ~
         any
 
+Iter:flatten({depth})                                         *Iter:flatten()*
+    Flattens a |list-iterator|, un-nesting nested values up to the given
+    {depth}. Errors if it attempts to flatten a dict-like value.
+
+    Examples: >lua
+        vim.iter({ 1, { 2 }, { { 3 } } }):flatten():totable()
+        -- { 1, 2, { 3 } }
+
+        vim.iter({1, { { a = 2 } }, { 3 } }):flatten():totable()
+        -- { 1, { a = 2 }, 3 }
+
+        vim.iter({ 1, { { a = 2 } }, { 3 } }):flatten(math.huge):totable()
+        -- error: attempt to flatten a dict-like table
+<
+
+    Parameters: ~
+      • {depth}  (number|nil) Depth to which |list-iterator| should be
+                 flattened (defaults to 1)
+
+    Return: ~
+        Iter
+
 Iter:fold({init}, {f})                                           *Iter:fold()*
     Folds ("reduces") an iterator into a single value.
 

runtime/lua/vim/iter.lua

@@ -112,6 +112,35 @@ local function sanitize(t)
   return t
 end
 
+--- Flattens a single list-like table. Errors if it attempts to flatten a
+--- dict-like table
+---@param v table table which should be flattened
+---@param max_depth number depth to which the table should be flattened
+---@param depth number current iteration depth
+---@param result table output table that contains flattened result
+---@return table|nil flattened table if it can be flattened, otherwise nil
+local function flatten(v, max_depth, depth, result)
+  if depth < max_depth and type(v) == 'table' then
+    local i = 0
+    for _ in pairs(v) do
+      i = i + 1
+
+      if v[i] == nil then
+        -- short-circuit: this is not a list like table
+        return nil
+      end
+
+      if flatten(v[i], max_depth, depth + 1, result) == nil then
+        return nil
+      end
+    end
+  else
+    result[#result + 1] = v
+  end
+
+  return result
+end
+
 --- Determine if the current iterator stage should continue.
 ---
 --- If any arguments are passed to this function, then return those arguments
@@ -179,6 +208,54 @@ function ListIter.filter(self, f)
   return self
 end
 
+--- Flattens a |list-iterator|, un-nesting nested values up to the given {depth}.
+--- Errors if it attempts to flatten a dict-like value.
+---
+--- Examples:
+---
+--- ```lua
+--- vim.iter({ 1, { 2 }, { { 3 } } }):flatten():totable()
+--- -- { 1, 2, { 3 } }
+---
+--- vim.iter({1, { { a = 2 } }, { 3 } }):flatten():totable()
+--- -- { 1, { a = 2 }, 3 }
+---
+--- vim.iter({ 1, { { a = 2 } }, { 3 } }):flatten(math.huge):totable()
+--- -- error: attempt to flatten a dict-like table
+--- ```
+---
+---@param depth? number Depth to which |list-iterator| should be flattened
+---                        (defaults to 1)
+---@return Iter
+function Iter.flatten(self, depth) -- luacheck: no unused args
+  error('flatten() requires a list-like table')
+end
+
+---@private
+function ListIter.flatten(self, depth)
+  depth = depth or 1
+  local inc = self._head < self._tail and 1 or -1
+  local target = {}
+
+  for i = self._head, self._tail - inc, inc do
+    local flattened = flatten(self._table[i], depth, 0, {})
+
+    -- exit early if we try to flatten a dict-like table
+    if flattened == nil then
+      error('flatten() requires a list-like table')
+    end
+
+    for _, v in pairs(flattened) do
+      target[#target + 1] = v
+    end
+  end
+
+  self._head = 1
+  self._tail = #target + 1
+  self._table = target
+  return self
+end
+
 --- Maps the items of an iterator pipeline to the values returned by `f`.
 ---
 --- If the map function returns nil, the value is filtered from the iterator.
@@ -461,9 +538,8 @@ end
 --- ```
 ---
 ---@return Iter
-function Iter.rev(self)
+function Iter.rev(self) -- luacheck: no unused args
   error('rev() requires a list-like table')
-  return self
 end
 
 ---@private
@@ -733,7 +809,6 @@ end
 ---@diagnostic disable-next-line: unused-local
 function Iter.skipback(self, n) -- luacheck: no unused args
   error('skipback() requires a list-like table')
-  return self
 end
 
 ---@private
@@ -800,7 +875,6 @@ end
 ---@diagnostic disable-next-line: unused-local
 function Iter.slice(self, first, last) -- luacheck: no unused args
   error('slice() requires a list-like table')
-  return self
 end
 
 ---@private

test/functional/lua/iter_spec.lua

@@ -462,6 +462,38 @@ describe('vim.iter', function()
     )
   end)
 
+  it('flatten()', function()
+    local t = { { 1, { 2 } }, { { { { 3 } } }, { 4 } }, { 5 } }
+
+    eq(t, vim.iter(t):flatten(-1):totable())
+    eq(t, vim.iter(t):flatten(0):totable())
+    eq({ 1, { 2 }, { { { 3 } } }, { 4 }, 5 }, vim.iter(t):flatten():totable())
+    eq({ 1, 2, { { 3 } }, 4, 5 }, vim.iter(t):flatten(2):totable())
+    eq({ 1, 2, { 3 }, 4, 5 }, vim.iter(t):flatten(3):totable())
+    eq({ 1, 2, 3, 4, 5 }, vim.iter(t):flatten(4):totable())
+
+    local m = { a = 1, b = { 2, 3 }, d = { 4 } }
+    local it = vim.iter(m)
+
+    local flat_err = 'flatten%(%) requires a list%-like table'
+    matches(flat_err, pcall_err(it.flatten, it))
+
+    -- cases from the documentation
+    local simple_example = { 1, { 2 }, { { 3 } } }
+    eq({ 1, 2, { 3 } }, vim.iter(simple_example):flatten():totable())
+
+    local not_list_like = vim.iter({ [2] = 2 })
+    matches(flat_err, pcall_err(not_list_like.flatten, not_list_like))
+
+    local also_not_list_like = vim.iter({ nil, 2 })
+    matches(flat_err, pcall_err(not_list_like.flatten, also_not_list_like))
+
+    local nested_non_lists = vim.iter({ 1, { { a = 2 } }, { { nil } }, { 3 } })
+    eq({ 1, { a = 2 }, { nil }, 3 }, nested_non_lists:flatten():totable())
+    -- only error if we're going deep enough to flatten a dict-like table
+    matches(flat_err, pcall_err(nested_non_lists.flatten, nested_non_lists, math.huge))
+  end)
+
   it('handles map-like tables', function()
     local it = vim.iter({ a = 1, b = 2, c = 3 }):map(function(k, v)
       if v % 2 ~= 0 then