Implement accumulate and friends for arbitrary iterators

Author: tkf

NEWS.md

@@ -65,7 +65,7 @@ New library features
 
 * `isapprox` (or `≈`) now has a one-argument "curried" method `isapprox(x)` which returns a function, like `isequal` (or `==`)` ([#32305]).
 * `Ref{NTuple{N,T}}` can be passed to `Ptr{T}`/`Ref{T}` `ccall` signatures ([#34199])
-* `accumulate`, `cumsum`, and `cumprod` now support `Tuple` ([#34654]).
+* `accumulate`, `cumsum`, and `cumprod` now support `Tuple` ([#34654]) and arbitrary iterators ([#34656]).
 
 
 Standard library changes

base/accumulate.jl

@@ -92,14 +92,14 @@ function cumsum(A::AbstractArray{T}; dims::Integer) where T
 end
 
 """
-    cumsum(itr::Union{AbstractVector,Tuple})
+    cumsum(itr)
 
 Cumulative sum an iterator. See also [`cumsum!`](@ref)
 to use a preallocated output array, both for performance and to control the precision of the
 output (e.g. to avoid overflow).
 
 !!! compat "Julia 1.5"
-    `cumsum` on a tuple requires at least Julia 1.5.
+    `cumsum` on a non-array iterator requires at least Julia 1.5.
 
 # Examples
 ```jldoctest
@@ -117,6 +117,12 @@ julia> cumsum([fill(1, 2) for i in 1:3])
 
 julia> cumsum((1, 1, 1))
 (1, 2, 3)
+
+julia> cumsum(x^2 for x in 1:3)
+3-element Array{Int64,1}:
+  1
+  5
+ 14
 ```
 """
 cumsum(x::AbstractVector) = cumsum(x, dims=1)
@@ -170,14 +176,14 @@ function cumprod(A::AbstractArray; dims::Integer)
 end
 
 """
-    cumprod(itr::Union{AbstractVector,Tuple})
+    cumprod(itr)
 
 Cumulative product of an iterator. See also
 [`cumprod!`](@ref) to use a preallocated output array, both for performance and
 to control the precision of the output (e.g. to avoid overflow).
 
 !!! compat "Julia 1.5"
-    `cumprod` on a tuple requires at least Julia 1.5.
+    `cumprod` on a non-array iterator requires at least Julia 1.5.
 
 # Examples
 ```jldoctest
@@ -195,6 +201,12 @@ julia> cumprod([fill(1//3, 2, 2) for i in 1:3])
 
 julia> cumprod((1, 2, 1))
 (1, 2, 2)
+
+julia> cumprod(x^2 for x in 1:3)
+3-element Array{Int64,1}:
+  1
+  4
+ 36
 ```
 """
 cumprod(x::AbstractVector) = cumprod(x, dims=1)
@@ -250,6 +262,10 @@ julia> accumulate(+, fill(1, 3, 3), dims=2)
 ```
 """
 function accumulate(op, A; dims::Union{Nothing,Integer}=nothing, kw...)
+    if dims === nothing && !(A isa AbstractVector)
+        # This branch takes care of the cases not handled by `_accumulate!`.
+        return collect(Iterators.accumulate(op, A; kw...))
+    end
     nt = kw.data
     if nt isa NamedTuple{()}
         out = similar(A, promote_op(op, eltype(A), eltype(A)))

base/iterators.jl

@@ -443,13 +443,14 @@ reverse(f::Filter) = Filter(f.flt, reverse(f.itr))
 
 # Accumulate -- partial reductions of a function over an iterator
 
-struct Accumulate{F,I}
+struct Accumulate{F,I,T}
     f::F
     itr::I
+    init::T
 end
 
 """
-    Iterators.accumulate(f, itr)
+    Iterators.accumulate(f, itr; [init])
 
 Given a 2-argument function `f` and an iterator `itr`, return a new
 iterator that successively applies `f` to the previous value and the
@@ -459,24 +460,31 @@ 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> f = Iterators.accumulate(+, [1,2,3,4]);
 
 julia> foreach(println, f)
 1
 3
 6
 10
+
+julia> f = Iterators.accumulate(+, [1,2,3]; init = 100);
+
+julia> foreach(println, f)
+101
+103
+106
 ```
 """
-accumulate(f, itr) = Accumulate(f, itr)
+accumulate(f, itr; init = Base._InitialValue()) = Accumulate(f, itr, init)
 
 function iterate(itr::Accumulate)
     state = iterate(itr.itr)
     if state === nothing
         return nothing
     end
-    return (state[1], state)
+    val = Base.BottomRF(itr.f)(itr.init, state[1])
+    return (val, (val, state[2]))
 end
 
 function iterate(itr::Accumulate, state)
@@ -491,7 +499,7 @@ end
 length(itr::Accumulate) = length(itr.itr)
 size(itr::Accumulate) = size(itr.itr)
 
-IteratorSize(::Type{Accumulate{F,I}}) where {F,I} = IteratorSize(I)
+IteratorSize(::Type{<:Accumulate{F,I}}) where {F,I} = IteratorSize(I)
 IteratorEltype(::Type{<:Accumulate}) = EltypeUnknown()
 
 # Rest -- iterate starting at the given state

test/iterators.jl

@@ -798,3 +798,9 @@ end
     @test Base.IteratorSize(Iterators.accumulate(max, rand(2,3))) === Base.IteratorSize(rand(2,3))
     @test Base.IteratorEltype(Iterators.accumulate(*, ())) isa Base.EltypeUnknown
 end
+
+@testset "Base.accumulate" begin
+    @test cumsum(x^2 for x in 1:3) == [1, 5, 14]
+    @test cumprod(x + 1 for x in 1:3) == [2, 6, 24]
+    @test accumulate(+, (x^2 for x in 1:3); init=100) == [101, 105, 114]
+end