IteratorSize for Tensorizer and BlockInterlacer (#365)

* IteratorSize for Tensorizer and BlockInterlacer

* keys and IteratorSize for CachedIterator

Author: jishnub

src/ApproxFunBase.jl

@@ -135,6 +135,11 @@ function assert_integer(k::Number)
     return nothing
 end
 
+function _IteratorSize(::Type{T}) where {T<:Tuple}
+    s = ntuple(i-> Base.IteratorSize(fieldtype(T, i)), fieldcount(T))
+    any(x -> x isa Base.IsInfinite, s) ? Base.IsInfinite() : Base.HasLength()
+end
+
 include("LinearAlgebra/LinearAlgebra.jl")
 include("Fun.jl")
 include("onehotvector.jl")

src/LinearAlgebra/helper.jl

@@ -571,6 +571,10 @@ end
 
 eltype(it::Type{<:CachedIterator{T}}) where {T} = T
 
+Base.IteratorSize(::Type{<:CachedIterator{<:Any,IT}}) where {IT} = Base.IteratorSize(IT)
+
+Base.keys(c::CachedIterator) = keys(c.iterator)
+
 iterate(it::CachedIterator) = iterate(it,1)
 function iterate(it::CachedIterator,st::Int)
     if  st == it.length + 1 && iterate(it.iterator,it.state...) === nothing
@@ -702,8 +706,9 @@ eltype(::Type{<:BlockInterlacer}) = Tuple{Int,Int}
 
 dimensions(b::BlockInterlacer) = map(sum,b.blocks)
 dimension(b::BlockInterlacer,k) = sum(b.blocks[k])
-Base.length(b::BlockInterlacer) = mapreduce(sum,+,b.blocks)
+length(b::BlockInterlacer) = mapreduce(sum,+,b.blocks)
 
+Base.IteratorSize(::Type{BlockInterlacer{T}}) where {T} = _IteratorSize(T)
 
 # the state is always (whichblock,curblock,cursubblock,curcoefficients)
 # start(it::BlockInterlacer) = (1,1,map(start,it.blocks),ntuple(zero,length(it.blocks)))

src/Multivariate/TensorSpace.jl

@@ -28,12 +28,15 @@ struct Tensorizer{DMS<:Tuple}
     blocks::DMS
 end
 
+const InfOnes = Ones{Int,1,Tuple{OneToInf{Int}}}
 const Tensorizer2D{AA, BB} = Tensorizer{Tuple{AA, BB}}
-const TrivialTensorizer{d} = Tensorizer{NTuple{d,Ones{Int,1,Tuple{OneToInf{Int}}}}}
+const TrivialTensorizer{d} = Tensorizer{NTuple{d,InfOnes}}
 
 eltype(::Type{<:Tensorizer{<:Tuple{Vararg{Any,N}}}}) where {N} = NTuple{N,Int}
 dimensions(a::Tensorizer) = map(sum,a.blocks)
-Base.length(a::Tensorizer) = reduce(*, dimensions(a)) # easier type-inference than mapreduce
+length(a::Tensorizer) = reduce(*, dimensions(a)) # easier type-inference than mapreduce
+
+Base.IteratorSize(::Type{Tensorizer{T}}) where {T<:Tuple} = _IteratorSize(T)
 
 Base.keys(a::Tensorizer) = oneto(length(a))
 
@@ -129,7 +132,9 @@ end
 
 cache(a::Tensorizer) = CachedIterator(a)
 
-function Base.findfirst(::TrivialTensorizer{2},kj::Tuple{Int,Int})
+@deprecate findfirst(t::Tensorizer, kj::NTuple{2,Int}) findfirst(kj, t)
+
+function Base.findfirst(kj::NTuple{2,Int}, ::TrivialTensorizer{2})
     k,j=kj
     if k > 0 && j > 0
         n=k+j-2
@@ -138,7 +143,7 @@ function Base.findfirst(::TrivialTensorizer{2},kj::Tuple{Int,Int})
         nothing
     end
 end
-function Base.findfirst(sp::Tensorizer{<:NTuple{2,Ones{Int}}}, kj::NTuple{2,Int})
+function Base.findfirst(kj::NTuple{2,Int}, sp::Tensorizer{<:NTuple{2,Ones{Int}}})
     k,j=kj
 
     len1, len2 = length(sp.blocks[1]), length(sp.blocks[2])
@@ -212,9 +217,9 @@ blocklengths(::TrivialTensorizer{2}) = 1:∞
 blocklengths(it::Tensorizer) = tensorblocklengths(it.blocks...)
 blocklengths(it::CachedIterator) = blocklengths(it.iterator)
 
-function getindex(it::TrivialTensorizer{2},n::Integer)
+function getindex(it::TrivialTensorizer{2}, n::Integer)
     m=Int(block(it,n))
-    p=findfirst(it,(1,m))
+    p=findfirst((1,m), it)
     j=1+n-p
     j,m-j+1
 end

test/runtests.jl

@@ -511,6 +511,15 @@ end
             @test findfirst(t, vi) == i
         end
     end
+    @testset "cache" begin
+        ax = Ones{Int}(4);
+        t = ApproxFunBase.Tensorizer((ax,ax))
+        c = ApproxFunBase.cache(t)
+        v = collect(c)
+        @testset for i in eachindex(c)
+            @test c[i] == v[i]
+        end
+    end
 end
 
 @time include("ETDRK4Test.jl")