From 2de628fcbdb79772d8eda9b8a2d70a2d00c90df0 Mon Sep 17 00:00:00 2001
From: chethega <chethega@gmail.com>
Date: Wed, 24 Oct 2018 12:27:09 +0200
Subject: [PATCH] speed up logical indexing by bitarray (#29746)

* speed up logical indexing by bitarray

* changed inbounds / OOB behavior to match array iterators

switched spelling of _blsr

* Update base/multidimensional.jl

Co-Authored-By: chethega <chethega@gmail.com>
(cherry picked from commit 44f2563d11f3a6fa56e9a059a8349df53cdc8eb8)
---
 base/bitarray.jl         |  1 +
 base/multidimensional.jl | 26 +++++++++++++++-----------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/base/bitarray.jl b/base/bitarray.jl
index 526f027c6aa9b5..aea54aa3385bfd 100644
--- a/base/bitarray.jl
+++ b/base/bitarray.jl
@@ -84,6 +84,7 @@ IndexStyle(::Type{<:BitArray}) = IndexLinear()
 const _msk64 = ~UInt64(0)
 @inline _div64(l) = l >> 6
 @inline _mod64(l) = l & 63
+@inline _blsr(x)= x & (x-1) #zeros the last set bit. Has native instruction on many archs. needed in multidimensional.jl
 @inline _msk_end(l::Integer) = _msk64 >>> _mod64(-l)
 @inline _msk_end(B::BitArray) = _msk_end(length(B))
 num_bit_chunks(n::Int) = _div64(n+63)
diff --git a/base/multidimensional.jl b/base/multidimensional.jl
index 1a78b84c1e6a31..a5d2dc30692e66 100644
--- a/base/multidimensional.jl
+++ b/base/multidimensional.jl
@@ -520,19 +520,23 @@ end
         L.mask[idx] && return (idx, s)
     end
 end
-# When wrapping a BitArray, lean heavily upon its internals -- this is a common
-# case. Just use the Int index and count as its state.
-@inline function iterate(L::LogicalIndex{Int,<:BitArray}, s=(0,1))
-    s[2] > length(L) && return nothing
-    i, n = s
+# When wrapping a BitArray, lean heavily upon its internals.
+@inline function iterate(L::Base.LogicalIndex{Int,<:BitArray})
+    L.sum == 0 && return nothing
     Bc = L.mask.chunks
-    while true
-        if Bc[_div64(i)+1] & (UInt64(1)<<_mod64(i)) != 0
-            i += 1
-            return (i, (i, n+1))
-        end
-        i += 1
+    return iterate(L, (1, @inbounds Bc[1]))
+end
+@inline function iterate(L::Base.LogicalIndex{Int,<:BitArray}, s)
+    Bc = L.mask.chunks
+    i1, c = s
+    while c==0
+        i1 % UInt >= length(Bc) % UInt && return nothing
+        i1 += 1
+        @inbounds c = Bc[i1]
     end
+    tz = trailing_zeros(c) + 1
+    c = _blsr(c)
+    return ((i1-1)<<6 + tz, (i1, c))
 end
 
 @inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex{<:Any,<:AbstractArray{Bool,1}}) =