perf(store): find chunk cutoffs and code fences by binary search

CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Fixed
 
+- Chunking time no longer grows quadratically with a document's break points
+  and code fences. A 1 MiB document of nothing but fence lines took about 4.4 s
+  and now takes about 0.1 s. Chunk boundaries do not change. #1102
 - `qmd doctor` selects vector sample identities before loading document bodies,
   avoiding excessive SQLite memory use on large indexes with duplicate paths.
   #978 (thanks @naveenspark)

src/store.ts

@@ -290,11 +290,31 @@ export function findCodeFences(text: string): CodeFenceRegion[] {
   return regions;
 }
 
+/**
+ * Index of the first item for which `isBefore` is false. `items` must hold every
+ * item for which it is true ahead of every item for which it is false.
+ */
+function partitionPoint<T>(items: T[], isBefore: (item: T) => boolean): number {
+  let lo = 0;
+  let hi = items.length;
+  while (lo < hi) {
+    const mid = (lo + hi) >>> 1;
+    if (isBefore(items[mid]!)) lo = mid + 1;
+    else hi = mid;
+  }
+  return lo;
+}
+
 /**
  * Check if a position is inside a code fence region.
+ *
+ * `fences` must be sorted and non-overlapping, as findCodeFences returns them.
+ * Only the last fence starting before `pos` can then contain it: every earlier
+ * fence has ended by the time that one starts.
  */
 export function isInsideCodeFence(pos: number, fences: CodeFenceRegion[]): boolean {
-  return fences.some(f => pos > f.start && pos < f.end);
+  const next = partitionPoint(fences, f => f.start < pos);
+  return next > 0 && pos < fences[next - 1]!.end;
 }
 
 /**
@@ -303,7 +323,7 @@ export function isInsideCodeFence(pos: number, fences: CodeFenceRegion[]): boole
  * Uses squared distance for gentler early decay - headings far back still win
  * over low-quality breaks near the target.
  *
- * @param breakPoints - Pre-scanned break points from scanBreakPoints()
+ * @param breakPoints - Pre-scanned break points from scanBreakPoints(), sorted by position
  * @param targetCharPos - The ideal cut position (e.g., maxChars boundary)
  * @param windowChars - How far back to search for break points (default ~200 tokens)
  * @param decayFactor - How much to penalize distance (0.7 = 30% score at window edge)
@@ -321,8 +341,11 @@ export function findBestCutoff(
   let bestScore = -1;
   let bestPos = targetCharPos;
 
-  for (const bp of breakPoints) {
-    if (bp.pos < windowStart) continue;
+  // Sorted, so jump to the window: a document's chunks would otherwise each walk
+  // every break point before it.
+  const first = partitionPoint(breakPoints, bp => bp.pos < windowStart);
+  for (let i = first; i < breakPoints.length; i++) {
+    const bp = breakPoints[i]!;
     if (bp.pos > targetCharPos) break;  // sorted, so we can stop
 
     // Skip break points inside code fences

test/store.test.ts

@@ -877,6 +877,21 @@ describe("findCodeFences", () => {
   });
 });
 
+// Documents that put fences and break points where an off-by-one would show.
+const CHUNKER_EDGE_DOCS = [
+  "",
+  "```",
+  "```\ncode\n```\nafter\n\nmore",
+  "intro\n```\nnever closed\n\n# heading\n\nmore\n",
+  "a\n```\nb\n```\nc\n```\nd\n\ne",
+  "\n```\n```\n```\n```\n```\n```\n",
+  "\n``````\n``````\n```",
+  "x\n````\ncode\n```\ninner\n````\ny\n",
+  "intro\r\n\r\n```\r\ncode\r\n```\r\n\r\n# head\r\n\r\nbody\r\n---\r\nend\r\n",
+  "é\n\n```\n😀😀\n```\n\n# 見出し\n\n- 😀\n1. é\n",
+  "# a\n```\n# not a heading\n\n\n```\n# b\n",
+];
+
 describe("isInsideCodeFence", () => {
   test("returns true for position inside fence", () => {
     const fences: CodeFenceRegion[] = [{ start: 10, end: 30 }];
@@ -905,6 +920,37 @@ describe("isInsideCodeFence", () => {
     expect(isInsideCodeFence(60, fences)).toBe(true);
     expect(isInsideCodeFence(40, fences)).toBe(false);
   });
+
+  test("handles no fences", () => {
+    expect(isInsideCodeFence(5, [])).toBe(false);
+  });
+
+  test("treats the edge shared by two tight fences as outside both", () => {
+    const fences = findCodeFences("\n```\n```\n```\n```");
+    expect(fences).toEqual([{ start: 0, end: 8 }, { start: 8, end: 16 }]);
+    expect(isInsideCodeFence(7, fences)).toBe(true);
+    expect(isInsideCodeFence(8, fences)).toBe(false);
+    expect(isInsideCodeFence(9, fences)).toBe(true);
+  });
+
+  test("agrees with a linear scan at every position", () => {
+    for (const text of CHUNKER_EDGE_DOCS) {
+      const fences = findCodeFences(text);
+      for (let pos = -1; pos <= text.length + 1; pos++) {
+        expect(isInsideCodeFence(pos, fences)).toBe(fences.some(f => pos > f.start && pos < f.end));
+      }
+    }
+  });
+
+  test("reads only a few of many fences", () => {
+    let reads = 0;
+    const fences: CodeFenceRegion[] = Array.from({ length: 100_000 }, (_, i) => ({
+      get start() { reads++; return i * 10; },
+      end: i * 10 + 8,
+    }));
+    expect(isInsideCodeFence(999_995, fences)).toBe(true);
+    expect(reads).toBeLessThan(100);
+  });
 });
 
 describe("findBestCutoff", () => {
@@ -966,6 +1012,63 @@ describe("findBestCutoff", () => {
     const cutoff = findBestCutoff([], 200, 100, 0.7);
     expect(cutoff).toBe(200);
   });
+
+  test("considers a break point exactly on the window edge, and not one before it", () => {
+    const onEdge: BreakPoint = { pos: 100, score: 100, type: 'h1' };
+    const beforeEdge: BreakPoint = { pos: 99, score: 100, type: 'h1' };
+    expect(findBestCutoff([onEdge], 200, 100, 0.7)).toBe(100);
+    expect(findBestCutoff([beforeEdge], 200, 100, 0.7)).toBe(200);
+  });
+
+  test("considers a break point exactly on the target, and not one after it", () => {
+    const onTarget: BreakPoint = { pos: 200, score: 1, type: 'newline' };
+    const afterTarget: BreakPoint = { pos: 201, score: 100, type: 'h1' };
+    expect(findBestCutoff([onTarget, afterTarget], 200, 100, 0.7)).toBe(200);
+    expect(findBestCutoff([afterTarget], 200, 100, 0.7)).toBe(200);
+  });
+
+  test("agrees with a linear walk for every target and window", () => {
+    const linearCutoff = (
+      breakPoints: BreakPoint[], target: number, windowChars: number, codeFences: CodeFenceRegion[]
+    ): number => {
+      let bestScore = -1;
+      let bestPos = target;
+      for (const bp of breakPoints) {
+        if (bp.pos < target - windowChars) continue;
+        if (bp.pos > target) break;
+        if (codeFences.some(f => bp.pos > f.start && bp.pos < f.end)) continue;
+        const normalizedDist = (target - bp.pos) / windowChars;
+        const finalScore = bp.score * (1.0 - (normalizedDist * normalizedDist) * 0.7);
+        if (finalScore > bestScore) {
+          bestScore = finalScore;
+          bestPos = bp.pos;
+        }
+      }
+      return bestPos;
+    };
+
+    for (const text of CHUNKER_EDGE_DOCS) {
+      const breakPoints = scanBreakPoints(text);
+      const fences = findCodeFences(text);
+      for (let target = 0; target <= text.length + 1; target++) {
+        for (const windowChars of [0, 1, 8, 50]) {
+          expect(findBestCutoff(breakPoints, target, windowChars, 0.7, fences))
+            .toBe(linearCutoff(breakPoints, target, windowChars, fences));
+        }
+      }
+    }
+  });
+
+  test("reads only the break points near its window", () => {
+    let reads = 0;
+    const breakPoints: BreakPoint[] = Array.from({ length: 100_000 }, (_, i) => ({
+      get pos() { reads++; return i * 10; },
+      score: 1,
+      type: 'newline',
+    }));
+    expect(findBestCutoff(breakPoints, 999_000, 100, 0.7)).toBe(999_000);
+    expect(reads).toBeLessThan(100);
+  });
 });
 
 describe("Smart Chunking Integration", () => {