From 630fa67ee1d5d848b936f389110a285c8b296e51 Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Thu, 10 Mar 2016 16:01:13 +0100
Subject: [PATCH] length() of Take and Drop iterators

---
 base/iterator.jl   | 26 ++++++++++++++++++++------
 test/functional.jl |  7 +++++++
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/base/iterator.jl b/base/iterator.jl
index 5523f7b898a2b0..121c6de073c697 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(b) # inherit behaviour, error
+_diff_length(a, b, A, B) = max(length(a)-length(b), 0) 
+
 # enumerate
 
 immutable Enumerate{I}
@@ -31,10 +41,7 @@ zip_iteratorsize(a, b) = and_iteratorsize(a,b) # as `and_iteratorsize` but inher
 zip_iteratorsize(::HasLength, ::IsInfinite) = HasLength()
 zip_iteratorsize(::HasShape, ::IsInfinite) = HasLength()
 zip_iteratorsize(a::IsInfinite, b) = zip_iteratorsize(b,a)
-_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))
+
 
 immutable Zip1{I} <: AbstractZipIterator
     a::I
@@ -182,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))
 
@@ -207,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)
diff --git a/test/functional.jl b/test/functional.jl
index 842163cacd695b..ca1e2dbcbb74ad 100644
--- a/test/functional.jl
+++ b/test/functional.jl
@@ -122,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
 # ----
 
@@ -133,6 +137,9 @@ let i = 0
     @test i == 4
 end
 
+@test length(drop(1:3,typemax(Int))) == 0
+@test Base.iteratorsize(drop(countfrom(1),3)) == Base.IsInfinite()
+
 # cycle
 # -----