Specialization for faster count(::Array{Bool}) (JuliaLang#34060)

N5N3 · Jan 28, 2020 · 03804d6 · 03804d6
1 parent baaf63e
commit 03804d6
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/base/reduce.jl b/base/reduce.jl
@@ -861,3 +861,19 @@ function count(pred, a::AbstractArray)
  return n
 end
 count(itr) = count(identity, itr)
+
+function count(::typeof(identity), x::Array{Bool})
+ n = 0
+ chunks = length(x) ÷ sizeof(UInt)
+ mask = 0x0101010101010101 % UInt
+ GC.@preserve x begin
+ ptr = Ptr{UInt}(pointer(x))
+ for i in 1:chunks
+ n += count_ones(unsafe_load(ptr, i) & mask)
+ end
+ end
+ for i in sizeof(UInt)*chunks+1:length(x)
+ n += x[i]
+ end
+ return n
+end
diff --git a/test/reduce.jl b/test/reduce.jl
@@ -445,6 +445,11 @@ struct SomeFunctor end
 @test count(x->x>0, Int[]) == count(Bool[]) == 0
 @test count(x->x>0, -3:5) == count((-3:5) .> 0) == 5
 @test count([true, true, false, true]) == count(BitVector([true, true, false, true])) == 3
+let x = repeat([false, true, false, true, true, false], 7)
+ @test count(x) == 21
+ GC.@preserve x (unsafe_store!(Ptr{UInt8}(pointer(x)), 0xfe, 3))
+ @test count(x) == 21
+end
 @test_throws TypeError count(sqrt, [1])
 @test_throws TypeError count([1])
 let itr = (x for x in 1:10 if x < 7)