Improve bit-twiddling for counting (#335)

This adds a more generic encoding counting function and uses it.

Author: jakobnissen

Project.toml

@@ -1,7 +1,7 @@
 name = "BioSequences"
 uuid = "7e6ae17a-c86d-528c-b3b9-7f778a29fe59"
 authors = ["Sabrina Jaye Ward <sabrinajward@protonmail.com>", "Jakob Nissen <jakobnybonissen@gmail.com>"]
-version = "3.4.1"
+version = "3.4.2"
 
 [deps]
 BioSymbols = "3c28c6f8-a34d-59c4-9654-267d177fcfa9"

src/bit-manipulation/bit-manipulation.jl

@@ -68,6 +68,26 @@ end
     @inline sum(map(i -> f(i.value), z))
 end
 
+pattern(::BitsPerSymbol{1})   = typemax(UInt128)
+pattern(::BitsPerSymbol{2})   = 0x55555555555555555555555555555555
+pattern(::BitsPerSymbol{4})   = 0x11111111111111111111111111111111
+pattern(::BitsPerSymbol{8})   = 0x01010101010101010101010101010101
+pattern(::BitsPerSymbol{16})  = 0x00010001000100010001000100010001
+pattern(::BitsPerSymbol{32})  = 0x00000001000000010000000100000001
+pattern(::BitsPerSymbol{64})  = 0x00000000000000010000000000000001
+pattern(::BitsPerSymbol{128}) = 0x00000000000000000000000000000001
+
+function count_encoding(chunk::T, encoding::T, b::BitsPerSymbol{B}) where {T <: Unsigned, B}
+    pat = pattern(b) % typeof(encoding)
+    u = chunk ⊻ (encoding * pat)
+    for i in 1:trailing_zeros(B)
+        shift = (1 << (i - 1)) & (8*sizeof(T) - 1)
+        u |= (u >> shift)
+    end
+    u = ~u & pat
+    return count_ones(u)
+end
+
 @inline function certain_bitcount(x::UInt64, ::T) where {T<:NucleicAcidAlphabet{4}}
     x = enumerate_nibbles(x)
     x = x ⊻ 0x1111111111111111

src/counting.jl

@@ -21,7 +21,6 @@ end
         y += @inline body(data[i])
     end
     y
-    y
 end
 
 """
@@ -356,44 +355,22 @@ function count_symbol(seq::BioSequence, sym::BioSymbol)
     n
 end
 
-function Base.count(
-    pred::Base.Fix2{<:Union{typeof(==), typeof(isequal)}, <: BioSymbol},
-    s::BioSequence
-)
-    count_symbol(s, pred.x)
-end
-
-function count_symbol(seq::FourBit, s::Union{RNA, DNA})
-    pattern = encode(Alphabet(seq), s)::UInt64 * 0x1111111111111111
+function count_symbol(seq::Union{LongSubSeq, LongSequence}, sym::BioSymbol)
+    enc = encode(Alphabet(seq), sym)
     tail = (chunk, rm) -> begin
         mask = UInt64(1) << (rm & 63) - 1
-        masked = iszero(pattern) ? chunk | ~mask : mask & chunk
-        count_0000_nibbles(masked ⊻ pattern)
+        masked = iszero(enc) ? chunk | ~mask : mask & chunk
+        count_encoding(masked, enc, BitsPerSymbol(seq))
     end
-    body = i -> count_0000_nibbles(i ⊻ pattern)
+    body = i -> count_encoding(i, enc, BitsPerSymbol(seq))
     counter_1seq(tail, body, seq)
 end
 
-function count_symbol(seq::TwoBit, s::Union{RNA, DNA})
-    pattern = encode(Alphabet(seq), s)::UInt64 * 0x5555555555555555
-    tail = (chunk, rm) -> begin
-        mask = UInt64(1) << (rm & 63) - 1
-        masked = iszero(pattern) ? chunk | ~mask : mask & chunk
-        count_00_bitpairs(masked ⊻ pattern)
-    end
-    body = i -> count_00_bitpairs(i ⊻ pattern)
-    counter_1seq(tail, body, seq)
-end
-
-function count_symbol(seq::SeqOrView{AminoAcidAlphabet}, s::AminoAcid)
-    byte = encode(Alphabet(seq), s) % UInt8
-    tail = (chunk, rm) -> begin
-        mask = UInt64(1) << (rm & 63) - 1
-        masked = iszero(byte) ? chunk | ~mask : mask & chunk
-        count_compared_bytes(==(byte), masked)
-    end
-    body = i -> count_compared_bytes(==(byte), i)
-    counter_1seq(tail, body, seq)
+function Base.count(
+    pred::Base.Fix2{<:Union{typeof(==), typeof(isequal)}, <: BioSymbol},
+    s::BioSequence
+)
+    count_symbol(s, pred.x)
 end
 
 ## Deprecate weird two-arg methods