path: root/runtime/lua/vim/iter.lua

--- @brief
---
--- [vim.iter()]() is an interface for [iterable]s: it wraps a 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 to create iterator "pipelines": the output
--- of each pipeline stage is input to the next stage. The first stage depends on the type passed to
--- `vim.iter()`:
---
--- - Lists or arrays (|lua-list|) yield only the value of each element.
---   - Holes (nil values) are allowed (but discarded).
---   - Use pairs() to treat array/list tables as dicts (preserve holes and non-contiguous integer
---     keys): `vim.iter(pairs(…))`.
---   - Use |Iter:enumerate()| to also pass the index to the next stage.
---     - Or initialize with ipairs(): `vim.iter(ipairs(…))`.
--- - Non-list tables (|lua-dict|) yield both the key and value of each element.
--- - Function |iterator|s yield all values returned by the underlying function.
--- - Tables with a |__call()| metamethod are treated as function iterators.
---
--- The iterator pipeline terminates when the underlying |iterable| is exhausted (for function
--- iterators this means it returned nil).
---
--- Note: `vim.iter()` scans table input to decide if it is a list or a dict; to avoid this cost you
--- can wrap the table with an iterator e.g. `vim.iter(ipairs({…}))`, but that precludes the use of
--- |list-iterator| operations such as |IterArray:rev()|).
---
--- Examples:
---
--- ```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 }
---
--- -- ipairs() is a function iterator which returns both the index (i) and the value (v)
--- 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
---
--- local rb = vim.ringbuf(3)
--- rb:push("a")
--- rb:push("b")
--- vim.iter(rb):totable()
--- -- { "a", "b" }
--- ```

-- `never` represents an empty tail for single-value iterators.
--- @nodoc
--- @class vim.IterModule
--- @operator call: vim.Iter<any, any...>
--- @overload fun<T>(src: T[]): vim.IterArray<T, never>
--- @overload fun<K, V>(src: table<K, V>): vim.Iter<K, V>
--- @overload fun(src: table, ...): vim.Iter<any, any...>
--- @overload fun<S, C, R1, R...>(src: (fun(state: S, control: C): R1, R...), state?: S, control?: C): vim.Iter<R1, R...>
local M = {}

--- @nodoc
--- @class vim.Iter<V1, V...>
--- @field private _peeked V1|[V1, V...]? Cached peek value, if any. Nil if no value is peeked.
--- @field private _next fun(): V1, V... The underlying function that returns the next packed value from the source.
--- @overload fun(self: vim.Iter<V1, V...>): V1?, V...
local Iter = {}
Iter.__index = Iter

function Iter:__call()
  return self:next()
end

-- Special case implementations for iterators on list tables.
--- @nodoc
--- @class vim.IterArray<V1, V...> : vim.Iter<V1, V...>
--- @field private _table (V1|[V1, V...])[] Underlying table data. Items may be packed tuples after transforms.
--- @field private _head integer Index to the front of a table iterator
--- @field private _tail integer Index to the end of a table iterator (exclusive)
local IterArray = setmetatable({}, Iter)
IterArray.__index = IterArray
IterArray.__call = Iter.__call

--- Packed tables use this as their metatable
local packedmt = {}

--- @generic V1, V...
--- @param t V1|([V1, V...] & { n: integer })
--- @return V1?, V...
local function unpack(t)
  if type(t) == 'table' and getmetatable(t) == packedmt then
    return _G.unpack(t, 1, t.n)
  end
  return t
end

--- @generic T...
--- @param ... T... Values to pack
--- @return T|[T...]
local function pack(...)
  local n = select('#', ...)
  if n > 1 then
    return setmetatable({ n = n, ... }, packedmt)
  end
  return ...
end

--- @generic T
--- @param t T
local function sanitize(t)
  if type(t) == 'table' and getmetatable(t) == packedmt then
    --- @cast t table
    -- Remove length tag and metatable
    t.n = nil
    setmetatable(t, nil)
  end
  return t
end

--- Flattens a single array-like table. Errors if it attempts to flatten a
--- dict-like table
--- @param t any value which should be flattened
--- @param max_depth integer depth to which the table should be flattened
--- @param depth integer current iteration depth
--- @param result table output table that contains flattened result
--- @return table? flattened table if it can be flattened, otherwise nil
local function flatten(t, max_depth, depth, result)
  if depth < max_depth and type(t) == 'table' then
    for k, v in pairs(t) do
      if type(k) ~= 'number' or k <= 0 or math.floor(k) ~= k then
        -- short-circuit: this is not a list like table
        return nil
      end

      if flatten(v, max_depth, depth + 1, result) == nil then
        return nil
      end
    end
  elseif t ~= nil then
    result[#result + 1] = t
  end

  return result
end

--- Determine if the current iterator stage should continue.
---
--- If any arguments are passed to this function, then return those arguments
--- and stop the current iterator stage. Otherwise, return true to signal that
--- the current stage should continue.
---
--- @param ... any Function arguments.
--- @return boolean True if the iterator stage should continue, false otherwise
--- @return any Function arguments.
local function continue(...)
  if select(1, ...) ~= nil then
    return false, ...
  end
  return true
end

--- If no input arguments are given return false, indicating the current
--- iterator stage should stop. Otherwise, apply the arguments to the function
--- f. If that function returns no values, the current iterator stage continues.
--- Otherwise, those values are returned.
---
--- @generic T, R
--- @param f fun(...:T...): R... Function to call with the given arguments
--- @param ... T... Arguments to apply to f
--- @return boolean : True if the iterator pipeline should continue, false otherwise
--- @return R... : Return values of f