Faster iterate method for TrivialInterlacer (#503)

* Faster iterate method for TrivialInterlacer

* Version bump to 0.8.43

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "ApproxFunBase"
 uuid = "fbd15aa5-315a-5a7d-a8a4-24992e37be05"
-version = "0.8.42"
+version = "0.8.43"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/LinearAlgebra/helper.jl

@@ -686,7 +686,7 @@ struct BlockInterlacer{DMS<:Tuple{Vararg{AbstractVector{Int}}}}
 end
 
 
-const TrivialInterlacer{d} = BlockInterlacer{<:NTuple{d,Ones{Int}}}
+const TrivialInterlacer{d,Ax} = BlockInterlacer{<:NTuple{d,Ones{Int,1,Tuple{Ax}}}}
 
 BlockInterlacer(v::AbstractVector) = BlockInterlacer(Tuple(v))
 
@@ -732,23 +732,18 @@ function iterate(it::BlockInterlacer, (N,k,blkst,lngs))
         return iterate(it,(1,1,blkst,lngs))
     end
 
-    Bnxtb = iterate(it.blocks[N],blkst[N]...)  # B is block size
+    Bnxtb = iterate(it.blocks[N],blkst[N]...)
 
-    if Bnxtb === nothing
+    if Bnxtb === nothing || k > Bnxtb[1]
         # increment to next N
         return iterate(it,(N+1,1,blkst,lngs))
     end
 
-    B,nxtb = Bnxtb
-
-    if k > B
-        #increment to next N
-        return iterate(it,(N+1,1,blkst,lngs))
-    end
-
-
     lngs = Base.setindex(lngs, lngs[N]+1, N)
     return (N,lngs[N]),(N,k+1,blkst,lngs)
 end
 
+iterate(it::TrivialInterlacer{N,OneToInf{Int}}, st...) where {N} =
+    iterate(Iterators.product(1:N, axes(it.blocks[1],1)), st...)
+
 cache(Q::BlockInterlacer) = CachedIterator(Q)

test/runtests.jl

@@ -55,6 +55,22 @@ end
 
         @test collect(ApproxFunBase.BlockInterlacer(([2],[2],[2]))) ==
             [(1,1),(1,2),(2,1),(2,2),(3,1),(3,2)]
+
+        @testset "TrivialInterlacer" begin
+            o = Ones{Int}(ℵ₀)
+            f = Fill(1, ℵ₀)
+            function test_nD(nD)
+                B1 = ApproxFunBase.BlockInterlacer(ntuple(_->o, nD))
+                B2 = ApproxFunBase.BlockInterlacer(ntuple(_->f, nD))
+                it = Iterators.take(zip(B1, B2), 100_000)
+                for (a,b) in it
+                    @test a == b
+                end
+            end
+            test_nD(1)
+            test_nD(2)
+            test_nD(3)
+        end
     end
 
     @testset "issue #94" begin