fix performance (#39)

* fix performance

* optimize collect

* fix for v1.11's cartesian logical arrays

* Update test/runtests.jl

src/InvertedIndices.jl

@@ -101,17 +101,49 @@ end
 InvertedIndexIterator(skips, picks) = InvertedIndexIterator{eltype(picks), typeof(skips), typeof(picks)}(skips, picks)
 Base.size(III::InvertedIndexIterator) = (length(III.picks) - length(III.skips),)
 
-@inline Base.iterate(I::InvertedIndexIterator) = iterate(I, (iterate(I.skips), iterate(I.picks)))
-Base.iterate(I::InvertedIndexIterator, ::Tuple{Any, Nothing}) = nothing
-@inline function Base.iterate(I::InvertedIndexIterator, (skipitr, pickitr))
+@inline function Base.iterate(I::InvertedIndexIterator)
+    skipitr = iterate(I.skips)
+    pickitr = iterate(I.picks)
+    pickitr === nothing && return nothing
+    while should_skip(skipitr, pickitr)
+        skipitr = iterate(I.skips, skipitr[2])
+        pickitr = iterate(I.picks, pickitr[2])
+        pickitr === nothing && return nothing
+    end
+    # This is a little silly, but splitting the tuple here allows inference to normalize
+    # Tuple{Union{Nothing, Tuple}, Tuple} to Union{Tuple{Nothing, Tuple}, Tuple{Tuple, Tuple}}
+    return skipitr === nothing ?
+            (pickitr[1], (nothing, pickitr[2])) :
+            (pickitr[1], (skipitr, pickitr[2]))
+end
+@inline function Base.iterate(I::InvertedIndexIterator, (_, pickstate)::Tuple{Nothing, Any})
+    pickitr = iterate(I.picks, pickstate)
+    pickitr === nothing && return nothing
+    return (pickitr[1], (nothing, pickitr[2]))
+end
+@inline function Base.iterate(I::InvertedIndexIterator, (skipitr, pickstate)::Tuple)
+    pickitr = iterate(I.picks, pickstate)
+    pickitr === nothing && return nothing
     while should_skip(skipitr, pickitr)
         skipitr = iterate(I.skips, tail(skipitr)...)
         pickitr = iterate(I.picks, tail(pickitr)...)
         pickitr === nothing && return nothing
     end
-    return (pickitr[1], (skipitr, iterate(I.picks, tail(pickitr)...)))
+    return skipitr === nothing ?
+            (pickitr[1], (nothing, pickitr[2])) :
+            (pickitr[1], (skipitr, pickitr[2]))
+end
+function Base.collect(III::InvertedIndexIterator{T}) where {T}
+    !isconcretetype(T) && return [i for i in III] # use widening if T is not concrete
+    v = Vector{T}(undef, length(III))
+    i = 0
+    for elt in III
+        i += 1
+        @inbounds v[i] = elt
+    end
+    i != length(v) && throw(AssertionError("length of inverted index does not match iterated count"))
+    return v
 end
-Base.collect(III::InvertedIndexIterator) = [i for i in III]
 
 should_skip(::Nothing, ::Any) = false
 should_skip(s::Tuple, p::Tuple) = _should_skip(s[1], p[1])

test/runtests.jl

@@ -185,3 +185,21 @@ end
     @test x isa InvertedIndex{InvertedIndices.NotMultiIndex}
     @test_throws ArgumentError v[x]
 end
+
+returns(val) = _->val
+@testset "type stability" begin
+    for arr in (
+            1:5,
+            reshape(1:5*3, 5, 3),
+            reshape(1:5*3*7, 5, 3, 7),
+            reshape(1:5*3*7*11, 5, 3, 7, 11),
+        )
+        I = to_indices(arr, (Not(iseven.(arr)),))[1]
+        @test all(isodd, arr[I])
+        @test all(isodd, LinearIndices(arr)[I])
+        @test all(isodd, LinearIndices(arr)[collect(I)])
+        @allocated(foreach(returns(nothing), I))
+        @test @allocated(foreach(returns(nothing), I)) == 0
+        @test @inferred(LinearIndices(arr)[collect(I)]) == vec(filter(!iseven, arr))
+    end
+end