Add missing optimization for `iterate(::LogicalIndex{<:CartesianIndex…

…,<:BitArray})` (#43448) * make `collect(::LogicalIndex)` faster by call `findall` * borrow code from findall borrow code from findall * Update base/multidimensional.jl Co-authored-by: Kristoffer Carlsson <[email protected]> Co-authored-by: Kristoffer Carlsson <[email protected]>
JuliaLang · Feb 16, 2022 · d85108a · d85108a
1 parent 1a3da30
commit d85108a
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -781,22 +781,28 @@ end
  end
 end
 # When wrapping a BitArray, lean heavily upon its internals.
-@inline function iterate(L::Base.LogicalIndex{Int,<:BitArray})
+@inline function iterate(L::LogicalIndex{Int,<:BitArray})
  L.sum == 0 && return nothing
  Bc = L.mask.chunks
- return iterate(L, (1, @inbounds Bc[1]))
+ return iterate(L, (1, 1, (), @inbounds Bc[1]))
 end
-@inline function iterate(L::Base.LogicalIndex{Int,<:BitArray}, s)
+@inline function iterate(L::LogicalIndex{<:CartesianIndex,<:BitArray})
+ L.sum == 0 && return nothing
+ Bc = L.mask.chunks
+ irest = ntuple(one, ndims(L.mask)-1)
+ return iterate(L, (1, 1, irest, @inbounds Bc[1]))
+end
+@inline function iterate(L::LogicalIndex{<:Any,<:BitArray}, (i1, Bi, irest, c))
  Bc = L.mask.chunks
- i1, c = s
- while c==0
- i1 % UInt >= length(Bc) % UInt && return nothing
- i1 += 1
- @inbounds c = Bc[i1]
+ while c == 0
+ Bi >= length(Bc) && return nothing
+ i1 += 64
+ @inbounds c = Bc[Bi+=1]
  end
- tz = trailing_zeros(c) + 1
+ tz = trailing_zeros(c)
  c = _blsr(c)
- return ((i1-1)<<6 + tz, (i1, c))
+ i1, irest = _overflowind(i1 + tz, irest, size(L.mask))
+ return eltype(L)(i1, irest...), (i1 - tz, Bi, irest, c)
 end
 
 @inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex{<:Any,<:AbstractArray{Bool,1}}) =