rsjohnson performance improvements

README.md

@@ -1,3 +1,4 @@
+# rsjohnson
 # Python Performance Lab: Sharpening Your Instincts
 
 A PyCon US 2026 hands-on tutorial. You optimize intentionally slow Python code

rounds/1_histogram/baseline.py

@@ -4,21 +4,24 @@
 tokens) in a binary payload.
 """
 
+from pathlib import Path
 
-def compute_histogram(path: str) -> dict[bytes, int]:
-    """Frequency of every 2-byte bigram in the file at ``path``."""
-    # Step 1: read the whole file into memory as a single bytes object.
-    with open(path, "rb") as f:
-        data = f.read()
+# -------------------------------------------------------------------------------------------------
 
-    # Step 2: slide a 2-byte window across the buffer. For ``b"ABCD"`` the
-    # iterations produce ``b"AB"``, ``b"BC"``, then ``b"CD"``. For each window,
-    # bump the matching bucket in a ``dict`` keyed by the bigram itself.
-    counts: dict[bytes, int] = {}
-    for i in range(len(data) - 1):
-        bigram = data[i : i + 2]
-        if bigram in counts:
-            counts[bigram] += 1
-        else:
-            counts[bigram] = 1
-    return counts
+def compute_histogram(path):
+    """Return frequency of every 2-byte bigram in the file at path."""
+    counts = [0] * 65536
+    previous = None
+
+    with Path(path).open("rb") as file:
+        while chunk := file.read(1024 * 1024):
+            for byte in chunk:
+                if previous is not None:
+                    counts[(previous << 8) | byte] += 1
+                previous = byte
+
+    return {
+        bigram.to_bytes(2, "big"): count
+        for bigram, count in enumerate(counts)
+        if count
+    }

rounds/2_corruption/baseline.py

@@ -1,46 +1,128 @@
-"""Round 2 baseline: corruption scanner.
+from __future__ import annotations
 
-Compares two equally-sized binary files and reports every contiguous run of
-differing bytes as ``(offset, length)``.
-"""
+import mmap
+import os
+from os import PathLike
+from typing import Union
 
-from __future__ import annotations
+import numpy as np
+
+
+Pathish = Union[str, bytes, PathLike[str], PathLike[bytes]]
+
+
+def find_corruptions(
+    ref_path: Pathish,
+    cor_path: Pathish,
+    *,
+    chunk_size: int = 1 << 26,  # 64 MiB
+) -> list[tuple[int, int]]:
+    """
+    Return [(offset, length), ...] for every differing byte range.
+
+    Optimizations:
+    - checks file sizes before reading
+    - memory-maps both files
+    - compares bytes using NumPy's native vectorized code
+    - records only transition points, not every differing offset
+    - handles corruption ranges that cross chunk boundaries
+    """
 
+    ref_size = os.path.getsize(ref_path)
+    cor_size = os.path.getsize(cor_path)
 
-def find_corruptions(ref_path: str, cor_path: str) -> list[tuple[int, int]]:
-    """Return ``[(offset, length), ...]`` for every differing byte range."""
-    # Step 1: read both files fully into memory as bytes objects.
-    with open(ref_path, "rb") as f:
-        ref = f.read()
-    with open(cor_path, "rb") as f:
-        cor = f.read()
-    if len(ref) != len(cor):
+    if ref_size != cor_size:
         raise ValueError("reference and corrupted files differ in length")
 
-    # Step 2: walk both buffers in lockstep and record every position where
-    # the two files disagree. The result is a sorted list of standalone byte
-    # offsets, e.g. [3, 4, 5, 17, 18].
-    diffs: list[int] = []
-    for i in range(len(ref)):
-        if ref[i] != cor[i]:
-            diffs.append(i)
-
-    # Step 3: collapse runs of consecutive offsets into (start, length) ranges.
-    # The list from step 2 becomes [(3, 3), (17, 2)]: starting at 3 there are
-    # three differing bytes, then starting at 17 there are two more.
-    if not diffs:
+    if ref_size == 0:
         return []
+
+    if chunk_size <= 0:
+        raise ValueError("chunk_size must be positive")
+
+    chunk_size = min(chunk_size, ref_size)
+
     ranges: list[tuple[int, int]] = []
-    start = diffs[0]
-    prev = diffs[0]
-    for pos in diffs[1:]:
-        if pos == prev + 1:
-            # Still inside the current run; extend it.
-            prev = pos
-        else:
-            # Gap. Close the current run and start a new one.
-            ranges.append((start, prev - start + 1))
-            start = pos
-            prev = pos
-    ranges.append((start, prev - start + 1))  # Close the final run.
+    append = ranges.append
+
+    in_run = False
+    run_start = 0
+
+    # Reuse this buffer so we do not allocate a new boolean array per chunk.
+    diff_buffer = np.empty(chunk_size, dtype=np.bool_)
+
+    with open(ref_path, "rb") as ref_file, open(cor_path, "rb") as cor_file:
+        with (
+            mmap.mmap(ref_file.fileno(), 0, access=mmap.ACCESS_READ) as ref_map,
+            mmap.mmap(cor_file.fileno(), 0, access=mmap.ACCESS_READ) as cor_map,
+        ):
+            for offset in range(0, ref_size, chunk_size):
+                stop = min(offset + chunk_size, ref_size)
+                length = stop - offset
+
+                ref_chunk = np.frombuffer(
+                    ref_map,
+                    dtype=np.uint8,
+                    count=length,
+                    offset=offset,
+                )
+                cor_chunk = np.frombuffer(
+                    cor_map,
+                    dtype=np.uint8,
+                    count=length,
+                    offset=offset,
+                )
+
+                diff = diff_buffer[:length]
+                np.not_equal(ref_chunk, cor_chunk, out=diff)
+
+                # Fast path: this entire chunk is identical.
+                if not bool(diff.any()):
+                    if in_run:
+                        append((run_start, offset - run_start))
+                        in_run = False
+
+                    del ref_chunk, cor_chunk, diff
+                    continue
+
+                # Fast path: this entire chunk differs.
+                if bool(diff.all()):
+                    if not in_run:
+                        run_start = offset
+                        in_run = True
+
+                    del ref_chunk, cor_chunk, diff
+                    continue
+
+                # Handle a transition at the chunk boundary.
+                first_is_diff = bool(diff[0])
+                if first_is_diff != in_run:
+                    if in_run:
+                        append((run_start, offset - run_start))
+                        in_run = False
+                    else:
+                        run_start = offset
+                        in_run = True
+
+                # Internal transitions:
+                # False -> True starts a corruption range.
+                # True -> False closes a corruption range.
+                transitions = np.flatnonzero(diff[1:] != diff[:-1]) + 1
+
+                for transition in transitions:
+                    pos = offset + int(transition)
+
+                    if in_run:
+                        append((run_start, pos - run_start))
+                        in_run = False
+                    else:
+                        run_start = pos
+                        in_run = True
+
+                # Release mmap-backed NumPy views before closing mmap objects.
+                del ref_chunk, cor_chunk, diff, transitions
+
+            if in_run:
+                append((run_start, ref_size - run_start))
+
     return ranges