hoge

Author: delphinus

lua/plenary/iterators.lua

@@ -16,13 +16,14 @@ local f = require "plenary.functional"
 local exports = {}
 
 ---@generic V
----@alias PlenaryIteratorsIterator fun(table: V[], i?: integer): integer, V?
+---@alias PlenaryIteratorsIterator fun(param: string|V[], i?: integer): integer?, string|V|nil
 
 ---@class PlenaryIterator
 ---@field gen PlenaryIteratorsIterator
----@field param table
+---@field param string|table
 ---@field state? integer
----@overload fun(param?: table, state?: integer): integer, any?
+---@overload fun(param?: string|table, state?: integer): integer?, any?
+
 local Iterator = {}
 Iterator.__index = Iterator
 
@@ -35,6 +36,11 @@ Iterator.__index = Iterator
 ---So instead we do not return param and state as multivals when doing wrap
 ---This causes the first loop iteration to call param and state with nil because we didn't return them as multivals
 ---We have to use or to check for nil and default to interal starting state and param
+---@generic T
+---@param param? string|T[]
+---@param state? integer
+---@return integer i
+---@return string|T v
 function Iterator:__call(param, state)
   return self.gen(param or self.param, state or self.state)
 end
@@ -46,6 +52,8 @@ end
 
 -- A special hack for zip/chain to skip last two state, if a wrapped iterator
 -- has been passed
+---@param ... any
+---@return integer
 local numargs = function(...)
   local n = select("#", ...)
   if n >= 3 then
@@ -63,13 +71,21 @@ local numargs = function(...)
   return n
 end
 
+---@param state_x? integer
+---@param ... any
+---@return ...
 local return_if_not_empty = function(state_x, ...)
   if state_x == nil then
     return nil
   end
   return ...
 end
 
+---@param fun function
+---@param state_x? integer
+---@param ... any
+---@return integer?
+---@return ...
 local call_if_not_empty = function(fun, state_x, ...)
   if state_x == nil then
     return nil
@@ -80,12 +96,21 @@ end
 --------------------------------------------------------------------------------
 -- Basic Functions
 --------------------------------------------------------------------------------
+---@param _param any
+---@param _state any
+---@return nil
 local nil_gen = function(_param, _state)
   return nil
 end
 
 local pairs_gen = pairs {}
 
+---@generic K, V
+---@param map table<K, V>
+---@param key K
+---@return K key
+---@return K key
+---@return V value
 local map_gen = function(map, key)
   local value
   key, value = pairs_gen(map, key)
@@ -94,8 +119,8 @@ end
 
 ---@param param string
 ---@param state integer
----@return integer? state
----@return string? r
+---@return integer state
+---@return string r
 local string_gen = function(param, state)
   state = state + 1
   if state > #param then
@@ -105,18 +130,12 @@ local string_gen = function(param, state)
   return state, r
 end
 
----@generic T: table, U: table
----@alias PlenaryIteratorsRawiterTable fun(obj: T|PlenaryIterator, param?: string, state?: integer): PlenaryIteratorsIterator, T|U|nil, integer?
-
----@generic T: table, U: table
----@param obj T|PlenaryIterator
----@param param? string
+---@param obj string|function|table|PlenaryIterator
+---@param param? string|table
 ---@param state? integer
 ---@return PlenaryIteratorsIterator gen
----@return T|U|nil param
+---@return string|table|nil param
 ---@return integer? state
----@overload fun(obj: PlenaryIteratorsIterator, param: any, state?: integer): PlenaryIteratorsIterator, any, integer?
----@overload fun(obj: string): PlenaryIteratorsIterator, string?, integer?
 local rawiter = function(obj, param, state)
   assert(obj ~= nil, "invalid iterator")
 
@@ -151,9 +170,9 @@ end
 ---Wraps the iterator triplet into a table to allow metamethods and calling with method form
 ---Important! We do not return param and state as multivals like the original luafun
 ---See the __call metamethod for more information
----@param gen any
----@param param any
----@param state any
+---@param gen PlenaryIteratorsIterator
+---@param param? string|table
+---@param state? integer
 ---@return PlenaryIterator
 local function wrap(gen, param, state)
   return setmetatable({
@@ -165,17 +184,17 @@ end
 
 ---Unwrap an iterator metatable into the iterator triplet
 ---@param self PlenaryIterator
----@return any
----@return any
----@return any
+---@return PlenaryIteratorsIterator gen
+---@return string|table param
+---@return integer? state
 local unwrap = function(self)
   return self.gen, self.param, self.state
 end
 
 ---Create an iterator from an object
----@param obj any
----@param param any (optional)
----@param state any (optional)
+---@param obj string|function|table|PlenaryIterator
+---@param param? table
+---@param state? integer
 ---@return PlenaryIterator
 local iter = function(obj, param, state)
   return wrap(rawiter(obj, param, state))
@@ -185,13 +204,15 @@ exports.iter = iter
 exports.wrap = wrap
 exports.unwrap = unwrap
 
+---@param fn PlenaryIteratorsIterator
 function Iterator:for_each(fn)
   local param, state = self.param, self.state
   repeat
     state = call_if_not_empty(fn, self.gen(param, state))
   until state == nil
 end
 
+---@return PlenaryIterator
 function Iterator:stateful()
   return wrap(
     co.wrap(function()
@@ -228,6 +249,10 @@ end
 --------------------------------------------------------------------------------
 -- Generators
 --------------------------------------------------------------------------------
+---@param param { [1]: integer, [2]: integer }
+---@param state? integer
+---@return integer? state
+---@return integer? state
 local range_gen = function(param, state)
   local stop, step = param[1], param[2]
   state = state + step
@@ -237,6 +262,10 @@ local range_gen = function(param, state)
   return state, state
 end
 
+---@param param { [1]: integer, [2]: integer }
+---@param state? integer
+---@return integer? state
+---@return integer? state
 local range_rev_gen = function(param, state)
   local stop, step = param[1], param[2]
   state = state + step
@@ -247,9 +276,9 @@ local range_rev_gen = function(param, state)
 end
 
 ---Creates a range iterator
----@param start number
----@param stop number
----@param step number
+---@param start integer
+---@param stop? integer
+---@param step? integer
 ---@return PlenaryIterator
 local range = function(start, stop, step)
   if step == nil then
@@ -276,21 +305,35 @@ local range = function(start, stop, step)
 end
 exports.range = range
 
+---@generic T
+---@param param_x T[]
+---@param state_x integer
+---@return integer state
+---@return T ...
 local duplicate_table_gen = function(param_x, state_x)
   return state_x + 1, unpack(param_x)
 end
 
+---@param param_x fun(state: integer): ...
+---@param state_x integer
+---@return integer state
+---@return any ...
 local duplicate_fun_gen = function(param_x, state_x)
   return state_x + 1, param_x(state_x)
 end
 
+---@generic T
+---@param param_x T
+---@param state_x integer
+---@return integer state
+---@return T param
 local duplicate_gen = function(param_x, state_x)
   return state_x + 1, param_x
 end
 
 ---Creates an infinite iterator that will yield the arguments
 ---If multiple arguments are passed, the args will be packed and unpacked
----@param ...: the arguments to duplicate
+---@param ... any the arguments to duplicate
 ---@return PlenaryIterator
 local duplicate = function(...)
   if select("#", ...) <= 1 then
@@ -303,7 +346,7 @@ exports.duplicate = duplicate
 
 ---Creates an iterator from a function
 ---NOTE: if the function is a closure and modifies state, the resulting iterator will not be stateless
----@param fun function
+---@param fun fun(state: integer): ...
 ---@return PlenaryIterator
 local from_fun = function(fun)
   assert(type(fun) == "function")
@@ -327,17 +370,25 @@ local ones = function()
 end
 exports.ones = ones
 
+---@param param_x { [1]: integer, [2]: integer }
+---@param _state_x any
+---@return 0
+---@return integer
 local rands_gen = function(param_x, _state_x)
   return 0, math.random(param_x[1], param_x[2])
 end
 
+---@param _param_x any
+---@param _state_x any
+---@return 0
+---@return float
 local rands_nil_gen = function(_param_x, _state_x)
   return 0, math.random()
 end
 
 ---Creates an infinite iterator that will yield random values.
----@param n number
----@param m number
+---@param n integer
+---@param m integer
 ---@return PlenaryIterator
 local rands = function(n, m)
   if n == nil and m == nil then
@@ -355,6 +406,10 @@ local rands = function(n, m)
 end
 exports.rands = rands
 
+---@param param { [1]: string, [2]: string }
+---@param state? integer
+---@return integer?
+---@return string?
 local split_gen = function(param, state)
   local input, sep = param[1], param[2]
   local input_len = #input
@@ -375,8 +430,8 @@ local split_gen = function(param, state)
 end
 
 ---Return an iterator of substrings separated by a string
----@param input string: the string to split
----@param sep string: the separator to find and split based on
+---@param input string the string to split
+---@param sep string the separator to find and split based on
 ---@return PlenaryIterator
 local split = function(input, sep)
   return wrap(split_gen, { input, sep }, 1)
@@ -385,13 +440,15 @@ exports.split = split
 
 ---Splits a string based on a single space
 ---An alias for split(input, " ")
----@param input any
----@return any
+---@param input string
+---@return PlenaryIterator
 local words = function(input)
   return split(input, " ")
 end
 exports.words = words
 
+---@param input string
+---@return PlenaryIterator
 local lines = function(input)
   -- TODO: platform specific linebreaks
   return split(input, "\n")
@@ -401,27 +458,45 @@ exports.lines = lines
 --------------------------------------------------------------------------------
 -- Transformations
 --------------------------------------------------------------------------------
+---@alias PlenaryIteratorsLoop fun(i?: integer, v?: any)
+
+---@param param { [1]: PlenaryIteratorsIterator, [2]: string|table, [3]: PlenaryIteratorsLoop }
+---@param state? integer
+---@return integer?
+---@return ...
 local map_gen2 = function(param, state)
   local gen_x, param_x, fun = param[1], param[2], param[3]
   return call_if_not_empty(fun, gen_x(param_x, state))
 end
 
 ---Iterator adapter that maps the previous iterator with a function
----@param fun function: The function to map with. Will be called on each element
+---@param fun PlenaryIteratorsLoop The function to map with. Will be called on each element
 ---@return PlenaryIterator
 function Iterator:map(fun)
   return wrap(map_gen2, { self.gen, self.param, fun }, self.state)
 end
 
+---@alias PlenaryIteratorsFlattenParam { [1]: PlenaryIteratorsIterator, [2]: string|table, [3]: integer }
+
 local flatten_gen1
 do
+  ---@param new_iter PlenaryIterator
+  ---@param state_x? integer
+  ---@param ...
+  ---@return PlenaryIteratorsFlattenParam?
+  ---@return ...
   local it = function(new_iter, state_x, ...)
     if state_x == nil then
       return nil
     end
     return { new_iter.gen, new_iter.param, state_x }, ...
   end
 
+  ---@param state PlenaryIteratorsFlattenParam
+  ---@param state_x? integer
+  ---@param ... unknown
+  ---@return PlenaryIteratorsFlattenParam?
+  ---@return ...
   flatten_gen1 = function(state, state_x, ...)
     if state_x == nil then
       return nil
@@ -432,7 +507,7 @@ do
     -- experimental part
     if getmetatable(first_arg) == Iterator then
       -- attach the iterator to the rest
-      local new_iter = (first_arg .. wrap(state[1], state[2], state_x)):flatten()
+      local new_iter = (first_arg .. wrap(state[1], state[2], state_x)):flatten() --[[@as PlenaryIterator]]
       -- advance the iterator by one
       return it(new_iter, new_iter.gen(new_iter.param, new_iter.state))
     end
@@ -441,6 +516,10 @@ do
   end
 end
 
+---@param _ any
+---@param state? PlenaryIteratorsFlattenParam
+---@return PlearyIteratorsFlattenParam?
+---@return ...
 local flatten_gen = function(_, state)
   if state == nil then
     return