add init argument to count

Author: simeonschaub

base/bitarray.jl

@@ -1386,15 +1386,15 @@ circshift!(B::BitVector, i::Integer) = circshift!(B, B, i)
 
 ## count & find ##
 
-function bitcount(Bc::Vector{UInt64})
-    n = 0
+function bitcount(Bc::Vector{UInt64}; init::T=0) where {T}
+    n::T = init
     @inbounds for i = 1:length(Bc)
-        n += count_ones(Bc[i])
+        n += count_ones(Bc[i]) % T
     end
     return n
 end
 
-count(B::BitArray) = bitcount(B.chunks)
+count(B::BitArray; init=0) = bitcount(B.chunks; init)
 
 function unsafe_bitfindnext(Bc::Vector{UInt64}, start::Int)
     chunk_start = _div64(start-1)+1

base/reduce.jl

@@ -938,11 +938,10 @@ end
 _bool(f) = x->f(x)::Bool
 
 """
-    count(p, itr) -> Integer
-    count(itr) -> Integer
+    count([f=identity,] itr; init=0) -> Integer
 
-Count the number of elements in `itr` for which predicate `p` returns `true`.
-If `p` is omitted, counts the number of `true` elements in `itr` (which
+Count the number of elements in `itr` for which the function `f` returns `true`.
+If `f` is omitted, counts the number of `true` elements in `itr` (which
 should be a collection of boolean values).
 
 # Examples
@@ -954,30 +953,30 @@ julia> count([true, false, true, true])
 3
 ```
 """
-count(itr) = count(identity, itr)
+count(itr; init=0) = count(identity, itr; init)
 
-count(f, itr) = _simple_count(f, itr)
+count(f, itr; init=0) = _simple_count(f, itr, init)
 
-function _simple_count(pred, itr)
-    n = 0
+function _simple_count(pred, itr, init)
+    n = init
     for x in itr
         n += pred(x)::Bool
     end
     return n
 end
 
-function count(::typeof(identity), x::Array{Bool})
-    n = 0
+function _simple_count(::typeof(identity), x::Array{Bool}, init::T=0) where {T}
+    n::T = init
     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)
+            n += count_ones(unsafe_load(ptr, i) & mask) % T
         end
     end
     for i in sizeof(UInt)*chunks+1:length(x)
-        n += x[i]
+        n += x[i] % T
     end
     return n
 end

base/reducedim.jl

@@ -386,11 +386,11 @@ julia> count(<=(2), A, dims=2)
  0
 ```
 """
-count(A::AbstractArrayOrBroadcasted; dims=:) = count(identity, A, dims=dims)
-count(f, A::AbstractArrayOrBroadcasted; dims=:) = _count(f, A, dims)
+count(A::AbstractArrayOrBroadcasted; dims=:, init=0) = count(identity, A; dims, init)
+count(f, A::AbstractArrayOrBroadcasted; dims=:, init=0) = _count(f, A, dims, init)
 
-_count(f, A::AbstractArrayOrBroadcasted, dims::Colon) = _simple_count(f, A)
-_count(f, A::AbstractArrayOrBroadcasted, dims) = mapreduce(_bool(f), add_sum, A, dims=dims, init=0)
+_count(f, A::AbstractArrayOrBroadcasted, dims::Colon, init) = _simple_count(f, A, init)
+_count(f, A::AbstractArrayOrBroadcasted, dims, init) = mapreduce(_bool(f), add_sum, A; dims, init)
 
 """
     count!([f=identity,] r, A)

test/bitarray.jl

@@ -1219,6 +1219,8 @@ timesofar("datamove")
         @check_bit_operation findall(falses(t)) ret_type
         @check_bit_operation findall(bitrand(t)) ret_type
     end
+
+    @test count(trues(2, 2), init=0x00) isa UInt8
 end
 
 timesofar("find")

test/reduce.jl

@@ -520,6 +520,12 @@ struct NonFunctionIsZero end
 @test count(NonFunctionIsZero(), [0]) == 1
 @test count(NonFunctionIsZero(), [1]) == 0
 
+@test count(Iterators.repeated(true, 3), init=0x00) isa UInt8
+@test count(!=(2), Iterators.take(1:7, 3), init=Int32(0)) isa Int32
+@test count(identity, [true, false], init=Int8(0)) isa Int8
+@test count(!, [true false; false true], dims=:, init=Int16(0)) isa Int16
+@test count(identity, [true false; false true], dims=2, init=UInt(0)) isa Matrix{UInt}
+
 ## cumsum, cummin, cummax
 
 z = rand(10^6)

test/reducedim.jl

@@ -77,6 +77,9 @@ safe_minabs(A::Array{T}, region) where {T} = safe_mapslices(minimum, abs.(A), re
     @test @inferred(maximum(abs, Areduc, dims=region)) ≈ safe_maxabs(Areduc, region)
     @test @inferred(minimum(abs, Areduc, dims=region)) ≈ safe_minabs(Areduc, region)
     @test @inferred(count(!, Breduc, dims=region)) ≈ safe_count(.!Breduc, region)
+
+    @test eltype(@inferred(count(Breduc, dims=region, init=0x00))) == UInt8
+    @test eltype(@inferred(count(!, Breduc, dims=region, init=Int16(0)))) == Int16
 end
 
 # Combining dims and init
@@ -446,3 +449,6 @@ end
     @test_throws TypeError count([1], dims=1)
     @test_throws TypeError count!([1], [1])
 end
+
+@test @inferred(count(false:true, dims=:, init=0x0000)) isa UInt16
+@test @inferred(count(isodd, reshape(1:9, 3, 3), dims=:, init=Int128(0))) isa Int128