From 5701e79a39640e57641cdcd455a1a00c5670e027 Mon Sep 17 00:00:00 2001 From: Harmen Stoppels Date: Sun, 3 Jun 2018 23:39:59 +0300 Subject: [PATCH] Simplify zip iterator --- base/iterators.jl | 163 ++++++++++++++-------------------------------- 1 file changed, 49 insertions(+), 114 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index b11a6dfb636ba..6de4ff316c304 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -255,142 +255,77 @@ get(f::Base.Callable, collection::Pairs, key) = get(f, v.data, key) abstract type AbstractZipIterator end -zip_iteratorsize(a, b) = and_iteratorsize(a,b) # as `and_iteratorsize` but inherit `Union{HasLength,IsInfinite}` of the shorter iterator +# as `and_iteratorsize` but inherit `Union{HasLength,IsInfinite}` of the shorter iterator +zip_iteratorsize(a, b) = and_iteratorsize(a,b) 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 +struct ZipHead{I} <: AbstractZipIterator a::I end -zip(a) = Zip1(a) -length(z::Zip1) = length(z.a) -size(z::Zip1) = size(z.a) -axes(z::Zip1) = axes(z.a) -eltype(::Type{Zip1{I}}) where {I} = Tuple{eltype(I)} -@propagate_inbounds function iterate(z::Zip1, state...) - n = iterate(z.a, state...) - n === nothing && return n - return ((n[1],), n[2]) -end -@inline isdone(z::Zip1, state...) = isdone(z.a, state...) - -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)) -axes(z::Zip2) = promote_shape(axes(z.a), axes(z.b)) -eltype(::Type{Zip2{I1,I2}}) where {I1,I2} = Tuple{eltype(I1), eltype(I2)} -@inline isdone(z::Zip2) = isdone(z.a) | isdone(z.b) -@inline isdone(z::Zip2, (sa, sb)::Tuple{Any, Any}) = isdone(z.a, sa) | isdone(z.b, sb) -function zip_iterate(a, b, sta, stb) # the states are either Tuple{} or Tuple{Any} - da, db = isdone(a), isdone(b) - da === true && return nothing - db === true && return nothing - if da === missing - ya = iterate(a, sta...) - ya === nothing && return nothing - end - if db === missing - yb = iterate(b, stb...) - yb === nothing && return nothing - end - if da === false - ya = iterate(a, sta...) - ya === nothing && return nothing - end - if db === false - yb = iterate(b, stb...) - yb === nothing && return nothing - end - return (ya, yb) -end -let interleave(a, b) = ((a[1], b[1]), (a[2], b[2])) - global iterate - @propagate_inbounds function iterate(z::Zip2) - ys = zip_iterate(z.a, z.b, (), ()) - ys === nothing && return nothing - return interleave(ys...) - end - @propagate_inbounds function iterate(z::Zip2, st::Tuple{Any, Any}) - ys = zip_iterate(z.a, z.b, (st[1],), (st[2],)) - ys === nothing && return nothing - return interleave(ys...) - end -end - -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...) +@inline isdone(it::ZipHead, state...) = isdone(it.a, state...) +@inline isdone(it::Zip) = isdone(it.a) | isdone(it.z) +@inline isdone(it::Zip, (sa, sz)) = isdone(it.a, sa) | isdone(it.z, sz) -For a set of iterable objects, return an iterable of tuples, where the `i`th tuple contains -the `i`th component of each input iterable. +length(it::ZipHead) = length(it.a) +length(it::Zip) = _min_length(it.a, it.z, IteratorSize(it.a), IteratorSize(it.z)) -# Examples -```jldoctest -julia> a = 1:5 -1:5 +size(z::ZipHead) = size(z.a) +size(it::Zip) = promote_shape(size(it.a), size(it.z)) -julia> b = ["e","d","b","c","a"] -5-element Array{String,1}: - "e" - "d" - "b" - "c" - "a" +axes(z::ZipHead) = axes(z.a) +axes(it::Zip) = promote_shape(axes(it.a), axes(it.z)) -julia> c = zip(a,b) -Base.Iterators.Zip2{UnitRange{Int64},Array{String,1}}(1:5, ["e", "d", "b", "c", "a"]) +eltype(::Type{ZipHead{I}}) where {I} = Tuple{eltype(I)} +eltype(::Type{Zip{I,Z}}) where {I,Z} = tuple_type_cons(eltype(I), eltype(Z)) -julia> length(c) -5 +IteratorSize(::Type{ZipHead{I}}) where {I} = IteratorSize(I) +IteratorSize(::Type{Zip{I,Z}}) where {I,Z} = zip_iteratorsize(IteratorSize(I),IteratorSize(Z)) -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)) -axes(z::Zip) = promote_shape(axes(z.a), axes(z.z)) -eltype(::Type{Zip{I,Z}}) where {I,Z} = tuple_type_cons(eltype(I), eltype(Z)) -@inline isdone(z::Zip) = isdone(z.a) | isdone(z.z) -@inline isdone(z::Zip, (sa, sz)) = isdone(z.a, sa) | isdone(z.a, sz) -let interleave(a, b) = ((a[1], b[1]...), (a[2], b[2])) - global iterate - @propagate_inbounds function iterate(z::Zip) - ys = zip_iterate(z.a, z.z, (), ()) - ys === nothing && return nothing - return interleave(ys...) - end - @propagate_inbounds function iterate(z::Zip, st::Tuple{Any, Any}) - ys = zip_iterate(z.a, z.z, (st[1],), (st[2],)) - ys === nothing && return nothing - return interleave(ys...) - end +IteratorEltype(::Type{ZipHead{I}}) where {I} = IteratorEltype(I) +IteratorEltype(::Type{Zip{I,Z}}) where {I,Z} = and_iteratoreltype(IteratorEltype(I),IteratorEltype(Z)) + +reverse(it::ZipHead) = ZipHead(reverse(it.a)) +reverse(it::Zip) = Zip(reverse(it.a), reverse(it.z)) + +zip(a) = ZipHead(a) +zip(a, b...) = Zip(a, zip(b...)) + +# Iteration +@inline function iterate(it::ZipHead, state...) + isdone(it.a) === true && return nothing + a = iterate(it.a, state...) + a === nothing && return nothing + ((a[1],), a[2]) end -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)) +@inline function iterate(it::Zip) + isdone(it.a) === true && return nothing + isdone(it.z) === true && return nothing + a = iterate(it.a) + z = iterate(it.z) + a === nothing && return nothing + z === nothing && return nothing + ((a[1], z[1]...), (a[2], z[2])) +end -reverse(z::Zip1) = Zip1(reverse(z.a)) -reverse(z::Zip2) = Zip2(reverse(z.a), reverse(z.b)) -reverse(z::Zip) = Zip(reverse(z.a), reverse(z.z)) +@inline function iterate(it::Zip, state::Tuple{Any, Any}) + isdone(it.a) === true && return nothing + isdone(it.z) === true && return nothing + a = iterate(it.a, state[1]) + z = iterate(it.z, state[2]) + a === nothing && return nothing + z === nothing && return nothing + ((a[1], z[1]...), (a[2], z[2])) +end # filter