# This file is a part of Julia. License is MIT: https://julialang.org/license module Iterators import Base: start, done, next, isempty, length, size, eltype, iteratorsize, iteratoreltype, indices, ndims using Base: tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, @propagate_inbounds export enumerate, zip, rest, countfrom, take, drop, cycle, repeated, product, flatten, partition _min_length(a, b, ::IsInfinite, ::IsInfinite) = min(length(a),length(b)) # inherit behaviour, error _min_length(a, b, A, ::IsInfinite) = length(a) _min_length(a, b, ::IsInfinite, B) = length(b) _min_length(a, b, A, B) = min(length(a),length(b)) _diff_length(a, b, A, ::IsInfinite) = 0 _diff_length(a, b, ::IsInfinite, ::IsInfinite) = 0 _diff_length(a, b, ::IsInfinite, B) = length(a) # inherit behaviour, error _diff_length(a, b, A, B) = max(length(a)-length(b), 0) and_iteratorsize{T}(isz::T, ::T) = isz and_iteratorsize(::HasLength, ::HasShape) = HasLength() and_iteratorsize(::HasShape, ::HasLength) = HasLength() and_iteratorsize(a, b) = SizeUnknown() and_iteratoreltype(iel::T, ::T) where {T} = iel and_iteratoreltype(a, b) = EltypeUnknown() # enumerate struct Enumerate{I} itr::I end """ enumerate(iter) An iterator that yields `(i, x)` where `i` is a counter starting at 1, and `x` is the `i`th value from the given iterator. It's useful when you need not only the values `x` over which you are iterating, but also the number of iterations so far. Note that `i` may not be valid for indexing `iter`; it's also possible that `x != iter[i]`, if `iter` has indices that do not start at 1. See the `enumerate(IndexLinear(), iter)` method if you want to ensure that `i` is an index. # Example ```jldoctest julia> a = ["a", "b", "c"]; julia> for (index, value) in enumerate(a) println("\$index \$value") end 1 a 2 b 3 c ``` """ enumerate(iter) = Enumerate(iter) length(e::Enumerate) = length(e.itr) size(e::Enumerate) = size(e.itr) @inline start(e::Enumerate) = (1, start(e.itr)) @inline function next(e::Enumerate, state) n = next(e.itr,state[2]) (state[1],n[1]), (state[1]+1,n[2]) end @inline done(e::Enumerate, state) = done(e.itr, state[2]) eltype(::Type{Enumerate{I}}) where {I} = Tuple{Int, eltype(I)} iteratorsize(::Type{Enumerate{I}}) where {I} = iteratorsize(I) iteratoreltype(::Type{Enumerate{I}}) where {I} = iteratoreltype(I) struct IndexValue{I,A<:AbstractArray} data::A itr::I end """ enumerate(IndexLinear(), A) enumerate(IndexCartesian(), A) enumerate(IndexStyle(A), A) An iterator that accesses each element of the array `A`, returning `(i, x)`, where `i` is the index for the element and `x = A[i]`. This is similar to `enumerate(A)`, except `i` will always be a valid index for `A`. Specifying `IndexLinear()` ensures that `i` will be an integer; specifying `IndexCartesian()` ensures that `i` will be a `CartesianIndex`; specifying `IndexStyle(A)` chooses whichever has been defined as the native indexing style for array `A`. # Examples ```jldoctest julia> A = ["a" "d"; "b" "e"; "c" "f"]; julia> for (index, value) in enumerate(IndexStyle(A), A) println("\$index \$value") end 1 a 2 b 3 c 4 d 5 e 6 f julia> S = view(A, 1:2, :); julia> for (index, value) in enumerate(IndexStyle(S), S) println("\$index \$value") end CartesianIndex{2}((1, 1)) a CartesianIndex{2}((2, 1)) b CartesianIndex{2}((1, 2)) d CartesianIndex{2}((2, 2)) e ``` Note that `enumerate(A)` returns `i` as a *counter* (always starting at 1), whereas `enumerate(IndexLinear(), A)` returns `i` as an *index* (starting at the first linear index of `A`, which may or may not be 1). See also: [`IndexStyle`](@ref), [`indices`](@ref). """ enumerate(::IndexLinear, A::AbstractArray) = IndexValue(A, linearindices(A)) enumerate(::IndexCartesian, A::AbstractArray) = IndexValue(A, CartesianRange(indices(A))) length(v::IndexValue) = length(v.itr) indices(v::IndexValue) = indices(v.itr) size(v::IndexValue) = size(v.itr) @inline start(v::IndexValue) = start(v.itr) @propagate_inbounds function next(v::IndexValue, state) indx, n = next(v.itr, state) item = v.data[indx] (indx, item), n end @inline done(v::IndexValue, state) = done(v.itr, state) eltype(::Type{IndexValue{I,A}}) where {I,A} = Tuple{eltype(I), eltype(A)} iteratorsize(::Type{IndexValue{I}}) where {I} = iteratorsize(I) iteratoreltype(::Type{IndexValue{I}}) where {I} = iteratoreltype(I) # zip abstract type AbstractZipIterator end zip_iteratorsize(a, b) = and_iteratorsize(a,b) # as `and_iteratorsize` but inherit `Union{HasLength,IsInfinite}` of the shorter iterator zip_iteratorsize(::HasLength, ::IsInfinite) = HasLength() zip_iteratorsize(::HasShape, ::IsInfinite) = HasLength() zip_iteratorsize(a::IsInfinite, b) = zip_iteratorsize(b,a) zip_iteratorsize(a::IsInfinite, b::IsInfinite) = IsInfinite() struct Zip1{I} <: AbstractZipIterator a::I end zip(a) = Zip1(a) length(z::Zip1) = length(z.a) size(z::Zip1) = size(z.a) indices(z::Zip1) = indices(z.a) eltype(::Type{Zip1{I}}) where {I} = Tuple{eltype(I)} @inline start(z::Zip1) = start(z.a) @inline function next(z::Zip1, st) n = next(z.a,st) return ((n[1],), n[2]) end @inline done(z::Zip1, st) = done(z.a,st) iteratorsize(::Type{Zip1{I}}) where {I} = iteratorsize(I) iteratoreltype(::Type{Zip1{I}}) where {I} = iteratoreltype(I) struct Zip2{I1, I2} <: AbstractZipIterator a::I1 b::I2 end zip(a, b) = Zip2(a, b) length(z::Zip2) = _min_length(z.a, z.b, iteratorsize(z.a), iteratorsize(z.b)) size(z::Zip2) = promote_shape(size(z.a), size(z.b)) indices(z::Zip2) = promote_shape(indices(z.a), indices(z.b)) eltype(::Type{Zip2{I1,I2}}) where {I1,I2} = Tuple{eltype(I1), eltype(I2)} @inline start(z::Zip2) = (start(z.a), start(z.b)) @inline function next(z::Zip2, st) n1 = next(z.a,st[1]) n2 = next(z.b,st[2]) return ((n1[1], n2[1]), (n1[2], n2[2])) end @inline done(z::Zip2, st) = done(z.a,st[1]) | done(z.b,st[2]) iteratorsize(::Type{Zip2{I1,I2}}) where {I1,I2} = zip_iteratorsize(iteratorsize(I1),iteratorsize(I2)) iteratoreltype(::Type{Zip2{I1,I2}}) where {I1,I2} = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2)) struct Zip{I, Z<:AbstractZipIterator} <: AbstractZipIterator a::I z::Z end """ zip(iters...) For a set of iterable objects, returns an iterable of tuples, where the `i`th tuple contains the `i`th component of each input iterable. Note that [`zip`](@ref) is its own inverse: `collect(zip(zip(a...)...)) == collect(a)`. # Example ```jldoctest julia> a = 1:5 1:5 julia> b = ["e","d","b","c","a"] 5-element Array{String,1}: "e" "d" "b" "c" "a" julia> c = zip(a,b) Base.Iterators.Zip2{UnitRange{Int64},Array{String,1}}(1:5, String["e", "d", "b", "c", "a"]) julia> length(c) 5 julia> first(c) (1, "e") ``` """ zip(a, b, c...) = Zip(a, zip(b, c...)) length(z::Zip) = _min_length(z.a, z.z, iteratorsize(z.a), iteratorsize(z.z)) size(z::Zip) = promote_shape(size(z.a), size(z.z)) indices(z::Zip) = promote_shape(indices(z.a), indices(z.z)) eltype(::Type{Zip{I,Z}}) where {I,Z} = tuple_type_cons(eltype(I), eltype(Z)) @inline start(z::Zip) = tuple(start(z.a), start(z.z)) @inline function next(z::Zip, st) n1 = next(z.a, st[1]) n2 = next(z.z, st[2]) (tuple(n1[1], n2[1]...), (n1[2], n2[2])) end @inline done(z::Zip, st) = done(z.a,st[1]) | done(z.z,st[2]) iteratorsize(::Type{Zip{I1,I2}}) where {I1,I2} = zip_iteratorsize(iteratorsize(I1),iteratorsize(I2)) iteratoreltype(::Type{Zip{I1,I2}}) where {I1,I2} = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2)) # filter struct Filter{F,I} flt::F itr::I end """ Iterators.filter(flt, itr) Given a predicate function `flt` and an iterable object `itr`, return an iterable object which upon iteration yields the elements `x` of `itr` that satisfy `flt(x)`. The order of the original iterator is preserved. This function is *lazy*; that is, it is guaranteed to return in ``Θ(1)`` time and use ``Θ(1)`` additional space, and `flt` will not be called by an invocation of `filter`. Calls to `flt` will be made when iterating over the returned iterable object. These calls are not cached and repeated calls will be made when reiterating. See [`Base.filter`](@ref) for an eager implementation of filtering for arrays. """ filter(flt, itr) = Filter(flt, itr) start(f::Filter) = start_filter(f.flt, f.itr) function start_filter(pred, itr) s = start(itr) while !done(itr,s) v,t = next(itr,s) if pred(v) return (false, v, t) end s=t end (true,) end next(f::Filter, s) = advance_filter(f.flt, f.itr, s) function advance_filter(pred, itr, st) _, v, s = st while !done(itr,s) w,t = next(itr,s) if pred(w) return v, (false, w, t) end s=t end v, (true, v, s) end done(f::Filter, s) = s[1] eltype(::Type{Filter{F,I}}) where {F,I} = eltype(I) iteratoreltype(::Type{Filter{F,I}}) where {F,I} = iteratoreltype(I) iteratorsize(::Type{<:Filter}) = SizeUnknown() # Rest -- iterate starting at the given state struct Rest{I,S} itr::I st::S end """ rest(iter, state) An iterator that yields the same elements as `iter`, but starting at the given `state`. """ rest(itr,state) = Rest(itr,state) start(i::Rest) = i.st next(i::Rest, st) = next(i.itr, st) done(i::Rest, st) = done(i.itr, st) eltype(::Type{Rest{I}}) where {I} = eltype(I) iteratoreltype(::Type{Rest{I,S}}) where {I,S} = iteratoreltype(I) rest_iteratorsize(a) = SizeUnknown() rest_iteratorsize(::IsInfinite) = IsInfinite() iteratorsize(::Type{Rest{I,S}}) where {I,S} = rest_iteratorsize(iteratorsize(I)) # Count -- infinite counting struct Count{S<:Number} start::S step::S end """ countfrom(start=1, step=1) An iterator that counts forever, starting at `start` and incrementing by `step`. """ countfrom(start::Number, step::Number) = Count(promote(start, step)...) countfrom(start::Number) = Count(start, oneunit(start)) countfrom() = Count(1, 1) eltype(::Type{Count{S}}) where {S} = S start(it::Count) = it.start next(it::Count, state) = (state, state + it.step) done(it::Count, state) = false iteratorsize(::Type{<:Count}) = IsInfinite() # Take -- iterate through the first n elements struct Take{I} xs::I n::Int end """ take(iter, n) An iterator that generates at most the first `n` elements of `iter`. # Example ```jldoctest julia> a = 1:2:11 1:2:11 julia> collect(a) 6-element Array{Int64,1}: 1 3 5 7 9 11 julia> collect(Iterators.take(a,3)) 3-element Array{Int64,1}: 1 3 5 ``` """ take(xs, n::Integer) = Take(xs, Int(n)) take(xs::Take, n::Integer) = Take(xs.xs, min(Int(n), xs.n)) eltype(::Type{Take{I}}) where {I} = eltype(I) iteratoreltype(::Type{Take{I}}) where {I} = iteratoreltype(I) take_iteratorsize(a) = HasLength() take_iteratorsize(::SizeUnknown) = SizeUnknown() iteratorsize(::Type{Take{I}}) where {I} = take_iteratorsize(iteratorsize(I)) length(t::Take) = _min_length(t.xs, 1:t.n, iteratorsize(t.xs), HasLength()) start(it::Take) = (it.n, start(it.xs)) function next(it::Take, state) n, xs_state = state v, xs_state = next(it.xs, xs_state) return v, (n - 1, xs_state) end function done(it::Take, state) n, xs_state = state return n <= 0 || done(it.xs, xs_state) end # Drop -- iterator through all but the first n elements struct Drop{I} xs::I n::Int end """ drop(iter, n) An iterator that generates all but the first `n` elements of `iter`. # Example ```jldoctest julia> a = 1:2:11 1:2:11 julia> collect(a) 6-element Array{Int64,1}: 1 3 5 7 9 11 julia> collect(Iterators.drop(a,4)) 2-element Array{Int64,1}: 9 11 ``` """ drop(xs, n::Integer) = Drop(xs, Int(n)) drop(xs::Take, n::Integer) = Take(drop(xs.xs, Int(n)), max(0, xs.n - Int(n))) drop(xs::Drop, n::Integer) = Drop(xs.xs, Int(n) + xs.n) eltype(::Type{Drop{I}}) where {I} = eltype(I) iteratoreltype(::Type{Drop{I}}) where {I} = iteratoreltype(I) drop_iteratorsize(::SizeUnknown) = SizeUnknown() drop_iteratorsize(::Union{HasShape, HasLength}) = HasLength() drop_iteratorsize(::IsInfinite) = IsInfinite() iteratorsize(::Type{Drop{I}}) where {I} = drop_iteratorsize(iteratorsize(I)) length(d::Drop) = _diff_length(d.xs, 1:d.n, iteratorsize(d.xs), HasLength()) function start(it::Drop) xs_state = start(it.xs) for i in 1:it.n if done(it.xs, xs_state) break end _, xs_state = next(it.xs, xs_state) end xs_state end next(it::Drop, state) = next(it.xs, state) done(it::Drop, state) = done(it.xs, state) # Cycle an iterator forever struct Cycle{I} xs::I end """ cycle(iter) An iterator that cycles through `iter` forever. """ cycle(xs) = Cycle(xs) eltype(::Type{Cycle{I}}) where {I} = eltype(I) iteratoreltype(::Type{Cycle{I}}) where {I} = iteratoreltype(I) iteratorsize(::Type{Cycle{I}}) where {I} = IsInfinite() function start(it::Cycle) s = start(it.xs) return s, done(it.xs, s) end function next(it::Cycle, state) s, d = state if done(it.xs, s) s = start(it.xs) end v, s = next(it.xs, s) return v, (s, false) end done(it::Cycle, state) = state[2] # Repeated - repeat an object infinitely many times struct Repeated{O} x::O end repeated(x) = Repeated(x) """ repeated(x[, n::Int]) An iterator that generates the value `x` forever. If `n` is specified, generates `x` that many times (equivalent to `take(repeated(x), n)`). # Example ```jldoctest julia> a = Iterators.repeated([1 2], 4); julia> collect(a) 4-element Array{Array{Int64,2},1}: [1 2] [1 2] [1 2] [1 2] ``` """ repeated(x, n::Integer) = take(repeated(x), Int(n)) eltype(::Type{Repeated{O}}) where {O} = O start(it::Repeated) = nothing next(it::Repeated, state) = (it.x, nothing) done(it::Repeated, state) = false iteratorsize(::Type{<:Repeated}) = IsInfinite() iteratoreltype(::Type{<:Repeated}) = HasEltype() # Product -- cartesian product of iterators abstract type AbstractProdIterator end length(p::AbstractProdIterator) = prod(size(p)) _length(p::AbstractProdIterator) = prod(map(unsafe_length, indices(p))) size(p::AbstractProdIterator) = _prod_size(p.a, p.b, iteratorsize(p.a), iteratorsize(p.b)) indices(p::AbstractProdIterator) = _prod_indices(p.a, p.b, iteratorsize(p.a), iteratorsize(p.b)) ndims(p::AbstractProdIterator) = length(indices(p)) # generic methods to handle size of Prod* types _prod_size(a, ::HasShape) = size(a) _prod_size(a, ::HasLength) = (length(a),) _prod_size(a, A) = throw(ArgumentError("Cannot compute size for object of type $(typeof(a))")) _prod_size(a, b, ::HasLength, ::HasLength) = (length(a), length(b)) _prod_size(a, b, ::HasLength, ::HasShape) = (length(a), size(b)...) _prod_size(a, b, ::HasShape, ::HasLength) = (size(a)..., length(b)) _prod_size(a, b, ::HasShape, ::HasShape) = (size(a)..., size(b)...) _prod_size(a, b, A, B) = throw(ArgumentError("Cannot construct size for objects of types $(typeof(a)) and $(typeof(b))")) _prod_indices(a, ::HasShape) = indices(a) _prod_indices(a, ::HasLength) = (OneTo(length(a)),) _prod_indices(a, A) = throw(ArgumentError("Cannot compute indices for object of type $(typeof(a))")) _prod_indices(a, b, ::HasLength, ::HasLength) = (OneTo(length(a)), OneTo(length(b))) _prod_indices(a, b, ::HasLength, ::HasShape) = (OneTo(length(a)), indices(b)...) _prod_indices(a, b, ::HasShape, ::HasLength) = (indices(a)..., OneTo(length(b))) _prod_indices(a, b, ::HasShape, ::HasShape) = (indices(a)..., indices(b)...) _prod_indices(a, b, A, B) = throw(ArgumentError("Cannot construct indices for objects of types $(typeof(a)) and $(typeof(b))")) # one iterator struct Prod1{I} <: AbstractProdIterator a::I end product(a) = Prod1(a) eltype(::Type{Prod1{I}}) where {I} = Tuple{eltype(I)} size(p::Prod1) = _prod_size(p.a, iteratorsize(p.a)) indices(p::Prod1) = _prod_indices(p.a, iteratorsize(p.a)) @inline start(p::Prod1) = start(p.a) @inline function next(p::Prod1, st) n, st = next(p.a, st) (n,), st end @inline done(p::Prod1, st) = done(p.a, st) iteratoreltype(::Type{Prod1{I}}) where {I} = iteratoreltype(I) iteratorsize(::Type{Prod1{I}}) where {I} = iteratorsize(I) # two iterators struct Prod2{I1, I2} <: AbstractProdIterator a::I1 b::I2 end """ product(iters...) Returns an iterator over the product of several iterators. Each generated element is a tuple whose `i`th element comes from the `i`th argument iterator. The first iterator changes the fastest. Example: # Example ```jldoctest julia> collect(Iterators.product(1:2,3:5)) 2×3 Array{Tuple{Int64,Int64},2}: (1, 3) (1, 4) (1, 5) (2, 3) (2, 4) (2, 5) ``` """ product(a, b) = Prod2(a, b) eltype(::Type{Prod2{I1,I2}}) where {I1,I2} = Tuple{eltype(I1), eltype(I2)} iteratoreltype(::Type{Prod2{I1,I2}}) where {I1,I2} = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2)) iteratorsize(::Type{Prod2{I1,I2}}) where {I1,I2} = prod_iteratorsize(iteratorsize(I1),iteratorsize(I2)) function start(p::AbstractProdIterator) s1, s2 = start(p.a), start(p.b) s1, s2, Nullable{eltype(p.b)}(), (done(p.a,s1) || done(p.b,s2)) end @inline function prod_next(p, st) s1, s2 = st[1], st[2] v1, s1 = next(p.a, s1) nv2 = st[3] if isnull(nv2) v2, s2 = next(p.b, s2) else v2 = nv2.value end if done(p.a, s1) return (v1,v2), (start(p.a), s2, Nullable{eltype(nv2)}(), done(p.b,s2)) end return (v1,v2), (s1, s2, Nullable(v2), false) end @inline next(p::Prod2, st) = prod_next(p, st) @inline done(p::AbstractProdIterator, st) = st[4] # n iterators struct Prod{I1, I2<:AbstractProdIterator} <: AbstractProdIterator a::I1 b::I2 end product(a, b, c...) = Prod(a, product(b, c...)) eltype(::Type{Prod{I1,I2}}) where {I1,I2} = tuple_type_cons(eltype(I1), eltype(I2)) iteratoreltype(::Type{Prod{I1,I2}}) where {I1,I2} = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2)) iteratorsize(::Type{Prod{I1,I2}}) where {I1,I2} = prod_iteratorsize(iteratorsize(I1),iteratorsize(I2)) @inline function next{I1,I2}(p::Prod{I1,I2}, st) x = prod_next(p, st) ((x[1][1],x[1][2]...), x[2]) end prod_iteratorsize(::Union{HasLength,HasShape}, ::Union{HasLength,HasShape}) = HasShape() # products can have an infinite iterator prod_iteratorsize(::IsInfinite, ::IsInfinite) = IsInfinite() prod_iteratorsize(a, ::IsInfinite) = IsInfinite() prod_iteratorsize(::IsInfinite, b) = IsInfinite() prod_iteratorsize(a, b) = SizeUnknown() # flatten an iterator of iterators struct Flatten{I} it::I end """ flatten(iter) Given an iterator that yields iterators, return an iterator that yields the elements of those iterators. Put differently, the elements of the argument iterator are concatenated. # Example ```jldoctest julia> collect(Iterators.flatten((1:2, 8:9))) 4-element Array{Int64,1}: 1 2 8 9 ``` """ flatten(itr) = Flatten(itr) eltype(::Type{Flatten{I}}) where {I} = eltype(eltype(I)) iteratorsize(::Type{Flatten{I}}) where {I} = SizeUnknown() iteratoreltype(::Type{Flatten{I}}) where {I} = _flatteneltype(I, iteratoreltype(I)) _flatteneltype(I, ::HasEltype) = iteratoreltype(eltype(I)) _flatteneltype(I, et) = EltypeUnknown() function start(f::Flatten) local inner, s2 s = start(f.it) d = done(f.it, s) # this is a simple way to make this function type stable d && throw(ArgumentError("argument to Flatten must contain at least one iterator")) while !d inner, s = next(f.it, s) s2 = start(inner) !done(inner, s2) && break d = done(f.it, s) end return s, inner, s2 end @inline function next(f::Flatten, state) s, inner, s2 = state val, s2 = next(inner, s2) while done(inner, s2) && !done(f.it, s) inner, s = next(f.it, s) s2 = start(inner) end return val, (s, inner, s2) end @inline function done(f::Flatten, state) s, inner, s2 = state return done(f.it, s) && done(inner, s2) end """ partition(collection, n) Iterate over a collection `n` elements at a time. # Example ```jldoctest julia> collect(Iterators.partition([1,2,3,4,5], 2)) 3-element Array{Array{Int64,1},1}: [1, 2] [3, 4] [5] ``` """ partition(c::T, n::Integer) where {T} = PartitionIterator{T}(c, Int(n)) mutable struct PartitionIterator{T} c::T n::Int end eltype(::Type{PartitionIterator{T}}) where {T} = Vector{eltype(T)} function length(itr::PartitionIterator) l = length(itr.c) return div(l, itr.n) + ((mod(l, itr.n) > 0) ? 1 : 0) end start(itr::PartitionIterator) = start(itr.c) done(itr::PartitionIterator, state) = done(itr.c, state) function next(itr::PartitionIterator{<:Vector}, state) l = state r = min(state + itr.n-1, length(itr.c)) return view(itr.c, l:r), r + 1 end function next(itr::PartitionIterator, state) v = Vector{eltype(itr.c)}(itr.n) i = 0 while !done(itr.c, state) && i < itr.n i += 1 v[i], state = next(itr.c, state) end return resize!(v, i), state end end