Add more optimised find methods for BioSequences (#331)

This adds SIMD optimised methods for `findprev` and `findnext` (and therefore, by extension `findfirst` and `findlast`) for the predicates `!isgap`, `iscertain`, `!iscertain`, `isambiguous` and `!isambiguous` for the default alphabets for LongSequence and LongSubSeq

The purpose of these is to find e.g. regions of N's in a sequence, or find the first symbol in a long DNA sequence which is not a gap or an unambiguous symbol.

Some benchmarks for AA - nucleotide benchmarks are even better:
```julia
s = randaaseq(100_000)
@Btime findfirst(isgap, s)
@Btime findfirst(isambiguous, s)
@Btime findfirst(!iscertain, s)
for i in eachindex(s)
   s[i] = AA_B
end
@Btime findfirst(!isambiguous, s)
@Btime findfirst(iscertain, s)
```
Before:
  1.738 μs (0 allocations: 0 bytes)
  39.195 μs (0 allocations: 0 bytes)
  39.195 μs (0 allocations: 0 bytes)
  39.205 μs (0 allocations: 0 bytes)
  39.205 μs (0 allocations: 0 bytes)

After:
  1.894 μs (0 allocations: 0 bytes)
  3.270 μs (0 allocations: 0 bytes)
  5.500 μs (0 allocations: 0 bytes)
  3.240 μs (0 allocations: 0 bytes)
  4.530 μs (0 allocations: 0 bytes)

Author: jakobnissen

src/BioSequences.jl

@@ -219,6 +219,7 @@ include("biosequence/biosequence.jl")
 include("longsequences/longsequence.jl")
 include("longsequences/hash.jl")
 include("longsequences/randseq.jl")
+include("longsequences/find.jl")
 
 include("counting.jl")
 

src/biosequence/find.jl

@@ -7,9 +7,24 @@
 ### This file is a part of BioJulia.
 ### License is MIT: https://github.com/BioJulia/BioSequences.jl/blob/master/LICENSE.md
 
+Base.findfirst(f::Function, seq::BioSequence) = findnext(f, seq, firstindex(seq))
+Base.findlast(f::Function, seq::BioSequence) = findprev(f, seq, lastindex(seq))
+
 function Base.findnext(f::Function, seq::BioSequence, start::Integer)
+    start = Int(start)::Int
     start > lastindex(seq) && return nothing
-    checkbounds(seq, start)
+    @boundscheck (start < 1 && throw(BoundsError(seq, start)))
+    @inline _findnext(f, seq, start)
+end
+
+function Base.findprev(f::Function, seq::BioSequence, start::Integer)
+    start = Int(start)::Int
+    start < 1 && return nothing
+    @boundscheck (start > lastindex(seq) && throw(BoundsError(seq, start)))
+    @inline _findprev(f, seq, start)
+end
+
+function _findnext(f::Function, seq::BioSequence, start::Int)
     @inbounds for i in start:lastindex(seq)
         if f(seq[i])
             return i
@@ -18,25 +33,31 @@ function Base.findnext(f::Function, seq::BioSequence, start::Integer)
     return nothing
 end
 
-# No ambiguous sites can exist in a nucleic acid sequence using the two-bit alphabet.
-Base.findnext(::typeof(isambiguous), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Integer) = nothing
-
-function Base.findprev(f::Function, seq::BioSequence, start::Integer)
-    start < firstindex(seq) && return nothing
-    checkbounds(seq, start)
-    for i in start:-1:firstindex(seq)
+function _findprev(f::Function, seq::BioSequence, start::Int)
+    @inbounds for i in start:-1:firstindex(seq)
         if f(seq[i])
             return i
         end
     end
     return nothing
 end
 
-Base.findfirst(f::Function, seq::BioSequence) = findnext(f, seq, firstindex(seq))
-Base.findlast(f::Function, seq::BioSequence) = findprev(f, seq, lastindex(seq))
+# For two-bit alphabet
+# N.B: The isgap and isambiguous have identical definitions but are separate methods to avoid
+# ambiguity errors.
+_findnext(::typeof(isambiguous), ::BioSequence{<:NucleicAcidAlphabet{2}}, ::Int) = nothing
+_findprev(::typeof(isambiguous), ::BioSequence{<:NucleicAcidAlphabet{2}}, ::Int) = nothing
+_findnext(::typeof(isgap), ::BioSequence{<:NucleicAcidAlphabet{2}}, ::Int) = nothing
+_findprev(::typeof(isgap), ::BioSequence{<:NucleicAcidAlphabet{2}}, ::Int) = nothing
 
-# Finding specific symbols
+_findnext(::typeof(iscertain), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Int) = from
+_findprev(::typeof(iscertain), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Int) = from
+_findnext(::typeof(!isambiguous), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Int) = from
+_findprev(::typeof(!isambiguous), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Int) = from
+_findnext(::typeof(!isgap), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Int) = from
+_findprev(::typeof(!isgap), seq::BioSequence{<:NucleicAcidAlphabet{2}}, from::Int) = from
 
+# Finding specific symbols
 Base.findnext(x::DNA, seq::BioSequence{<:DNAAlphabet}, start::Integer) = findnext(isequal(x), seq, start)
 Base.findnext(x::RNA, seq::BioSequence{<:RNAAlphabet}, start::Integer) = findnext(isequal(x), seq, start)
 Base.findnext(x::AminoAcid, seq::BioSequence{AminoAcidAlphabet}, start::Integer) = findnext(isequal(x), seq, start)
@@ -51,4 +72,16 @@ Base.findfirst(x::AminoAcid, seq::BioSequence{AminoAcidAlphabet}) = findfirst(is
 
 Base.findlast(x::DNA, seq::BioSequence{<:DNAAlphabet}) = findlast(isequal(x), seq)
 Base.findlast(x::RNA, seq::BioSequence{<:RNAAlphabet}) = findlast(isequal(x), seq)
-Base.findlast(x::AminoAcid, seq::BioSequence{AminoAcidAlphabet}) = findlast(isequal(x), seq)
+Base.findlast(x::AminoAcid, seq::BioSequence{AminoAcidAlphabet}) = findlast(isequal(x), seq)
+
+# Finding gaps
+_findnext(::typeof(isgap), seq::BioSequence, i::Int) = _findnext(==(gap(eltype(seq))), seq, i)
+_findprev(::typeof(isgap), seq::BioSequence, i::Int) = _findprev(==(gap(eltype(seq))), seq, i)
+
+function _findnext(::ComposedFunction{typeof(!), typeof(isgap)}, seq::BioSequence, i::Int)
+    _findnext(!=(gap(eltype(seq))), seq, i)
+end
+
+function _findprev(::ComposedFunction{typeof(!), typeof(isgap)}, seq::BioSequence, i::Int)
+    _findprev(!=(gap(eltype(seq))), seq, i)
+end

src/counting.jl

@@ -259,12 +259,6 @@ _n_certain(seq::BioSequence) = count_naive(iscertain, seq)
 
 _n_gc(seq::BioSequence) = count_naive(isGC, seq)
 
-const FourBit = SeqOrView{<:NucleicAcidAlphabet{4}}
-const TwoBit = SeqOrView{<:NucleicAcidAlphabet{2}}
-
-const TwoBitSeq = BioSequence{<:NucleicAcidAlphabet{2}}
-const FourBitSeq = BioSequence{<:NucleicAcidAlphabet{4}}
-
 # Trivial overloads
 _n_certain(s::TwoBitSeq) = length(s)
 _n_ambiguous(s::TwoBitSeq) = 0

src/longsequences/find.jl

@@ -0,0 +1,155 @@
+# Zeros out all nibbles except the ones with 4bit A,C,G or T
+@inline function iscertain_kernel(::NucleicAcidAlphabet{4}, x::UInt64)
+    x = enumerate_nibbles(x)
+    y = (x & 0x4444444444444444) >> 2
+    y |= (x & 0x2222222222222222) >> 1
+    return x & ~y & 0x1111111111111111
+end
+
+# Zeros out all which is not a certain (normal or AA_Term)
+@inline function iscertain_kernel(::AminoAcidAlphabet, x::UInt64)
+    # 1. Set normal to FF, others to 00
+    y = simd_lt_byte(x, 0x16)
+
+    # 2. Set Term to FF, others to 00
+    z = set_zero_encoding(BitsPerSymbol{8}(), x ⊻ 0x1a1a1a1a1a1a1a1a) * 0xFF
+
+    # 3: OR them
+    return y | z
+end
+
+function _findnext(::typeof(iscertain), seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}}, i::Int)
+    f = @inline u -> iscertain_kernel(Alphabet(seq), u)
+    return _findnext_nonzero(f, seq, i)
+end
+
+function _findprev(::typeof(iscertain), seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}}, i::Int)
+    f = @inline u -> iscertain_kernel(Alphabet(seq), u)
+    return _findprev_nonzero(f, seq, i)
+end
+
+@inline function simd_lt_byte(x::UInt64, byte::UInt8)
+    T = NTuple{8, VecElement{UInt8}}
+    x = reinterpret(T, x)
+    y = ntuple(i -> VecElement(byte), Val{8}())
+    s = """
+    %res = icmp ult <8 x i8> %0, %1
+    %resb = sext <8 x i1> %res to <8 x i8>
+    ret <8 x i8> %resb
+    """
+    z = Core.Intrinsics.llvmcall(s, NTuple{8, VecElement{UInt8}}, Tuple{NTuple{8, VecElement{UInt8}}, NTuple{8, VecElement{UInt8}}}, x, y)
+    return reinterpret(UInt64, z)
+end
+
+@inline function sub_byte(x::UInt64, byte::UInt8)
+    T = NTuple{8, VecElement{UInt8}}
+    x = reinterpret(T, x)
+    y = ntuple(i -> VecElement(byte), Val{8}())
+    s = """
+    %res = sub <8 x i8> %0, %1
+    ret <8 x i8> %res
+    """
+    z = Core.Intrinsics.llvmcall(s, NTuple{8, VecElement{UInt8}}, Tuple{NTuple{8, VecElement{UInt8}}, NTuple{8, VecElement{UInt8}}}, x, y)
+    return reinterpret(UInt64, z)
+end
+
+# Zeros out all nibbles encoding 4bit A,C,G or T
+@inline function uncertain_kernel(::NucleicAcidAlphabet{4}, x::UInt64)
+    return enumerate_nibbles(x) ⊻ 0x1111111111111111
+end
+
+# Zero out normal AAs and AA_Term
+@inline function uncertain_kernel(::AminoAcidAlphabet, x::UInt64)
+    # Zero out normal AA, set rest to 0xFF
+    y = ~simd_lt_byte(x, 0x16)
+
+    # Zero out 0x1a, set rest to other bitpatterns
+    z = x ⊻ 0x1a1a1a1a1a1a1a1a
+
+    return y & z
+end
+
+function _findnext(
+        ::ComposedFunction{typeof(!), typeof(iscertain)},
+        seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}},
+        i::Int,
+    )
+    f = @inline u -> uncertain_kernel(Alphabet(seq), u)
+    return _findnext_nonzero(f, seq, i)
+end
+
+function _findprev(
+        ::ComposedFunction{typeof(!), typeof(iscertain)},
+        seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}},
+        i::Int,
+    )
+    f = @inline u -> uncertain_kernel(Alphabet(seq), u)
+    return _findprev_nonzero(f, seq, i)
+end
+
+# Zeros out A, C, G, T or Gap
+@inline function ambiguous_kernel(A::NucleicAcidAlphabet{4}, x::UInt64)
+    # The y part makes every nibble 0xF, unless it's 0 to begin with
+    y = x | (x >>> 2)
+    y |= y >>> 1
+    y &= 0x1111111111111111
+    y *= 0x0F
+    return uncertain_kernel(A, x) & y
+end
+
+# Zero out all except ambiguous symbols (AA_B, AA_J, AA_Z, AA_X), 0x16:0x19
+@inline function ambiguous_kernel(::AminoAcidAlphabet, x::UInt64)
+    return simd_lt_byte(sub_byte(x, 0x16), 0x04)
+end
+
+function _findnext(::typeof(isambiguous), seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}}, i::Int)
+    f = @inline u -> ambiguous_kernel(Alphabet(seq), u)
+    return _findnext_nonzero(f, seq, i)
+end
+
+function _findprev(::typeof(isambiguous), seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}}, i::Int)
+    f = @inline u -> ambiguous_kernel(Alphabet(seq), u)
+    return _findprev_nonzero(f, seq, i)
+end
+
+# Zeros out all except A, C, G, T and Gap
+@inline function unambiguous_kernel(::NucleicAcidAlphabet{4}, x::UInt64)
+    y = enumerate_nibbles(x)
+    y = (y >>> 1) | (y >>> 2)
+    y &= 0x1111111111111111
+    return y ⊻= 0x1111111111111111
+end
+
+# Zero out the four ambiguous amino acids B, J, Z, X
+@inline function unambiguous_kernel(::AminoAcidAlphabet, x::UInt64)
+    return ~simd_lt_byte(sub_byte(x, 0x16), 0x04)
+end
+
+function _findnext(
+        ::ComposedFunction{typeof(!), typeof(isambiguous)},
+        seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}},
+        i::Int,
+    )
+    f = @inline u -> unambiguous_kernel(Alphabet(seq), u)
+    return _findnext_nonzero(f, seq, i)
+end
+
+function _findprev(
+        ::ComposedFunction{typeof(!), typeof(isambiguous)},
+        seq::SeqOrView{<:Union{<:NucleicAcidAlphabet{4}, AminoAcidAlphabet}},
+        i::Int,
+    )
+    f = @inline u -> unambiguous_kernel(Alphabet(seq), u)
+    return _findprev_nonzero(f, seq, i)
+end
+
+# For debug
+#=
+function make_all_nibbles()
+    x = UInt64(0)
+    for i in 0:15
+        x = (x << 4) | i
+    end
+    x
+end
+=#

src/longsequences/longsequence.jl

@@ -120,3 +120,9 @@ include("copying.jl")
 include("stringliterals.jl")
 include("transformations.jl")
 include("operators.jl")
+
+const FourBit = SeqOrView{<:NucleicAcidAlphabet{4}}
+const TwoBit = SeqOrView{<:NucleicAcidAlphabet{2}}
+
+const TwoBitSeq = BioSequence{<:NucleicAcidAlphabet{2}}
+const FourBitSeq = BioSequence{<:NucleicAcidAlphabet{4}}