Merge pull request #15400 from mschauer/isinfinite2

zips, products and takes with infinite components

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,15 +76,15 @@ 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
     a::I
     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))

test/functional.jl

@@ -76,13 +76,20 @@ 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"
     _, st = next(s, start(s))
     @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