fix: find image placeholders in linear time

plans/fix-43-linear-image-placeholders.md

@@ -0,0 +1,13 @@
+# #43 — find image placeholders in linear time
+
+`/!\[([^\]]+)\]\(\/?__too_be_replaced_image_path__\)/g` retries every `![` and scans to the end each time, so a document with many `![` and no `]` takes quadratic time.
+
+## How
+- `src/core/imagePlaceholders.ts`:
+  - `findImagePlaceholders` returns the pattern's matches with one forward scan. A prompt holds no `]`, so the only candidate ending at a given `]` starts at the first `![` after the previous `]`.
+  - `replaceImagePlaceholders` rebuilds the document from those matches.
+- `plugin.ts` uses both in place of `matchAll` and `replace`, and replaces from the matches it already found.
+- `tests/imagePlaceholders.test.ts` keeps the old pattern as the oracle. On seeded generated documents, the matches and the replaced document are compared against it, and a long input without `]` must finish within a budget.
+
+## Left out
+- What counts as a placeholder: unchanged.

src/core/imagePlaceholders.ts

@@ -0,0 +1,56 @@
+/** One `![prompt](__too_be_replaced_image_path__)` in a document, the path optionally with a leading slash. */
+export interface ImagePlaceholder {
+  start: number;
+  end: number;
+  prompt: string;
+}
+
+const PLACEHOLDER_TARGETS = [
+  "(__too_be_replaced_image_path__)",
+  "(/__too_be_replaced_image_path__)",
+];
+
+const targetLengthAt = (markdown: string, index: number): number =>
+  PLACEHOLDER_TARGETS.find((target) => markdown.startsWith(target, index))
+    ?.length ?? 0;
+
+/**
+ * The matches of /!\[([^\]]+)\]\(\/?__too_be_replaced_image_path__\)/g, left to right, in linear time.
+ * A prompt holds no "]", so the only candidate ending at a given "]" starts at the first "![" after
+ * the previous "]"; the pattern itself retries every "![" and is quadratic without a closing "]".
+ */
+export function findImagePlaceholders(markdown: string): ImagePlaceholder[] {
+  const found: ImagePlaceholder[] = [];
+  let from = 0;
+  let open = markdown.indexOf("![");
+  let close = markdown.indexOf("]");
+  while (close !== -1) {
+    if (open !== -1 && open < from) open = markdown.indexOf("![", from);
+    const targetLength =
+      open !== -1 && open + 2 < close ? targetLengthAt(markdown, close + 1) : 0;
+    if (targetLength > 0) {
+      const end = close + 1 + targetLength;
+      found.push({ start: open, end, prompt: markdown.slice(open + 2, close) });
+      from = end;
+    } else {
+      from = close + 1;
+    }
+    close = markdown.indexOf("]", from);
+  }
+  return found;
+}
+
+/** `markdown` with each placeholder replaced by what `replacement` returns for it and its position in the list. */
+export function replaceImagePlaceholders(
+  markdown: string,
+  placeholders: readonly ImagePlaceholder[],
+  replacement: (placeholder: ImagePlaceholder, index: number) => string,
+): string {
+  const pieces: string[] = [];
+  const copiedUpTo = placeholders.reduce((position, placeholder, index) => {
+    pieces.push(markdown.slice(position, placeholder.start), replacement(placeholder, index));
+    return placeholder.end;
+  }, 0);
+  pieces.push(markdown.slice(copiedUpTo));
+  return pieces.join("");
+}

src/core/plugin.ts

@@ -6,6 +6,10 @@
   readBlankImageBase64,
 } from "./hostResponse";
 import { TOOL_NAME, TOOL_DEFINITION } from "./definition";
+import {
+  findImagePlaceholders,
+  replaceImagePlaceholders,
+} from "./imagePlaceholders";
 
 // context is nullable on purpose: hosts that run the plugin without client-side
 // state (MulmoClaude's server bridge) pass an empty or missing context, and
@@ -25,8 +29,7 @@
 
   // Detect all image placeholders in the format: ![prompt](__too_be_replaced_image_path__)
   // Also handle variant with leading slash: ![prompt](/__too_be_replaced_image_path__)
-  const imageRegex = /!\[([^\]]+)\]\(\/?__too_be_replaced_image_path__\)/g;
-  const matches = [...markdown.matchAll(imageRegex)];
+  const matches = findImagePlaceholders(markdown);
 
   if (matches.length > 0 && app?.generateImageWithBackend && app?.saveImages) {
     // Generate UUID for this document
@@ -40,11 +43,11 @@
 
     // Generate images for each placeholder in parallel
     const imagePromises = matches.map(async (match, i) => {
-      const prompt = `${match[1]}. Use the last image as the output dimension.`;
+      const prompt = `${match.prompt}. Use the last image as the output dimension.`;
       const imageId = `image_${i}`;
 
       try {
         const result = await app!.generateImageWithBackend!(
           prompt,
           blankImageBase64 ? [blankImageBase64] : [],
           context,
@@ -69,12 +72,12 @@
           const imageUrls = data.imageUrls;
 
           // Replace placeholders with actual image URLs
-          let imageIndex = 0;
-          markdown = markdown.replace(imageRegex, (match, prompt) => {
-            const imageId = `image_${imageIndex}`;
-            const imageUrl = imageUrls[imageId];
-            imageIndex++;
-            return imageUrl ? `![${prompt}](${imageUrl})` : match;
+          const original = markdown;
+          markdown = replaceImagePlaceholders(original, matches, (placeholder, index) => {
+            const imageUrl = imageUrls[`image_${index}`];
+            return imageUrl
+              ? `![${placeholder.prompt}](${imageUrl})`
+              : original.slice(placeholder.start, placeholder.end);
           });
         }
       } catch (error) {

tests/imagePlaceholders.test.ts

@@ -0,0 +1,94 @@
+import { test, describe } from "node:test";
+import assert from "node:assert";
+import {
+  findImagePlaceholders,
+  replaceImagePlaceholders,
+} from "../src/core/imagePlaceholders.js";
+
+// The pattern the scanner replaced: correct, but quadratic, so only an oracle for small inputs.
+const PLACEHOLDER_PATTERN =
+  /!\[([^\]]+)\]\(\/?__too_be_replaced_image_path__\)/g;
+
+const TOKENS = [
+  "!", "[", "]", "(", ")", "/", "\\", "a", " ", "\n", "![", "](", "]]", "[[",
+  "(__too_be_replaced_image_path__)", "(/__too_be_replaced_image_path__)",
+  "__too_be_replaced_image_path__", "(__too_be_replaced_image_path_)",
+];
+const GENERATED_DOCUMENTS = 5000;
+const MAX_TOKENS = 24;
+const SEED = 43;
+
+const lcg = (seed: number) => {
+  let state = seed >>> 0;
+  return (below: number): number => {
+    state = (Math.imul(state, 1664525) + 1013904223) >>> 0;
+    return state % below;
+  };
+};
+
+const generatedDocuments = (): string[] => {
+  const next = lcg(SEED);
+  return Array.from({ length: GENERATED_DOCUMENTS }, () =>
+    Array.from({ length: next(MAX_TOKENS + 1) }, () => TOKENS[next(TOKENS.length)]).join(""),
+  );
+};
+
+const patternMatches = (markdown: string) =>
+  [...markdown.matchAll(PLACEHOLDER_PATTERN)].map((match) => ({
+    start: match.index,
+    end: match.index + match[0].length,
+    prompt: match[1],
+  }));
+
+const urlFor = (index: number) => (index % 2 === 0 ? `/images/${index}.png` : undefined);
+
+describe("findImagePlaceholders", () => {
+  test("finds both forms of the placeholder with their prompts", () => {
+    const markdown = "a ![a cat](__too_be_replaced_image_path__) b ![a dog](/__too_be_replaced_image_path__)";
+    assert.deepEqual(
+      findImagePlaceholders(markdown).map(({ prompt }) => prompt),
+      ["a cat", "a dog"],
+    );
+  });
+
+  test("ignores an empty prompt, another path, and an ordinary image", () => {
+    ["![](__too_be_replaced_image_path__)", "![x](other.png)", "![x](__too_be_replaced_image_path__", "plain text", ""].forEach(
+      (markdown) => assert.deepEqual(findImagePlaceholders(markdown), []),
+    );
+  });
+
+  test("returns the pattern's matches on generated documents", () => {
+    generatedDocuments().forEach((markdown) =>
+      assert.deepEqual(findImagePlaceholders(markdown), patternMatches(markdown), JSON.stringify(markdown)),
+    );
+  });
+
+  test("stays fast on many '![' with no closing ']'", () => {
+    const LONG_INPUT_REPEATS = 200_000;
+    // The pattern took seconds at a fifth of this length; a linear scan takes milliseconds.
+    const BUDGET_MS = 1000;
+    const markdown = "![" + "![\\".repeat(LONG_INPUT_REPEATS) + "](__too_be_replaced_image_path__)";
+    const startedAt = performance.now();
+    const [only] = findImagePlaceholders(markdown);
+    assert.ok(performance.now() - startedAt < BUDGET_MS);
+    assert.equal(only.start, 0);
+    assert.equal(only.end, markdown.length);
+  });
+});
+
+describe("replaceImagePlaceholders", () => {
+  test("gives the same document as replacing through the pattern", () => {
+    generatedDocuments().forEach((markdown) => {
+      let index = 0;
+      const expected = markdown.replace(PLACEHOLDER_PATTERN, (match, prompt: string) => {
+        const url = urlFor(index++);
+        return url ? `![${prompt}](${url})` : match;
+      });
+      const actual = replaceImagePlaceholders(markdown, findImagePlaceholders(markdown), (placeholder, position) => {
+        const url = urlFor(position);
+        return url ? `![${placeholder.prompt}](${url})` : markdown.slice(placeholder.start, placeholder.end);
+      });
+      assert.equal(actual, expected, JSON.stringify(markdown));
+    });
+  });
+});