add Iterators.accumulate (#34033)

Author: JeffBezanson

NEWS.md

@@ -36,6 +36,7 @@ New library functions
 * `readdir` output is now guaranteed to be sorted. The `sort` keyword allows opting out of sorting to get names in OS-native order ([#33542]).
 * The new `only(x)` function returns the one-and-only element of a collection `x`, and throws an `ArgumentError` if `x` contains zero or multiple elements. ([#33129])
 * `takewhile` and `dropwhile` have been added to the Iterators submodule ([#33437]).
+* `accumulate` has been added to the Iterators submodule.
 * `filter` can now act on a `Tuple` ([#32968]).
 * There is a now an `evalpoly` (generated) function meant to take the role of the `@evalpoly` macro. The function is just as efficient as the macro while giving added flexibility, so it should be preferred over `@evalpoly`. `evalpoly` takes a list of coefficients as a tuple, so where one might write `@evalpoly(x, p1, p2, p3)` one would instead write `evalpoly(x, (p1, p2, p3))`.
 

base/iterators.jl

@@ -442,6 +442,59 @@ IteratorSize(::Type{<:Filter}) = SizeUnknown()
 
 reverse(f::Filter) = Filter(f.flt, reverse(f.itr))
 
+# Accumulate -- partial reductions of a function over an iterator
+
+struct Accumulate{F,I}
+    f::F
+    itr::I
+end
+
+"""
+    Iterators.accumulate(f, itr)
+
+Given a 2-argument function `f` and an iterator `itr`, return a new
+iterator that successively applies `f` to the previous value and the
+next element of `itr`.
+
+This is effectively a lazy version of [`Base.accumulate`](@ref).
+
+# Examples
+```jldoctest
+julia> f = Iterators.accumulate(+, [1,2,3,4])
+Base.Iterators.Accumulate{typeof(+),Array{Int64,1}}(+, [1, 2, 3, 4])
+
+julia> foreach(println, f)
+1
+3
+6
+10
+```
+"""
+accumulate(f, itr) = Accumulate(f, itr)
+
+function iterate(itr::Accumulate)
+    state = iterate(itr.itr)
+    if state === nothing
+        return nothing
+    end
+    return (state[1], state)
+end
+
+function iterate(itr::Accumulate, state)
+    nxt = iterate(itr.itr, state[2])
+    if nxt === nothing
+        return nothing
+    end
+    val = itr.f(state[1], nxt[1])
+    return (val, (val, nxt[2]))
+end
+
+length(itr::Accumulate) = length(itr.itr)
+size(itr::Accumulate) = size(itr.itr)
+
+IteratorSize(::Type{Accumulate{F,I}}) where {F,I} = IteratorSize(I)
+IteratorEltype(::Type{<:Accumulate}) = EltypeUnknown()
+
 # Rest -- iterate starting at the given state
 
 struct Rest{I,S}

doc/src/base/iterators.md

@@ -14,6 +14,7 @@ Base.Iterators.product
 Base.Iterators.flatten
 Base.Iterators.partition
 Base.Iterators.filter
+Base.Iterators.accumulate
 Base.Iterators.reverse
 Base.Iterators.only
 Base.Iterators.peel

test/iterators.jl

@@ -784,3 +784,15 @@ end
 @testset "flatten empty tuple" begin
     @test isempty(collect(Iterators.flatten(())))
 end
+
+@testset "Iterators.accumulate" begin
+    @test collect(Iterators.accumulate(+, [])) == []
+    @test collect(Iterators.accumulate(+, [1])) == [1]
+    @test collect(Iterators.accumulate(+, [1,2])) == [1,3]
+    @test collect(Iterators.accumulate(+, [1,2,3])) == [1,3,6]
+    @test collect(Iterators.accumulate(=>, [:a,:b,:c])) == [:a, :a => :b, (:a => :b) => :c]
+    @test length(Iterators.accumulate(+, [10,20,30])) == 3
+    @test size(Iterators.accumulate(max, rand(2,3))) == (2,3)
+    @test Base.IteratorSize(Iterators.accumulate(max, rand(2,3))) === Base.IteratorSize(rand(2,3))
+    @test Base.IteratorEltype(Iterators.accumulate(*, ())) isa Base.EltypeUnknown
+end