Implement k-mer distance metrics and constrained generation

Adds support for Hamming distance, constrained k-mer generation, and optimized Longest Common Subsequence (LCS) calculations using scalar, batch, and SIMD methods.

This includes new methods for k-mer distance calculation, structures for k-mer constraints, and SIMD-based batch processing for LCS to improve performance and correctness.

pkg/obikmer/kmer_distance.go

@@ -0,0 +1,104 @@
+package obikmer
+
+import "math/bits"
+
+// KmerLCSResult is the result of a k-mer LCS alignment: the length of the
+// Longest Common Subsequence and the length of the shortest alignment
+// achieving it (same semantics as obialign.FastLCSScore).
+type KmerLCSResult struct {
+	LCS      int
+	AlignLen int
+}
+
+// Distance returns AlignLen - LCS, the LCS-mode edit distance.
+func (r KmerLCSResult) Distance() int {
+	return r.AlignLen - r.LCS
+}
+
+// HammingDistance returns the number of positions at which the two k-mers
+// differ. Both k-mers are encoded with 2 bits per base (see EncodeKmer);
+// bits above the 2*k lowest ones are ignored. k must be between 0 and 32.
+func HammingDistance(a, b uint64, k int) int {
+	diff := a ^ b
+	diff = (diff | diff>>1) & 0x5555555555555555
+	diff &= ^uint64(0) >> (64 - 2*k)
+	return bits.OnesCount64(diff)
+}
+
+// KmerLCS aligns two k-mers of the same size in LCS mode (match +1, mismatch
+// and indel 0) and returns the LCS length together with the length of the
+// shortest alignment among those reaching it, so that the distance is
+// AlignLen - LCS.
+//
+// It is the scalar reference implementation, straightforward on purpose:
+// it is the oracle against which the vectorised versions are tested.
+// It returns {-1, -1} if k is not between 1 and 31.
+func KmerLCS(a, b uint64, k int) KmerLCSResult {
+	if k < 1 || k > 31 {
+		return KmerLCSResult{-1, -1}
+	}
+
+	var ba, bb [32]uint8
+	for i := k - 1; i >= 0; i-- {
+		ba[i] = uint8(a & 3)
+		bb[i] = uint8(b & 3)
+		a >>= 2
+		b >>= 2
+	}
+
+	// One DP row: best (score, alignment length) per column.
+	var score, length [32]int
+	for j := 0; j <= k; j++ {
+		length[j] = j
+	}
+
+	better := func(s1, l1, s2, l2 int) bool {
+		return s1 > s2 || (s1 == s2 && l1 < l2)
+	}
+
+	for i := 1; i <= k; i++ {
+		diagS, diagL := score[0], length[0]
+		score[0], length[0] = 0, i
+
+		for j := 1; j <= k; j++ {
+			upS, upL := score[j], length[j]
+
+			bestS, bestL := diagS, diagL+1
+			if ba[j-1] == bb[i-1] {
+				bestS++
+			}
+			if better(upS, upL+1, bestS, bestL) {
+				bestS, bestL = upS, upL+1
+			}
+			if better(score[j-1], length[j-1]+1, bestS, bestL) {
+				bestS, bestL = score[j-1], length[j-1]+1
+			}
+
+			diagS, diagL = upS, upL
+			score[j], length[j] = bestS, bestL
+		}
+	}
+
+	return KmerLCSResult{score[k], length[k]}
+}
+
+// KmerLCSBatch computes KmerLCS(query, partners[i], k) for every partner and
+// stores it in out[i]. out must be at least as long as partners.
+//
+// When built with GOEXPERIMENT=simd the partners are aligned several at a
+// time, one per SIMD lane; otherwise the scalar reference is used.
+// It returns false, leaving out untouched, if k is not between 1 and 31 or
+// if out is too short.
+func KmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) bool {
+	if k < 1 || k > 31 || len(out) < len(partners) {
+		return false
+	}
+	kmerLCSBatch(query, partners, k, out)
+	return true
+}
+
+func kmerLCSBatchScalar(query uint64, partners []uint64, k int, out []KmerLCSResult) {
+	for i, p := range partners {
+		out[i] = KmerLCS(query, p, k)
+	}
+}