UPDATE the scan to drop a match before it builds one

CLAUDE.md

@@ -337,10 +337,12 @@ and narrowing it would buy nothing.
     - **`N` in a PAM pattern is satisfied by a genomic `N`.** `NGG` asks
       nothing of its first position, so an unsequenced base there is still a
       PAM; every other code names bases, and an `N` is not known to be one of
-      them. The rule lives in `base_matches` and in the Rust `matches_pam`,
+      them. The rule lives in `base_matches` and in the Rust `code_ok`,
       which disagreed with it until they were reconciled, and in the oracle
       of `tests/test_core_scan.py`, which had encoded the Rust reading.
-      Change one and the other two are wrong;
+      Change one and the other two are wrong. `code_ok` is read twice per
+      match, once by `pam_ok_at` on the raw bytes before anything is
+      allocated, so keep it branch-light;
     - classes count sites whose genomic PAM matches. A site whose PAM does
       not stays in the table with `pam_matches` `FALSE` and counts nowhere,
       so a spacer matching twice with one PAM is `unique`. `--pam-only`

rust/src/lib.rs

@@ -14,7 +14,7 @@ use std::io::{BufRead, BufReader};
 use pyo3::exceptions::{PyOSError, PyValueError};
 use pyo3::prelude::*;
 
-use scan::{matches_pam, pam_length, Spacers};
+use scan::{pam_length, Spacers};
 
 /// The crate version, so Python can check which core it loaded.
 #[pyfunction]
@@ -56,7 +56,8 @@ fn scan_fasta(
     contigs: Option<Vec<String>>,
     pam_only: bool,
 ) -> PyResult<(Vec<PyHit>, Vec<PyCounts>)> {
-    let width = pam_length(pam).map_err(PyValueError::new_err)?;
+    // Validates the pattern; the scan reads it base by base from here on.
+    pam_length(pam).map_err(PyValueError::new_err)?;
     let table = Spacers::new(&spacers).map_err(PyValueError::new_err)?;
     let wanted: Option<HashSet<String>> = contigs.map(|names| names.into_iter().collect());
     let mut seen: HashSet<String> = HashSet::new();
@@ -84,27 +85,19 @@ fn scan_fasta(
             if !name.is_empty() && wanted.as_ref().is_none_or(|set| set.contains(&name)) {
                 // Releasing the GIL here is what makes a whole genome bearable
                 // from Python: the scan touches nothing the interpreter owns.
-                let found = py.detach(|| scan::scan(&sequence, &table, width, upstream));
+                // The scan tallies and filters as it goes, so a match with
+                // no PAM costs neither a string nor a row of this list.
+                let found = py
+                    .detach(|| scan::scan(&sequence, &table, pam, upstream, pam_only, &mut counts));
                 for hit in found {
-                    let matched = matches_pam(&hit.pam, pam);
-                    let tally = &mut counts[hit.guide as usize];
-                    tally.0 += 1;
-                    if hit.pam.len() != width {
-                        tally.1 += 1;
-                    } else if !matched {
-                        tally.2 += 1;
-                    }
-                    if pam_only && !matched {
-                        continue;
-                    }
                     hits.push((
                         hit.guide,
                         name.clone(),
                         hit.start,
                         if hit.reverse { "-" } else { "+" },
                         hit.pam,
                         hit.upstream,
-                        matched,
+                        hit.matched,
                     ));
                 }
             }

rust/src/scan.rs

@@ -165,6 +165,7 @@ pub struct Hit {
     pub reverse: bool,
     pub pam: String,
     pub upstream: String,
+    pub matched: bool,
 }
 
 /// Every exact match of `spacers` in `sequence`, with its flanks.
@@ -173,8 +174,16 @@ pub struct Hit {
 /// it is the left end whichever strand the guide reads on. `pam` and
 /// `upstream` are in guide orientation, and are empty when the flank would
 /// run off the end of the contig.
-pub fn scan(sequence: &[u8], spacers: &Spacers, pam_length: usize, upstream: usize) -> Vec<Hit> {
+pub fn scan(
+    sequence: &[u8],
+    spacers: &Spacers,
+    pattern: &str,
+    upstream: usize,
+    pam_only: bool,
+    counts: &mut [(u64, u64, u64)],
+) -> Vec<Hit> {
     let length = spacers.length;
+    let pam_length = pattern.len();
     let mut hits = Vec::new();
     if sequence.len() < length {
         return hits;
@@ -198,16 +207,37 @@ pub fn scan(sequence: &[u8], spacers: &Spacers, pam_length: usize, upstream: usi
         };
         let start = index + 1 - length;
         for entry in entries {
+            // The guide reading right to left has its PAM 5' of the match on
+            // the forward strand, reverse-complemented.
+            let (pam_from, pam_limit) = if entry.reverse {
+                (start.wrapping_sub(pam_length), start)
+            } else {
+                (start + length, sequence.len())
+            };
+            let whole = fits(sequence, pam_from, pam_length, pam_limit);
+            let matched = whole && pam_ok_at(sequence, pam_from, pam_limit, pattern, entry.reverse);
+
+            let tally = &mut counts[entry.guide as usize];
+            tally.0 += 1;
+            if !whole {
+                tally.1 += 1;
+            } else if !matched {
+                tally.2 += 1;
+            }
+            // Nineteen of twenty matches stop here, before either string of
+            // the `Hit` they would otherwise have been built into.
+            if pam_only && !matched {
+                continue;
+            }
+
             let (pam, up) = if entry.reverse {
-                // The guide reads right to left here: its PAM is 5' of the
-                // match on the forward strand, reverse-complemented.
                 (
-                    revcomp_slice(sequence, start.wrapping_sub(pam_length), pam_length, start),
+                    revcomp_slice(sequence, pam_from, pam_length, pam_limit),
                     revcomp_slice(sequence, start + length, upstream, sequence.len()),
                 )
             } else {
                 (
-                    plain_slice(sequence, start + length, pam_length, sequence.len()),
+                    plain_slice(sequence, pam_from, pam_length, pam_limit),
                     plain_slice(sequence, start.wrapping_sub(upstream), upstream, start),
                 )
             };
@@ -217,6 +247,7 @@ pub fn scan(sequence: &[u8], spacers: &Spacers, pam_length: usize, upstream: usi
                 reverse: entry.reverse,
                 pam,
                 upstream: up,
+                matched,
             });
         }
     }
@@ -228,10 +259,7 @@ fn plain_slice(sequence: &[u8], from: usize, length: usize, limit: usize) -> Str
     if length == 0 {
         return String::new();
     }
-    if from > limit
-        || from.saturating_add(length) > limit
-        || from.saturating_add(length) > sequence.len()
-    {
+    if !fits(sequence, from, length, limit) {
         return String::new();
     }
     sequence[from..from + length]
@@ -283,21 +311,45 @@ pub fn pam_length(pattern: &str) -> Result<usize, String> {
     Ok(pattern.len())
 }
 
-/// Whether a stretch of genomic bases matches an IUPAC pattern.
-pub fn matches_pam(bases: &str, pattern: &str) -> bool {
-    if bases.len() != pattern.len() {
+/// Whether one base code satisfies one IUPAC symbol.
+///
+/// N stands for every base, so an unsequenced one satisfies it: NGG needs
+/// nothing of its first position. Any other code names bases, and an N in
+/// the reference is not known to be one of them. This is the rule
+/// `base_matches` states in Python; the two have to agree.
+fn code_ok(c: u8, symbol: u8) -> bool {
+    match iupac_mask(symbol) {
+        None => false,
+        Some(ANY) => true,
+        Some(mask) => c != INVALID && mask >> c & 1 == 1,
+    }
+}
+
+/// Whether the PAM at `from` matches `pattern`, read in guide orientation.
+///
+/// The point of it is what it does not do: a match whose PAM fails is the
+/// common case by nineteen to one, and this answers without building the
+/// two strings a `Hit` would carry.
+fn pam_ok_at(sequence: &[u8], from: usize, limit: usize, pattern: &str, reverse: bool) -> bool {
+    if !fits(sequence, from, pattern.len(), limit) {
         return false;
     }
-    bases.bytes().zip(pattern.bytes()).all(|(base, symbol)| {
-        let c = code(base);
-        match (c, iupac_mask(symbol)) {
-            (_, None) => false,
-            // N stands for every base, so an unsequenced one satisfies it:
-            // NGG needs nothing of its first position. Any other code names
-            // bases, and an N in the reference is not known to be one of them.
-            (_, Some(ANY)) => true,
-            (INVALID, _) => false,
-            (c, Some(mask)) => mask >> c & 1 == 1,
-        }
+    let bases = &sequence[from..from + pattern.len()];
+    pattern.bytes().enumerate().all(|(i, symbol)| {
+        // Reading right to left and complementing, as `revcomp_slice` does.
+        let raw = if reverse {
+            bases[pattern.len() - 1 - i]
+        } else {
+            bases[i]
+        };
+        let c = code(raw);
+        code_ok(if reverse && c != INVALID { 3 - c } else { c }, symbol)
     })
 }
+
+/// Whether `length` bases from `from` all exist within `limit`.
+fn fits(sequence: &[u8], from: usize, length: usize, limit: usize) -> bool {
+    from <= limit
+        && from.saturating_add(length) <= limit
+        && from.saturating_add(length) <= sequence.len()
+}

tests/test_core_scan.py

@@ -241,3 +241,44 @@ def test_an_unsequenced_base_satisfies_an_n_of_the_pattern(tmp_path):
 
     assert {hit[1]: hit[6] for hit in hits} == {"wild": True, "unknown": False}
     assert counts == [(2, 0, 1)]
+
+
+@pytest.mark.parametrize("seed", range(6))
+def test_pam_only_agrees_with_filtering_the_full_scan(tmp_path, seed):
+    """The filter moved into the scan; it must still change nothing but cost.
+
+    The oracle checks only the unfiltered path, so this is what would catch
+    the tally drifting once it is kept somewhere else.
+    """
+    rng = random.Random(seed)
+    spacers = [random_sequence(rng, 8) for _ in range(6)]
+    contigs = {f"chr{n}": random_sequence(rng, 400, "ACGTN") for n in range(3)}
+    path = tmp_path / "genome.fa"
+    write_fasta(path, contigs)
+
+    full, counts = _core.scan_fasta(str(path), spacers, "NGG", 1)
+    only, same = _core.scan_fasta(str(path), spacers, "NGG", 1, None, True)
+
+    assert counts == same == tally(oracle(contigs, spacers, "NGG", 1), spacers, "NGG")
+    assert sorted(only) == sorted(hit for hit in full if hit[6])
+
+
+def test_a_minus_strand_match_against_the_very_start_of_a_contig(tmp_path):
+    # Its PAM lies 5' of base 0, which does not exist: the bounds check has
+    # to catch `start - 3` wrapping round rather than read somewhere else.
+    spacer = "ACGTTGCAAGCTTCGATCGA"
+    contigs = {"chrEdge": reverse_complement_py(spacer) + "TTTTT"}
+    path = tmp_path / "genome.fa"
+    write_fasta(path, contigs)
+
+    hits, counts = _core.scan_fasta(str(path), [spacer], "NGG", 1)
+
+    assert [hit[3] for hit in hits] == ["-"]
+    assert hits[0][4] == ""
+    assert counts == [(1, 1, 0)]
+    only, same = _core.scan_fasta(str(path), [spacer], "NGG", 1, None, True)
+    assert only == [] and same == counts
+
+
+def reverse_complement_py(sequence):
+    return sequence[::-1].translate(str.maketrans("ACGT", "TGCA"))