Explicit loop in _findall to avoid allocations (#2771)

JuliaData · May 25, 2021 · 8ce4d2a · 8ce4d2a
1 parent e230855
commit 8ce4d2a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/src/other/utils.jl b/src/other/utils.jl
@@ -218,8 +218,10 @@ function _findall(B::BitVector)::Union{UnitRange{Int}, Vector{Int}}
         end
         if c == ~UInt64(0)
             if stop != -1
-                I = Vector{Int}(undef,nnzB)
-                I[1:i-1] .= start:stop
+                I = Vector{Int}(undef, nnzB)
+                for j in 1:i-1
+                    I[j] = start + j - 1
+                end
                 break
             end
             if start == -1
@@ -243,20 +245,26 @@ function _findall(B::BitVector)::Union{UnitRange{Int}, Vector{Int}}
             lz = leading_zeros(c)
             co = c >> tz == (one(UInt64) << (64 - lz - tz)) - 1 # block of countinous ones in c
             if stop != -1  # already found block of ones and zeros, just not materialized
-                I = Vector{Int}(undef,nnzB)
-                I[1:i-1] .= start:stop
+                I = Vector{Int}(undef, nnzB)
+                for j in 1:i-1
+                    I[j] = start + j - 1
+                end
                 break
             elseif !co # not countinous ones
-                I = Vector{Int}(undef,nnzB)
+                I = Vector{Int}(undef, nnzB)
                 if start != -1
-                    I[1:i-1] .= start:start + i - 2
+                    for j in 1:i-1
+                        I[j] = start + j - 1
+                    end
                 end
                 break
             else # countinous block of ones
                 if start != -1
                     if tz > 0 # like __1111__ or 111111__
-                        I = Vector{Int}(undef,nnzB)
-                        I[1:i-1] .= start:start + i - 2
+                        I = Vector{Int}(undef, nnzB)
+                        for j in 1:i-1
+                            I[j] = start + j - 1
+                        end
                         break
                     else # lz > 0, like __111111
                         stop = i1 + (64 - lz) - 1
@@ -310,7 +318,9 @@ function _findall(B::BitVector)::Union{UnitRange{Int}, Vector{Int}}
         end
 
         while c == ~UInt64(0)
-            I[i:i+64-1] .= i1:(i1+64-1)
+            for j in 0:64-1
+                I[i + j] = i1 + j
+            end
             i += 64
             if Bi == length(Bc)
                 @assert nnzB == i - 1

diff --git a/test/utils.jl b/test/utils.jl
@@ -167,7 +167,7 @@ end
         "TFT small" => ([trues(85); falses(100); trues(85)], Vector{Int}),
         "FTFFT small" => ([falses(64 + 32); trues(32); falses(128); trues(32)], Vector{Int}),
         "TFTF small" => ([falses(64); trues(64); falses(64); trues(64)], Vector{Int}),
-        "TFT small" => ([trues(64); falses(10); trues(100)], Vector{Int}),
+        "TFT small2" => ([trues(64); falses(10); trues(100)], Vector{Int}),
 
         "FTF Big" => ([falses(8500); trues(100000); falses(65000)], UnitRange{Int}),
         "TFT Big" => ([trues(8500); falses(100000); trues(65000)], Vector{Int}),