feat: add Iterator:fold and Iterator:enumerate (#505)

Author: willothy

lua/plenary/iterators.lua

@@ -486,6 +486,16 @@ function Iterator:filter(fun)
   return wrap(filter_gen, { self.gen, self.param, fun }, self.state)
 end
 
+---Iterator adapter that will provide numbers from 1 to n as the first multival
+---@return Iterator
+function Iterator:enumerate()
+  local i = 0
+  return self:map(function(...)
+    i = i + 1
+    return i, ...
+  end)
+end
+
 --------------------------------------------------------------------------------
 -- Reducing
 --------------------------------------------------------------------------------
@@ -531,6 +541,19 @@ function Iterator:find(val_or_fn)
   end
 end
 
+---Folds an iterator into a single value using a function.
+---@param init any
+---@param fun fun(acc: any, val: any): any
+---@return any
+function Iterator:fold(init, fun)
+  local acc = init
+  local gen, param, state = self.gen, self.param, self.state
+  for _, r in gen, param, state do
+    acc = fun(acc, r)
+  end
+  return acc
+end
+
 ---Turns an iterator into a list.
 ---If the iterator yields multivals only the first multival will be used.
 ---@return table

tests/plenary/iterators_spec.lua

@@ -31,6 +31,35 @@ describe("iterators", function()
     check_keys(tbl, { 1, 2, 3, 4 })
   end)
 
+  it("should be able to fold", function()
+    local numbers = { 1, 2, 3, 4 }
+    local result = i.iter(numbers):fold(0, function(a, v)
+      return a + v
+    end)
+    eq(result, 10)
+
+    local strings = { "hello", "world", "this", "is", "a", "test" }
+    result = i.iter(strings):fold("", function(a, v)
+      return a .. v
+    end)
+    eq(result, "helloworldthisisatest")
+  end)
+
+  it("should be able to enumerate", function()
+    local tbl = { 1, 2, 3, 4 }
+    local results = i.iter(tbl)
+      :enumerate()
+      :map(function(idx, v)
+        return { idx, v }
+      end)
+      :tolist()
+    eq(#results, 4)
+    eq(results[1], { 1, 1 })
+    eq(results[2], { 2, 2 })
+    eq(results[3], { 3, 3 })
+    eq(results[4], { 4, 4 })
+  end)
+
   it("should be able to find", function()
     local tbl = { 1, 2, 3, 4 }
     local tbl_iter = i.iter(tbl)