diff --git a/base/iterator.jl b/base/iterator.jl index 0323aa1548148..d80659013fc5d 100644 --- a/base/iterator.jl +++ b/base/iterator.jl @@ -2,6 +2,16 @@ isempty(itr) = done(itr, start(itr)) +_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) + # enumerate immutable Enumerate{I} @@ -27,6 +37,12 @@ iteratoreltype{I}(::Type{Enumerate{I}}) = iteratoreltype(I) abstract AbstractZipIterator +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) + + immutable Zip1{I} <: AbstractZipIterator a::I end @@ -49,7 +65,7 @@ immutable Zip2{I1, I2} <: AbstractZipIterator b::I2 end zip(a, b) = Zip2(a, b) -length(z::Zip2) = min(length(z.a), length(z.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)) eltype{I1,I2}(::Type{Zip2{I1,I2}}) = Tuple{eltype(I1), eltype(I2)} @inline start(z::Zip2) = (start(z.a), start(z.b)) @@ -60,7 +76,7 @@ eltype{I1,I2}(::Type{Zip2{I1,I2}}) = Tuple{eltype(I1), eltype(I2)} end @inline done(z::Zip2, st) = done(z.a,st[1]) | done(z.b,st[2]) -iteratorsize{I1,I2}(::Type{Zip2{I1,I2}}) = and_iteratorsize(iteratorsize(I1),iteratorsize(I2)) +iteratorsize{I1,I2}(::Type{Zip2{I1,I2}}) = zip_iteratorsize(iteratorsize(I1),iteratorsize(I2)) iteratoreltype{I1,I2}(::Type{Zip2{I1,I2}}) = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2)) immutable Zip{I, Z<:AbstractZipIterator} <: AbstractZipIterator @@ -68,7 +84,7 @@ immutable Zip{I, Z<:AbstractZipIterator} <: AbstractZipIterator z::Z end zip(a, b, c...) = Zip(a, zip(b, c...)) -length(z::Zip) = min(length(z.a), length(z.z)) +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)) tuple_type_cons{S}(::Type{S}, ::Type{Union{}}) = Union{} function tuple_type_cons{S,T<:Tuple}(::Type{S}, ::Type{T}) @@ -84,7 +100,7 @@ eltype{I,Z}(::Type{Zip{I,Z}}) = tuple_type_cons(eltype(I), eltype(Z)) end @inline done(z::Zip, st) = done(z.a,st[1]) | done(z.z,st[2]) -iteratorsize{I1,I2}(::Type{Zip{I1,I2}}) = and_iteratorsize(iteratorsize(I1),iteratorsize(I2)) +iteratorsize{I1,I2}(::Type{Zip{I1,I2}}) = zip_iteratorsize(iteratorsize(I1),iteratorsize(I2)) iteratoreltype{I1,I2}(::Type{Zip{I1,I2}}) = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2)) # filter @@ -141,7 +157,9 @@ done(i::Rest, st) = done(i.itr, st) eltype{I}(::Type{Rest{I}}) = eltype(I) iteratoreltype{I,S}(::Type{Rest{I,S}}) = iteratoreltype(I) -iteratorsize{T<:Rest}(::Type{T}) = SizeUnknown() +rest_iteratorsize(a) = SizeUnknown() +rest_iteratorsize(::IsInfinite) = IsInfinite() +iteratorsize{I,S}(::Type{Rest{I,S}}) = rest_iteratorsize(iteratorsize(I)) # Count -- infinite counting @@ -171,7 +189,10 @@ take(xs, n::Int) = Take(xs, n) eltype{I}(::Type{Take{I}}) = eltype(I) iteratoreltype{I}(::Type{Take{I}}) = iteratoreltype(I) -iteratorsize{T<:Take}(::Type{T}) = SizeUnknown() # TODO +take_iteratorsize(a) = HasLength() +take_iteratorsize(::SizeUnknown) = SizeUnknown() +iteratorsize{I}(::Type{Take{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)) @@ -196,7 +217,11 @@ drop(xs, n::Int) = Drop(xs, n) eltype{I}(::Type{Drop{I}}) = eltype(I) iteratoreltype{I}(::Type{Drop{I}}) = iteratoreltype(I) -iteratorsize{T<:Drop}(::Type{T}) = SizeUnknown() # TODO +drop_iteratorsize(::SizeUnknown) = SizeUnknown() +drop_iteratorsize(::Union{HasShape, HasLength}) = HasLength() +drop_iteratorsize(::IsInfinite) = IsInfinite() +iteratorsize{I}(::Type{Drop{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) @@ -327,6 +352,7 @@ iteratorsize{I1,I2}(::Type{Prod{I1,I2}}) = prod_iteratorsize(iteratorsize(I1),it end prod_iteratorsize(::Union{HasLength,HasShape}, ::Union{HasLength,HasShape}) = HasLength() +prod_iteratorsize(a, ::IsInfinite) = IsInfinite() # products can have an infinite last iterator (which moves slowest) prod_iteratorsize(a, b) = SizeUnknown() _size(p::Prod2) = (length(p.a), length(p.b)) diff --git a/test/functional.jl b/test/functional.jl index 826e4e8d2792c..4dc72d74aa9c9 100644 --- a/test/functional.jl +++ b/test/functional.jl @@ -76,6 +76,11 @@ let zeb = IOBuffer("1\n2\n3\n4\n5\n"), @test res == [(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd'), (5, 'e')] end +@test length(zip(cycle(1:3), 1:7)) == 7 +@test length(zip(cycle(1:3), 1:7, cycle(1:3))) == 7 +@test length(zip(1:3,Base.product(1:7,cycle(1:3)))) == 3 +@test length(zip(1:3,Base.product(1:7,cycle(1:3)),8)) == 1 + # rest # ---- let s = "hello" @@ -83,6 +88,8 @@ let s = "hello" @test collect(rest(s, st)) == ['e','l','l','o'] end +@test_throws MethodError collect(rest(countfrom(1), 5)) + # countfrom # --------- @@ -115,6 +122,10 @@ let i = 0 @test i == 10 end +@test length(take(1:3,typemax(Int))) == 3 +@test length(take(countfrom(1),3)) == 3 +@test length(take(1:6,3)) == 3 + # drop # ---- @@ -126,6 +137,10 @@ let i = 0 @test i == 4 end +@test length(drop(1:3,typemax(Int))) == 0 +@test Base.iteratorsize(drop(countfrom(1),3)) == Base.IsInfinite() +@test_throws MethodError length(drop(countfrom(1), 3)) + # cycle # ----- @@ -155,6 +170,7 @@ let i = 0 end end + # product # ------- @@ -165,6 +181,7 @@ end @test collect(Base.product(1:2,3:4)) == [(1,3),(2,3),(1,4),(2,4)] @test isempty(collect(Base.product(1:0,1:2))) @test length(Base.product(1:2,1:10,4:6)) == 60 +@test Base.iteratorsize(Base.product(1:2, countfrom(1))) == Base.IsInfinite() # foreach let