optimize byte code perf

Author: jrlouis21

rounds/3_dna/solution.py

@@ -7,77 +7,74 @@
 
 from __future__ import annotations
 
-import functools
 import os
 from concurrent.futures import ThreadPoolExecutor
+from mmap import ACCESS_READ, mmap
+from os import fstat
 
-_DELETE_TABLE = bytes.maketrans(b"", b"")
-_DELETE_CHARS = b"\n \r"
 _NUM_WORKERS = os.cpu_count() or 4
 
 
-@functools.lru_cache(maxsize=4)
-def _load(fasta_path: str) -> bytes:
-    with open(fasta_path, "rb") as f:
-        data = f.read()
-    boundaries = []
-    pos = data.find(b">")
-    while pos != -1:
-        nxt = data.find(b">", pos + 1)
-        boundaries.append((pos, nxt if nxt != -1 else len(data)))
-        pos = nxt
-    return data, boundaries
-
-
-def _search_chunk(
-    data: bytes,
-    pattern: bytes,
-    records: list[tuple[int, int]],
-) -> list[tuple[str, list[int]]]:
-    """Process a batch of (header_start, next_record_start) pairs."""
-    results: list[tuple[str, list[int]]] = []
-    for rec_start, rec_end in records:
-        nl = data.index(b"\n", rec_start)
-        seq = data[nl + 1 : rec_end].translate(_DELETE_TABLE, _DELETE_CHARS)
-
-        if pattern not in seq:
-            continue
-
-        record_id = data[rec_start + 1 : nl].strip().decode("ascii")
-
-        positions: list[int] = []
-        start = 0
-        _find = seq.find
-        while True:
-            idx = _find(pattern, start)
-            if idx == -1:
-                break
-            positions.append(idx)
-            start = idx + 1
-
-        if positions:
-            results.append((record_id, positions))
-    return results
+def _scan_record(record: bytes, pattern: bytes) -> tuple[str, list[int]] | None:
+    """Scan one FASTA record for all occurrences of ``pattern``.
 
+    Returns the record id and every zero-based match position, or ``None`` if
+    the record is empty or does not contain the pattern.
+    """
 
-def find_matches(fasta_path: str, pattern: bytes) -> list[tuple[str, list[int]]]:
-    """Find every FASTA record whose sequence contains ``pattern``.
+    if not record.strip():
+        return None
 
-    Returns ``[(record_id, [positions...]), ...]`` in file order.
-    """
-    data, boundaries = _load(fasta_path)
+    # Parition DNA record into header and DNA sequence
+    header, _, body = record.partition(b"\n")
+    record_id = header.strip().decode("ascii")
+
+    # Clean up data before parsing
+    sequence = body.replace(b"\n", b"").replace(b"\r", b"").replace(b" ", b"")
+
+    positions: list[int] = []
+    start = 0
 
-    if not boundaries:
-        return []
+    # Advance by one after each hit so overlapping matches are included.
+    while True:
+        pos = sequence.find(pattern, start)
+        if pos == -1:
+            break
+        positions.append(pos)
+        start = pos + 1
 
-    n = len(boundaries)
-    chunk_size = max(1, n // _NUM_WORKERS)
-    chunks = [boundaries[i : i + chunk_size] for i in range(0, n, chunk_size)]
+    if not positions:
+        return None
 
-    matches: list[tuple[str, list[int]]] = []
-    with ThreadPoolExecutor(max_workers=_NUM_WORKERS) as executor:
-        futures = [executor.submit(_search_chunk, data, pattern, chunk) for chunk in chunks]
-        for future in futures:
-            matches.extend(future.result())
+    return record_id, positions
 
-    return matches
+
+def find_matches(fasta_path: str, pattern: bytes) -> list[tuple[str, list[int]]]:
+    """Find every FASTA record whose sequence contains ``pattern``.
+
+    Returns ``[(record_id, [positions...]), ...]`` in file order.
+    """
+    with open(fasta_path, "rb") as f:
+        if fstat(f.fileno()).st_size == 0:
+            return []
+
+        with mmap(f.fileno(), 0, access=ACCESS_READ) as text:
+            # Read the file as an mmap and break it up into DNA records
+            records: list[bytes] = []
+            start = text.find(b">")
+            while start != -1:
+                end = text.find(b">", start + 1)
+                if end == -1:
+                    record = text[start + 1 :]
+                    start = -1
+                else:
+                    record = text[start + 1 : end]
+                    start = end
+
+                if record.strip():
+                    records.append(record)
+
+            with ThreadPoolExecutor(max_workers=_NUM_WORKERS) as executor:
+                results = executor.map(lambda record: _scan_record(record, pattern), records)
+
+    return [result for result in results if result is not None]