Filter regex start positions by their leading character sets (#1784)

packages/compiler/src/backend/llvm/regex-matchers.test.ts

@@ -110,6 +110,25 @@ describe("regex literal matcher code", () => {
     expect(body("a\\d{1,3}")).not.toContain("<16 x i8>");
   });
 
+  test("fixed leading characters filter 16 start positions at once", () => {
+    // Both alternatives' first characters are tested at every offset and
+    // combined, and the reads reach seven characters past the block.
+    const text = body("agggtaaa|tttaccct", "gi");
+    expect(text.match(/load <16 x i8>/g)!.length).toBeGreaterThanOrEqual(16);
+    expect(text).toMatch(/and i16 /);
+    expect(text).toMatch(/or i16 /);
+    expect(text).toMatch(/add i32 %t\d+, 7\n/);
+    // A class sequence with a fixed count extends the filter.
+    expect(body("[cg]x{3}y")).toMatch(/and i16 /);
+    // An optional first character, a single first character and a
+    // variable-length start keep the first-character search.
+    expect(body("[ab]?cd")).not.toMatch(/and i16 /);
+    expect(body("[ab]*cd")).not.toMatch(/and i16 /);
+    expect(body("xy|zw")).toMatch(/and i16 /);
+    expect(body("[ab]|cd")).not.toMatch(/and i16 /);
+    expect(body("agggtaaa|tttaccct", "gi", false)).not.toContain("<16 x i8>");
+  });
+
   test("loops no continuation can share push no backtrack entries", () => {
     // \d+ then a dot; [^\]]+ then ]; the run at the end of the pattern.
     expect(body("\\d+\\.\\d+")).not.toContain("@scr_regex_stack_grow");

packages/compiler/src/backend/llvm/regex-matchers.ts

@@ -589,6 +589,90 @@ function firstOf(node: RxNode): FirstInfo {
   }
 }
 
+/** Positions the start filter tests and alternatives it keeps. */
+const PREFIX_LENGTH = 8;
+const PREFIX_ALTERNATIVES = 8;
+
+/** The character sets a match's first characters must belong to: every
+ * match of the node begins with the characters of some entry of `alts`,
+ * one set per character. `complete` means every match is exactly as long as
+ * its entry (so a following node continues the sequences). Null when the
+ * node gives no usable condition (an alternative that can start with
+ * anything, or too many alternatives). Assertions add no characters: the
+ * result stays a necessary condition. */
+interface Prefix {
+  alts: CharSet[][];
+  complete: boolean;
+}
+
+/** Each head followed by each tail, cut to PREFIX_LENGTH sets; null past
+ * PREFIX_ALTERNATIVES entries. */
+function crossPrefixes(heads: CharSet[][], tails: CharSet[][]): CharSet[][] | null {
+  if (heads.length * tails.length > PREFIX_ALTERNATIVES) return null;
+  const out: CharSet[][] = [];
+  for (const head of heads)
+    for (const tail of tails) out.push([...head, ...tail].slice(0, PREFIX_LENGTH));
+  return out;
+}
+
+/** Whether no entry was cut at PREFIX_LENGTH. */
+function uncut(alts: CharSet[][]): boolean {
+  return alts.every((alt) => alt.length < PREFIX_LENGTH);
+}
+
+function prefixOf(node: RxNode): Prefix | null {
+  switch (node.kind) {
+    case "set":
+      return { alts: [[node.set]], complete: true };
+    case "assert":
+      return { alts: [[]], complete: true };
+    case "cap":
+      return prefixOf(node.body);
+    case "alt": {
+      const alts: CharSet[][] = [];
+      let complete = true;
+      for (const alt of node.alts) {
+        const p = prefixOf(alt);
+        if (p === null) return null;
+        alts.push(...p.alts);
+        complete &&= p.complete;
+      }
+      return alts.length <= PREFIX_ALTERNATIVES ? { alts, complete } : null;
+    }
+    case "seq": {
+      let alts: CharSet[][] = [[]];
+      for (const item of node.items) {
+        const p = prefixOf(item);
+        if (p === null) return { alts, complete: false };
+        const joined = crossPrefixes(alts, p.alts);
+        if (joined === null) return { alts, complete: false };
+        alts = joined;
+        if (!p.complete || !uncut(alts)) return { alts, complete: false };
+      }
+      return { alts, complete: true };
+    }
+    case "rep": {
+      if (node.max === 0) return { alts: [[]], complete: true };
+      if (node.min === 0) return { alts: [[]], complete: false };
+      const body = prefixOf(node.body);
+      if (body === null) return null;
+      // The first `min` iterations are required; they extend the sequences
+      // while the body has a fixed length.
+      let alts = body.alts;
+      let covered = 1;
+      while (covered < node.min && body.complete && uncut(alts)) {
+        const joined = crossPrefixes(alts, body.alts);
+        if (joined === null) break;
+        alts = joined;
+        covered++;
+      }
+      const complete =
+        body.complete && covered === node.min && node.min === node.max && uncut(alts);
+      return { alts, complete };
+    }
+  }
+}
+
 /** Whether a node can match without consuming (assertions included). */
 function canMatchEmpty(node: RxNode): boolean {
   switch (node.kind) {
@@ -721,16 +805,74 @@ class RxGen {
     return wide;
   }
 
+  /** The vector plans of a start filter (prefixOf): per alternative, the
+   * plans of its first sets up to the first set without one. Null unless
+   * every alternative keeps at least two positions, where the filter is
+   * more selective than the first character alone. */
+  prefixPlans(prefix: Prefix | null): { ranges: number[][]; negate: boolean }[][] | null {
+    if (prefix === null || !this.vector) return null;
+    const plans: { ranges: number[][]; negate: boolean }[][] = [];
+    for (const alt of prefix.alts) {
+      const row: { ranges: number[][]; negate: boolean }[] = [];
+      for (const set of alt) {
+        const plan = this.vectorPlan(set);
+        if (plan === null) break;
+        row.push(plan);
+      }
+      if (row.length < 2) return null;
+      plans.push(row);
+    }
+    return plans.length === 0 ? null : plans;
+  }
+
+  /** Start positions s[index..index+16) where some alternative of the
+   * filter matches its first sets: an i16 mask. */
+  testPrefix(plans: { ranges: number[][]; negate: boolean }[][], index: string): string {
+    let any = "";
+    for (const row of plans) {
+      let all = "";
+      for (let k = 0; k < row.length; k++) {
+        let at = index;
+        if (k > 0) {
+          at = this.tmp();
+          this.line(`${at} = add i32 ${index}, ${k}`);
+        }
+        const bits = this.testVector(row[k]!, at);
+        if (all === "") all = bits;
+        else {
+          const both = this.tmp();
+          this.line(`${both} = and i16 ${all}, ${bits}`);
+          all = both;
+        }
+      }
+      if (any === "") any = all;
+      else {
+        const either = this.tmp();
+        this.line(`${either} = or i16 ${any}, ${all}`);
+        any = either;
+      }
+    }
+    return any;
+  }
+
   /** The first-character search: from %at, store the next position whose
-   * character is in `set` and branch to %attempt, or branch to %nomatch. */
-  emitSeek(set: CharSet): void {
+   * character is in `set` and branch to %attempt, or branch to %nomatch.
+   * With a start filter (`prefix`), the vector loop tests the filter's
+   * first sets of each alternative at 16 start positions at once and stops
+   * only where one of them matches all of its sets; a match can start
+   * nowhere else, so the attempts that remain are the ones that can
+   * succeed, in the same order. */
+  emitSeek(set: CharSet, prefix: Prefix | null = null): void {
     this.block("seek");
     const a = this.tmp();
     this.line(`${a} = load i32, ptr %at`);
+    const filter = this.prefixPlans(prefix);
     const plan = this.vectorPlan(set);
+    // The filter reads s[i + reach + 15] at its farthest.
+    const reach = filter === null ? 0 : Math.max(...filter.map((row) => row.length)) - 1;
     let from = "seek",
       start = a;
-    if (plan !== null) {
+    if (plan !== null || filter !== null) {
       const head = this.label("vsh"),
         load = this.label("vsl"),
         next = this.label("vsn"),
@@ -740,17 +882,21 @@ class RxGen {
       this.block(head);
       const i = this.tmp(),
         i16 = this.tmp(),
+        rest = this.tmp(),
         fits = this.tmp();
       this.line(`${i} = phi i32 [ ${a}, %seek ], [ ${i16}, %${next} ]`);
-      this.line(`${i16} = add i32 ${i}, 16`);
-      this.line(`${fits} = icmp sle i32 ${i16}, %len`);
+      // The characters left (0 <= i <= len, so this cannot wrap; i + 16
+      // can for subjects near 2^31 characters).
+      this.line(`${rest} = sub i32 %len, ${i}`);
+      this.line(`${fits} = icmp sge i32 ${rest}, ${16 + reach}`);
       this.line(`br i1 ${fits}, label %${load}, label %${tail}`);
       this.block(load);
-      const mask = this.testVector(plan, i),
+      const mask = filter !== null ? this.testPrefix(filter, i) : this.testVector(plan!, i),
         any = this.tmp();
       this.line(`${any} = icmp ne i16 ${mask}, 0`);
       this.line(`br i1 ${any}, label %${hit}, label %${next}`);
       this.block(next);
+      this.line(`${i16} = add i32 ${i}, 16`);
       this.line(`br label %${head}`);
       this.block(hit);
       const k = this.lowestBit(mask),
@@ -1001,10 +1147,10 @@ class RxGen {
       check = this.label("chk");
     this.block(entry);
     const p = this.loadPos(),
-      end = this.tmp(),
+      rest = this.tmp(),
       inb = this.tmp();
-    this.line(`${end} = add i32 ${p}, ${sets.length}`);
-    this.line(`${inb} = icmp sle i32 ${end}, %len`);
+    this.line(`${rest} = sub i32 %len, ${p}`);
+    this.line(`${inb} = icmp sge i32 ${rest}, ${sets.length}`);
     this.line(`br i1 ${inb}, label %${check}, label %fail`);
     this.block(check);
     let acc = "";
@@ -1023,6 +1169,8 @@ class RxGen {
         acc = both;
       }
     }
+    const end = this.tmp();
+    this.line(`${end} = add i32 ${p}, ${sets.length}`);
     // Storing before the branch is safe: every resume sets the position.
     this.storePos(end);
     this.line(`br i1 ${acc}, label %${succ}, label %fail`);
@@ -1229,10 +1377,11 @@ class RxGen {
       this.block(vhead);
       const kv = this.tmp(),
         kv16 = this.tmp(),
+        left = this.tmp(),
         fits = this.tmp();
       this.line(`${kv} = phi i32 [ 0, %${entry} ], [ ${kv16}, %${vnext} ]`);
-      this.line(`${kv16} = add i32 ${kv}, 16`);
-      this.line(`${fits} = icmp sle i32 ${kv16}, ${lim}`);
+      this.line(`${left} = sub i32 ${lim}, ${kv}`);
+      this.line(`${fits} = icmp sge i32 ${left}, 16`);
       this.line(`br i1 ${fits}, label %${vload}, label %${tail}`);
       this.block(vload);
       const vat = this.tmp();
@@ -1242,6 +1391,7 @@ class RxGen {
       this.line(`${all} = icmp eq i16 ${mask}, -1`);
       this.line(`br i1 ${all}, label %${vnext}, label %${vexit}`);
       this.block(vnext);
+      this.line(`${kv16} = add i32 ${kv}, 16`);
       this.line(`br label %${vhead}`);
       this.block(vexit);
       const missing = this.tmp();
@@ -1769,7 +1919,10 @@ export function compileRegexMatcher(
       fn.push(`  store i32 %off.n, ptr %at`);
       fn.push(`  br label %attempt`);
     } else {
-      gen.emitSeek(first.set);
+      // A start filter over the first few characters (a single first byte
+      // keeps memchr).
+      const prefix = prefixOf(root);
+      gen.emitSeek(first.set, gen.prefixPlans(prefix) === null ? null : prefix);
     }
   }
 

packages/compiler/test/regex-matcher-bounds.test.ts

@@ -0,0 +1,91 @@
+/* Compiled regex matchers bound every read by the characters left in the
+ * subject, so subjects near the runtime's 2^31 - 2 character limit are
+ * read only inside the subject. Each matcher runs on a subject of exactly
+ * that length whose last byte sits just before an inaccessible page: a
+ * bounds check that wraps past INT32_MAX reads that page and faults. The
+ * subject is a lazily mapped zero region (no memory is committed for it)
+ * with the characters each case needs written near its end. */
+import { execFileSync } from "node:child_process";
+import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { expect, test } from "vitest";
+import { compileRegexMatcher, REGEX_MATCHER_DECLS } from "../src/backend/llvm/regex-matchers.js";
+
+const LIMIT = 2 ** 31 - 2;
+
+/** [symbol, pattern, flags, start offset from the end, tail text]. */
+const CASES: [string, string, string, number, string][] = [
+  // The first-character vector search, 10 characters before the end.
+  ["seek", "[b-d]", "", 10, "c"],
+  // The start filter (two 8-character alternatives), 20 before the end.
+  ["filter", "agggtaaa|tttaccct", "i", 20, "TTTACCCT"],
+  // A run of 8 literal characters beginning 2 before the end.
+  ["run", "xy234567", "", 30, "xy"],
+  // A greedy class over the whole subject, 16 characters per step.
+  ["greedy", "[^x]+", "", LIMIT, ""],
+];
+
+const DRIVER = `
+#include <stdint.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <unistd.h>
+void scr_regex_stack_grow(void *a, void *b, void *c, int32_t d) { (void)a; (void)b; (void)c; (void)d; _exit(3); }
+void scr_regex_stack_free(void *p) { (void)p; }
+${CASES.map(([sym]) => `int32_t ${sym}(const char *, int32_t, int32_t, const char **);`).join("\n")}
+static void check(const char *name, int32_t (*m)(const char *, int32_t, int32_t, const char **),
+                int32_t start, const char *tail) {
+  const int32_t len = ${LIMIT};
+  size_t page = (size_t)sysconf(_SC_PAGESIZE);
+  size_t body = ((size_t)len + page - 1) / page * page;
+  char *base = mmap(NULL, body + page, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0);
+  if (base == MAP_FAILED || mprotect(base + body, page, PROT_NONE) != 0) _exit(4);
+  char *s = base + (body - (size_t)len);
+  memcpy(s + len - strlen(tail), tail, strlen(tail));
+  const char *out[16] = {0};
+  int32_t found = m(s, len, start, out);
+  if (found) printf("%s %ld %ld\\n", name, (long)(out[0] - s), (long)(out[1] - s));
+  else printf("%s none\\n", name);
+  fflush(stdout);
+  munmap(base, body + page);
+}
+int main(void) {
+${CASES.map(([sym, , , back, tail]) => `  check("${sym}", ${sym}, ${LIMIT - back}, ${JSON.stringify(tail)});`).join("\n")}
+  return 0;
+}
+`;
+
+test.skipIf(process.platform === "win32")(
+  "compiled regex matchers stay inside subjects near 2^31 characters",
+  { timeout: 120_000 },
+  () => {
+    const lines: string[] = [];
+    const decls = new Set(REGEX_MATCHER_DECLS);
+    for (const [sym, pattern, flags] of CASES) {
+      const code = compileRegexMatcher(sym, pattern, flags, "i64", true);
+      if (code === null) throw new Error(`/${pattern}/${flags} did not compile`);
+      lines.push(...code.lines.map((line) => line.replace(/^define internal /, "define ")));
+      for (const decl of code.decls) decls.add(decl);
+    }
+    const dir = mkdtempSync(join(tmpdir(), "scriptc-regex-bounds-"));
+    try {
+      const ir = join(dir, "matchers.ll");
+      const driver = join(dir, "driver.c");
+      const binary = join(dir, "bounds");
+      writeFileSync(ir, [...decls, ...lines, "attributes #0 = { nounwind }", ""].join("\n"));
+      writeFileSync(driver, DRIVER);
+      execFileSync("clang", ["-O2", "-Wno-override-module", ir, driver, "-o", binary]);
+      const output = execFileSync(binary, { encoding: "utf8" });
+      expect(output.trim().split("\n")).toEqual([
+        `seek ${LIMIT - 1} ${LIMIT}`,
+        `filter ${LIMIT - 8} ${LIMIT}`,
+        "run none",
+        `greedy 0 ${LIMIT}`,
+      ]);
+    } finally {
+      rmSync(dir, { recursive: true, force: true });
+    }
+  },
+);

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6872-regex-start-filter.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6872-regex-start-filter.ts"
+ ],
+ "diags": []
+}

tests/corpus/6872-regex-start-filter.ts

@@ -0,0 +1,79 @@
+// Searches whose patterns begin with fixed sequences of characters and
+// classes (literal alternations, classes, fixed repetitions), over subjects
+// of every length from 0 to 80 and with matches at every offset: match,
+// matchAll, search, exec with captures, global and sticky lastIndex,
+// replace, case-insensitive letters, multiline anchors and word boundaries.
+// Results must not depend on where a match falls relative to 16-character
+// blocks or the end of the subject.
+const patterns: RegExp[] = [
+  /agggtaaa|tttaccct/gi,
+  /[cgt]gggtaaa|tttaccc[acg]/gi,
+  /ab|abc|b/g,
+  /(ab)(c|d)x/g,
+  /x[0-9]{3}y/g,
+  /(?:no|yes)(!+)/gi,
+  /^ab/gm,
+  /\bcat\b/g,
+  /a{2}(b+)c/g,
+  /[A-Z][a-z]{2}-[0-9]/g,
+];
+
+function describe(re: RegExp, subject: string): string {
+  const parts: string[] = [];
+  const all = subject.match(re);
+  parts.push(`match=${all === null ? "null" : JSON.stringify(all)}`);
+  parts.push(`search=${subject.search(re)}`);
+  const rows: string[] = [];
+  for (const m of subject.matchAll(re)) rows.push(`${m.index}:${JSON.stringify(m.slice(0))}`);
+  parts.push(`all=[${rows.join(" ")}]`);
+  parts.push(`replace=${JSON.stringify(subject.replace(re, "<$&>"))}`);
+  return parts.join(" ");
+}
+
+const filler = "qwertyuiopasdfghjklzxcvbnm";
+const needles = ["agggtaaa", "TTTACCCT", "tttacccg", "abcx", "abdx", "x123y", "NO!!", "yes!", "cat", "aabbbc", "Abc-1", "ab"];
+let checksum = 0;
+let lines = 0;
+for (let length = 0; length <= 80; length += 7) {
+  for (const needle of needles) {
+    for (const offset of [0, 1, 14, 15, 16, 17, 31, 32, 33]) {
+      if (offset + needle.length > length) continue;
+      let subject = "";
+      while (subject.length < length) subject += filler;
+      subject = subject.slice(0, offset) + needle + subject.slice(offset + needle.length, length);
+      for (const re of patterns) {
+        const text = describe(re, subject);
+        for (let i = 0; i < text.length; i++) checksum = (checksum * 31 + text.charCodeAt(i)) >>> 0;
+        lines++;
+        if (length === 35 && (offset === 15 || offset === 16)) console.log(re.source, offset, needle, text);
+      }
+    }
+  }
+}
+console.log("lines", lines, "checksum", checksum);
+
+// Overlapping candidates: every start position looks like a match for the
+// first characters, and only some complete.
+const dense = "agggtaa".repeat(9) + "agggtaaa" + "tttaccc".repeat(5) + "tttaccct";
+console.log(describe(/agggtaaa|tttaccct/g, dense));
+
+// exec with lastIndex, global and sticky.
+const re = /x[0-9]{3}y/g;
+const text = "zzx111yzzzzzzzzzzzzzzzzzzx222yx333y" + "z".repeat(20) + "x444y";
+let m: RegExpExecArray | null;
+const found: string[] = [];
+while ((m = re.exec(text)) !== null) found.push(`${m[0]}@${m.index}->${re.lastIndex}`);
+console.log("exec", found.join(" "));
+const sticky = /x[0-9]{3}y/y;
+sticky.lastIndex = 2;
+console.log("sticky", sticky.test(text), sticky.lastIndex, sticky.test(text), sticky.lastIndex);
+
+// Short subjects and the tail.
+for (const subject of ["", "a", "agggtaa", "agggtaaa", "zagggtaaa", "z".repeat(16) + "agggtaaa", "z".repeat(23) + "tttaccct"]) {
+  console.log(JSON.stringify(subject.slice(-12)), subject.length, describe(/agggtaaa|tttaccct/gi, subject));
+}
+
+// A long subject with matches spread through it.
+let long = "";
+for (let i = 0; i < 2000; i++) long += i % 97 === 0 ? "aGgGtAaA" : i % 89 === 0 ? "tttACCCt" : filler[i % filler.length];
+console.log("long", (long.match(/agggtaaa|tttaccct/gi) ?? []).length, long.search(/tttaccct/i));

tests/harness/corpus-costs/6872-regex-start-filter.ts.json

@@ -0,0 +1 @@
+{ "plain": 2430, "san": 8860 }