Repository: nvim-lua/plenary.nvim
Pull request #123 (2021-04-04 15:48:19)

feat: more iterators
- more docs
- flatten and zip iterator
- fixed stateful iterator
- spit iterator


Patch:
diff --git a/lua/plenary/async_lib/async.lua b/lua/plenary/async_lib/async.lua
index 027d640..3fce4a5 100644
--- a/lua/plenary/async_lib/async.lua
+++ b/lua/plenary/async_lib/async.lua
@@ -2,6 +2,7 @@ local co = coroutine
 local errors = require('plenary.errors')
 local traceback_error = errors.traceback_error
 local f = require('plenary.functional')
+local tbl = require('plenary.tbl')
 
 local M = {}
 
@@ -55,16 +56,19 @@ M.wrap = function(func, argc)
   end
 
   return function(...)
-    local params = {...}
+    local params = tbl.pack(...)
 
     local function future(step)
       if step then
         if type(argc) == "number" then
           params[argc] = step
+          params.n = argc
         else
           table.insert(params, step) -- change once not optional
+          params.n = params.n + 1
         end
-        return func(unpack(params))
+
+        return func(tbl.unpack(params))
       else
         return co.yield(future)
       end
@@ -180,10 +184,10 @@ M.async = function(func)
   end
 
   return function(...)
-    local args = {...}
+    local args = tbl.pack(...)
     local function future(step)
       if step == nil then
-        return func(unpack(args))
+        return func(tbl.unpack(args))
       else
         execute(future, step)
       end
diff --git a/lua/plenary/iterators.lua b/lua/plenary/iterators.lua
index 78648fa..e2fbde0 100644
--- a/lua/plenary/iterators.lua
+++ b/lua/plenary/iterators.lua
@@ -1,16 +1,36 @@
--- Adaptation of luafun for neovim
+---@brief [[
+---An adaptation of luafun for neovim.
+---This library will use neovim specific functions.
+---Some documentation is the same as from luafun.
+---Some extra functions are present that are not in luafun
+---@brief ]]
+
 local co = coroutine
+local f = require('plenary.functional')
 
 --------------------------------------------------------------------------------
 -- Tools
 --------------------------------------------------------------------------------
+
 local exports = {}
 
+---@class Iterator
+---@field gen function
+---@field param any
+---@field state any
 local Iterator = {}
 Iterator.__index = Iterator
 
+---Makes a for loop work
+---If not called without param or state, will just generate with the starting state
+---This is useful because the original luafun will also return param and state in addition to the iterator as a multival
+---This can cause problems because when using iterators as expressions the multivals can bleed
+---For example i.iter { 1, 2, i.iter { 3, 4 } } will not work because the inner iterator returns a multival thus polluting the list with internal values
+---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
 function Iterator:__call(param, state)
-  return self.gen(param, state)
+  return self.gen(param or self.param, state or self.state)
 end
 
 function Iterator:__tostring()
@@ -102,18 +122,35 @@ local rawiter = function(obj, param, state)
   error(string.format('object %s of type "%s" is not iterable', obj, type(obj)))
 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
+---Se the __call metamethod for more information
+---@param gen any
+---@param param any
+---@param state any
+---@return Iterator
 local function wrap(gen, param, state)
   return setmetatable({
     gen = gen,
     param = param,
     state = state
-  }, Iterator), param, state
+  }, Iterator)
 end
 
+---Unwrap an iterator metatable into the iterator triplet
+---@param self Iterator
+---@return any
+---@return any
+---@return any
 local unwrap = function(self)
-    return self.gen, self.param, self.state
+  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)
+---@return Iterator
 local iter = function(obj, param, state)
   return wrap(rawiter(obj, param, state))
 end
@@ -123,20 +160,41 @@ exports.wrap = wrap
 exports.unwrap = unwrap
 
 function Iterator:for_each(fn)
-  local state = self.state
-  local param = self.param
+  local param, state = self.param, self.state
   repeat
     state = call_if_not_empty(fn, self.gen(param, state))
   until state == nil
 end
 
 function Iterator:stateful()
-  return co.wrap(function()
+  return wrap(co.wrap(function()
     self:for_each(function(...)
-      co.yield(...)
+      co.yield(f.first(...), ...)
     end)
-  end)
+
+    -- too make sure that we always return nil if there are no more
+    while true do
+      co.yield()
+    end
+  end), nil, nil)
 end
+
+-- function Iterator:stateful()
+--   local gen, param, state = self.gen, self.param, self.state
+
+--   local function return_and_set_state(state_x, ...)
+--     state = state_x
+--     if state == nil then return end
+--     return state_x, ...
+--   end
+
+--   local stateful_gen = function()
+--     return return_and_set_state(gen(param, state))
+--   end
+
+--   return wrap(stateful_gen, false, false)
+-- end
+
 --------------------------------------------------------------------------------
 -- Generators
 --------------------------------------------------------------------------------
@@ -158,6 +216,11 @@ local range_rev_gen = function(param, state)
   return state, state
 end
 
+---Creates a range iterator
+---@param start number
+---@param stop number
+---@param step number
+---@return Iterator
 local range = function(start, stop, step)
   if step == nil then
     if stop == nil then
@@ -181,19 +244,181 @@ local range = function(start, stop, step)
     return wrap(range_rev_gen, {stop, step}, start - step)
   end
 end
-
 exports.range = range
 
+local duplicate_table_gen = function(param_x, state_x)
+  return state_x + 1, unpack(param_x)
+end
+
+local duplicate_fun_gen = function(param_x, state_x)
+  return state_x + 1, param_x(state_x)
+end
+
+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
+---@return Iterator
+local duplicate = function(...)
+  if select('#', ...) <= 1 then
+    return wrap(duplicate_gen, select(1, ...), 0)
+  else
+    return wrap(duplicate_table_gen, {...}, 0)
+  end
+end
+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
+---@return Iterator
+local from_fun = function(fun)
+  assert(type(fun) == "function")
+  return wrap(duplicate_fun_gen, fun, 0)
+end
+exports.from_fun = from_fun
+
+---Creates an infinite iterator that will yield zeros.
+---This is an alias to calling duplicate(0)
+---@return Iterator
+local zeros = function()
+  return wrap(duplicate_gen, 0, 0)
+end
+exports.zeros = zeros
+
+---Creates an infinite iterator that will yield ones.
+---This is an alias to calling duplicate(1)
+---@return Iterator
+local ones = function()
+  return wrap(duplicate_gen, 1, 0)
+end
+exports.ones = ones
+
+local rands_gen = function(param_x, _state_x)
+  return 0, math.random(param_x[1], param_x[2])
+end
+
+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
+---@return Iterator
+local rands = function(n, m)
+  if n == nil and m == nil then
+    return wrap(rands_nil_gen, 0, 0)
+  end
+  assert(type(n) == "number", "invalid first arg to rands")
+  if m == nil then
+    m = n
+    n = 0
+  else
+    assert(type(m) == "number", "invalid second arg to rands")
+  end
+  assert(n < m, "empty interval")
+  return wrap(rands_gen, {n, m - 1}, 0)
+end
+exports.rands = rands
+
+local split_gen = function(param, state)
+  local input, sep = param[1], param[2]
+  local input_len = #input
+
+  if state > input_len + 1 then
+    return
+  end
+
+  local start, finish = string.find(input, sep, state, true)
+  if not start then
+    start = input_len + 1
+    finish = input_len + 1
+  end
+
+  local sub_str = input:sub(state, start - 1)
+
+  return finish + 1, sub_str
+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
+---@return Iterator
+local split = function(input, sep)
+  return wrap(split_gen, {input, sep}, 1)
+end
+exports.split = split
+
+---Splits a string based on a single space
+---An alias for split(input, " ")
+---@param input any
+---@return any
+local words = function(input)
+  return split(input, " ")
+end
+exports.words = words
+
+local lines = function(input)
+  -- TODO: platform specific linebreaks
+  return split(input, "\n")
+end
+exports.lines = lines
+
 --------------------------------------------------------------------------------
 -- Transformations
 --------------------------------------------------------------------------------
 local map_gen = 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))
+  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
+---@return Iterator
+function Iterator:map(fun)
+  return wrap(map_gen, {self.gen, self.param, fun}, self.state)
 end
 
-function Iterator:map(fn)
-  return wrap(map_gen, {self.gen, self.param, fn}, self.state)
+local flatten_gen1
+do
+  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
+
+  flatten_gen1 = function(state, state_x, ...)
+    if state_x == nil then
+      return nil
+    end
+
+    local first_arg = f.first(...)
+
+    -- 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()
+      -- advance the iterator by one
+      return it(new_iter, new_iter.gen(new_iter.param, new_iter.state))
+    end
+
+    return {state[1], state[2], state_x}, ...
+  end
+end
+
+local flatten_gen = function(_, state)
+  if state == nil then return end
+  local gen_x, param_x, state_x = state[1], state[2], state[3]
+  return flatten_gen1(state, gen_x(param_x, state_x))
+end
+
+---Iterator adapter that will recursivley flatten nested iterator structure
+---@return Iterator
+function Iterator:flatten()
+  return wrap(flatten_gen, false, {self.gen, self.param, self.state})
 end
 
 --------------------------------------------------------------------------------
@@ -240,45 +465,61 @@ local filter_gen = function(param, state_x)
     return filter_detect(fun, gen_x, param_x, gen_x(param_x, state_x))
 end
 
-function Iterator:filter(fn)
-  return wrap(filter_gen, {self.gen, self.param, fn}, self.state)
+---Iterator adapter that will filter values
+---@param fun function: The function to filter values with. If the function returns true, the value will be kept.
+---@return Iterator
+function Iterator:filter(fun)
+  return wrap(filter_gen, {self.gen, self.param, fun}, self.state)
 end
 
 --------------------------------------------------------------------------------
 -- Reducing
 --------------------------------------------------------------------------------
-function Iterator:any(fn)
+
+---Returns true if any of the values in the iterator satisfy a predicate
+---@param fun function
+---@return boolean
+function Iterator:any(fun)
   local r
   local state, param, gen = self.state, self.param, self.gen
   repeat
-    state, r = call_if_not_empty(fn, gen(param, state))
+    state, r = call_if_not_empty(fun, gen(param, state))
   until state == nil or r
   return r
 end
 
-function Iterator:all(fn)
+---Returns true if all of the values in the iterator satisfy a predicate
+---@param fun function
+---@return boolean
+function Iterator:all(fun)
   local r
   local state, param, gen = self.state, self.param, self.gen
   repeat
-    state, r = call_if_not_empty(fn, gen(param, state))
+    state, r = call_if_not_empty(fun, gen(param, state))
   until state == nil or not r
   return state == nil
 end
 
+---Finds a value that is equal to the provided value of satisfies a predicate.
+---@param val_or_fn any
+---@return any
 function Iterator:find(val_or_fn)
-  local fn
-  if type(val_or_fn) ~= "function" then
-    fn = function(...)
-      return ... == val_or_fn
-    end
+  local gen, param, state = self.gen, self.param, self.state
+  if type(val_or_fn) == "function" then
+    return return_if_not_empty(filter_detect(val_or_fn, gen, param, gen(param, state)))
   else
-    fn = val_or_fn
+    for _, r in gen, param, state do
+      if r == val_or_fn then
+        return r
+      end
+    end
+    return nil
   end
-
-  local gen, param, state = self.gen, self.param, self.state
-  return return_if_not_empty(filter_detect(fn, gen, param, gen(param, state)))
 end
 
+---Turns an iterator into a list.
+---If the iterator yields multivals only the first multival will be used.
+---@return table
 function Iterator:tolist()
   local list = {}
   self:for_each(function(a)
@@ -287,6 +528,9 @@ function Iterator:tolist()
   return list
 end
 
+---Turns an iterator into a list.
+---If the iterator yields multivals all multivals will be used and packed into a table.
+---@return table
 function Iterator:tolistn()
   local list = {}
   self:for_each(function(...)
@@ -295,6 +539,10 @@ function Iterator:tolistn()
   return list
 end
 
+---Turns an iterator into a map.
+---The first multival that the iterator yields will be the key.
+---The second multival that the iterator yields will be the value.
+---@return table
 function Iterator:tomap()
   local map = {}
   self:for_each(function(key, value)
@@ -310,12 +558,11 @@ end
 local chain_gen_r1
 local chain_gen_r2 = function(param, state, state_x, ...)
   if state_x == nil then
-    local i = state[1]
-    i = i + 1
+    local i = state[1] + 1
     if param[3 * i - 1] == nil then
       return nil
     end
-    local state_x = param[3 * i]
+    state_x = param[3 * i]
     return chain_gen_r1(param, {i, state_x})
   end
   return {state[1], state_x}, ...
@@ -324,16 +571,24 @@ end
 chain_gen_r1 = function(param, state)
   local i, state_x = state[1], state[2]
   local gen_x, param_x = param[3 * i - 2], param[3 * i - 1]
-  return chain_gen_r2(param, state, gen_x(param_x, state[2]))
+  return chain_gen_r2(param, state, gen_x(param_x, state_x))
 end
 
+---Make an iterator that returns elements from the first iterator until it is exhausted, then proceeds to the next iterator,
+---until all of the iterators are exhausted.
+---Used for treating consecutive iterators as a single iterator.
+---Infinity iterators are supported, but are not recommended.
+---@param ...: the iterators to chain
+---@return Iterator
 local chain = function(...)
   local n = numargs(...)
+
   if n == 0 then
     return wrap(nil_gen, nil, nil)
   end
 
   local param = { [3 * n] = 0 }
+
   local i, gen_x, param_x, state_x
   for i = 1, n, 1 do
     local elem = select(i, ...)
@@ -350,52 +605,52 @@ Iterator.chain = chain
 Iterator.__concat = chain
 exports.chain = chain
 
---------------------------------------------------------------------------------
--- Operators
---------------------------------------------------------------------------------
-exports.operator = {
-    ----------------------------------------------------------------------------
-    -- Comparison operators
-    ----------------------------------------------------------------------------
-    lt  = function(a, b) return a < b end,
-    le  = function(a, b) return a <= b end,
-    eq  = function(a, b) return a == b end,
-    ne  = function(a, b) return a ~= b end,
-    ge  = function(a, b) return a >= b end,
-    gt  = function(a, b) return a > b end,
-
-    ----------------------------------------------------------------------------
-    -- Arithmetic operators
-    ----------------------------------------------------------------------------
-    add = function(a, b) return a + b end,
-    div = function(a, b) return a / b end,
-    floordiv = function(a, b) return math.floor(a/b) end,
-    intdiv = function(a, b)
-        local q = a / b
-        if a >= 0 then return math.floor(q) else return math.ceil(q) end
-    end,
-    mod = function(a, b) return a % b end,
-    mul = function(a, b) return a * b end,
-    neq = function(a) return -a end,
-    unm = function(a) return -a end, -- an alias
-    pow = function(a, b) return a ^ b end,
-    sub = function(a, b) return a - b end,
-    truediv = function(a, b) return a / b end,
-
-    ----------------------------------------------------------------------------
-    -- String operators
-    ----------------------------------------------------------------------------
-    concat = function(a, b) return a..b end,
-    len = function(a) return #a end,
-    length = function(a) return #a end, -- an alias
-
-    ----------------------------------------------------------------------------
-    -- Logical operators
-    ----------------------------------------------------------------------------
-    land = function(a, b) return a and b end,
-    lor = function(a, b) return a or b end,
-    lnot = function(a) return not a end,
-    truth = function(a) return not not a end,
-}
+local function zip_gen_r(param, state, state_new, ...)
+  if #state_new == #param / 2 then
+    return state_new, ...
+  end
+
+  local i = #state_new + 1
+  local gen_x, param_x = param[2 * i - 1], param[2 * i]
+  local state_x, r = gen_x(param_x, state[i])
+  if state_x == nil then
+    return nil
+  end
+  table.insert(state_new, state_x)
+  return zip_gen_r(param, state, state_new, r, ...)
+end
+
+local zip_gen = function(param, state)
+  return zip_gen_r(param, state, {})
+end
+
+---Return a new iterator where i-th return value contains the i-th element from each of the iterators.
+---The returned iterator is truncated in length to the length of the shortest iterator.
+---For multi-return iterators only the first variable is used.
+---@param ...: the iterators to zip
+---@return Iterator
+local zip = function(...)
+  local n = numargs(...)
+  if n == 0 then
+    return wrap(nil_gen, nil, nil)
+  end
+  local param = { [2 * n] = 0 }
+  local state = { [n] = 0 }
+
+  local i, gen_x, param_x, state_x
+  for i = 1, n, 1 do
+    local it = select(n - i + 1, ...)
+    gen_x, param_x, state_x = rawiter(it)
+    param[2 * i - 1] = gen_x
+    param[2 * i] = param_x
+    state[i] = state_x
+  end
+
+  return wrap(zip_gen, param, state)
+end
+
+Iterator.zip = zip
+Iterator.__div = zip
+exports.zip = zip
 
 return exports
diff --git a/lua/plenary/operators.lua b/lua/plenary/operators.lua
new file mode 100644
index 0000000..2f80ef1
--- /dev/null
+++ b/lua/plenary/operators.lua
@@ -0,0 +1,50 @@
+---@brief [[
+---Operators that are functions.
+---This is useful when you want to pass operators to higher order functions.
+---Lua has no currying so we have to make a function for each operator.
+---@brief ]]
+
+return {
+  ----------------------------------------------------------------------------
+  -- Comparison operators
+  ----------------------------------------------------------------------------
+  lt  = function(a, b) return a < b end,
+  le  = function(a, b) return a <= b end,
+  eq  = function(a, b) return a == b end,
+  ne  = function(a, b) return a ~= b end,
+  ge  = function(a, b) return a >= b end,
+  gt  = function(a, b) return a > b end,
+
+  ----------------------------------------------------------------------------
+  -- Arithmetic operators
+  ----------------------------------------------------------------------------
+  add = function(a, b) return a + b end,
+  div = function(a, b) return a / b end,
+  floordiv = function(a, b) return math.floor(a/b) end,
+  intdiv = function(a, b)
+    local q = a / b
+    if a >= 0 then return math.floor(q) else return math.ceil(q) end
+  end,
+  mod = function(a, b) return a % b end,
+  mul = function(a, b) return a * b end,
+  neq = function(a) return -a end,
+  unm = function(a) return -a end, -- an alias
+  pow = function(a, b) return a ^ b end,
+  sub = function(a, b) return a - b end,
+  truediv = function(a, b) return a / b end,
+
+  ----------------------------------------------------------------------------
+  -- String operators
+  ----------------------------------------------------------------------------
+  concat = function(a, b) return a..b end,
+  len = function(a) return #a end,
+  length = function(a) return #a end, -- an alias
+
+  ----------------------------------------------------------------------------
+  -- Logical operators
+  ----------------------------------------------------------------------------
+  land = function(a, b) return a and b end,
+  lor = function(a, b) return a or b end,
+  lnot = function(a) return not a end,
+  truth = function(a) return not not a end,
+}
diff --git a/lua/plenary/tbl.lua b/lua/plenary/tbl.lua
index 9de42b1..d1f018c 100644
--- a/lua/plenary/tbl.lua
+++ b/lua/plenary/tbl.lua
@@ -16,4 +16,13 @@ function tbl.apply_defaults(original, defaults)
   return original
 end
 
+function tbl.pack(...)
+  return {n=select('#',...); ...}
+end
+
+function tbl.unpack(t, i, j)
+  return unpack(t, i or 1, j or t.n or #t)
+end
+
+
 return tbl
diff --git a/scratch/iterators.lua b/scratch/iterators.lua
new file mode 100644
index 0000000..ebc1ea2
--- /dev/null
+++ b/scratch/iterators.lua
@@ -0,0 +1,4 @@
+local i = require('plenary.iterators')
+
+dump(i.split("  hello person dude ", " "):tolist())
+-- dump(("hello person"):find("person"))


Test patch:
diff --git a/tests/plenary/iterators_spec.lua b/tests/plenary/iterators_spec.lua
index ab726d2..3717fc2 100644
--- a/tests/plenary/iterators_spec.lua
+++ b/tests/plenary/iterators_spec.lua
@@ -54,4 +54,53 @@ describe('iterators', function()
   it('should make a range', function()
     eq({1, 2, 3, 4, 5}, i.range(5):tolist())
   end)
+
+  it('should be able to make a stateful', function()
+    local iter = i.range(3):stateful()
+    eq(iter(), 1)
+    eq(iter(), 2)
+    eq(iter(), 3)
+    assert(not iter())
+    assert(not iter())
+    assert(not iter())
+
+    local iter = i.range(3):stateful()
+    -- dump(iter:tolist())
+    eq(iter:tolist(), { 1, 2, 3 })
+  end)
+
+  it('should be able to flatten', function()
+    local iter = i.range(3):map(function(_) return i.iter { 5, 7, 9 } end):flatten():stateful()
+    eq(iter(), 5)
+    eq(iter(), 7)
+    eq(iter(), 9)
+    eq(iter(), 5)
+    eq(iter(), 7)
+    eq(iter(), 9)
+    eq(iter(), 5)
+    eq(iter(), 7)
+    eq(iter(), 9)
+    local iter = i.range(3):map(function(_) return i.iter { 5, 7, 9 } end):flatten()
+    eq(iter:tolist(), { 5, 7, 9, 5, 7, 9, 5, 7, 9 })
+  end)
+
+  it('should be able to flatten very nested stuff', function()
+    local iter = i.iter { 5, 6, i.iter { i.iter { 5, 5 }, 7, 8 }, i.iter { 9, 10, i.iter { 1, 2 } } }:flatten()
+    eq(iter:tolist(), { 5, 6, 5, 5, 7, 8, 9, 10, 1, 2 })
+  end)
+
+  it('chaining nil should work', function()
+    local iter = i.iter(""):chain(i.iter { 5, 7, 9 })
+    eq(#iter:tolist(), 3)
+  end)
+
+  describe('generators', function()
+    it('should be able to split', function()
+      local iter = i.split(" hello person world ", " ")
+      eq(iter:tolist(), {"", "hello", "person", "world", ""})
+
+      iter = i.split("\n\n\nfirst\nsecond\nthird\n\n", "\n")
+      eq(iter:tolist(), {"", "", "", "first", "second", "third", "", ""})
+    end)
+  end)
 end)