fix(vom): add new texts and fix bugs

packages/vom/src/__tests__/visual-surface-priority.test.ts

@@ -0,0 +1,35 @@
+import { describe, expect, it } from "vitest";
+import {
+  compareVisualSurfacePriority,
+  selectHighestPriorityVisualSurfaces,
+} from "../visual-surface-priority";
+
+function surface(backendNodeId: number, area: number, label?: string) {
+  return {
+    backendNodeId,
+    frameId: "main",
+    visibleRect: { x: backendNodeId, y: 0, w: area, h: 1 },
+    ...(label ? { label } : {}),
+  };
+}
+
+describe("visual surface priority selection", () => {
+  it("matches a full priority sort when selecting a bounded prefix", () => {
+    const surfaces = Array.from({ length: 600 }, (_, index) => surface(index, index + 1));
+    surfaces[599] = surface(599, 1, "Important status");
+
+    const selected = selectHighestPriorityVisualSurfaces(surfaces, 512);
+    const expected = [...surfaces].sort(compareVisualSurfacePriority).slice(0, 512);
+
+    expect(selected).toEqual(expected);
+    expect(selected[0]?.backendNodeId).toBe(599);
+  });
+
+  it("normalizes empty, fractional, and unbounded limits", () => {
+    const surfaces = [surface(1, 1), surface(2, 2), surface(3, 3)];
+
+    expect(selectHighestPriorityVisualSurfaces(surfaces, 0)).toEqual([]);
+    expect(selectHighestPriorityVisualSurfaces(surfaces, 1.9)).toEqual([surface(3, 3)]);
+    expect(selectHighestPriorityVisualSurfaces(surfaces, Number.POSITIVE_INFINITY)).toHaveLength(3);
+  });
+});

packages/vom/src/index.ts

@@ -22,5 +22,6 @@ export type {
 } from "./types";
 export {
   compareVisualSurfacePriority,
+  selectHighestPriorityVisualSurfaces,
   type VisualSurfacePriorityInput,
 } from "./visual-surface-priority";

packages/vom/src/render.ts

@@ -11,7 +11,7 @@ import type {
   VomScene,
   VomVisualSurface,
 } from "./types";
-import { compareVisualSurfacePriority } from "./visual-surface-priority";
+import { selectHighestPriorityVisualSurfaces } from "./visual-surface-priority";
 
 const SKIP_ROLES = new Set(["generic", "none", "presentation", "inlinetextbox"]);
 
@@ -822,17 +822,18 @@ function selectVisualSurfaces(
   surfaces: readonly VomVisualSurface[],
   nodesById: ReadonlyMap<number, VomNode>,
 ): { selected: VomVisualSurface[]; eligibleCount: number } {
-  const selected: VomVisualSurface[] = [];
   let eligibleCount = 0;
-
-  for (const surface of surfaces) {
-    if (surface.parentId !== null && !nodesById.has(surface.parentId)) continue;
-    eligibleCount += 1;
-
-    selected.push(surface);
+  function* eligibleSurfaces(): Generator<VomVisualSurface> {
+    for (const surface of surfaces) {
+      if (surface.parentId !== null && !nodesById.has(surface.parentId)) continue;
+      eligibleCount += 1;
+      yield surface;
+    }
   }
-  selected.sort(compareVisualSurfacePriority);
-  return { selected, eligibleCount };
+  return {
+    selected: selectHighestPriorityVisualSurfaces(eligibleSurfaces(), MAX_VISUAL_SURFACE_REFS),
+    eligibleCount,
+  };
 }
 
 function emitVisualAppendix(state: RenderState): void {

packages/vom/src/visual-surface-priority.ts

@@ -27,3 +27,54 @@ export function compareVisualSurfacePriority(
     a.backendNodeId - b.backendNodeId
   );
 }
+
+function siftWorstUp<T extends VisualSurfacePriorityInput>(heap: T[], start: number): void {
+  let index = start;
+  while (index > 0) {
+    const parent = Math.floor((index - 1) / 2);
+    if (compareVisualSurfacePriority(heap[index], heap[parent]) <= 0) break;
+    [heap[index], heap[parent]] = [heap[parent], heap[index]];
+    index = parent;
+  }
+}
+
+function siftWorstDown<T extends VisualSurfacePriorityInput>(heap: T[], start: number): void {
+  let index = start;
+  while (true) {
+    const left = index * 2 + 1;
+    if (left >= heap.length) return;
+    const right = left + 1;
+    let worst = left;
+    if (right < heap.length && compareVisualSurfacePriority(heap[right], heap[left]) > 0) {
+      worst = right;
+    }
+    if (compareVisualSurfacePriority(heap[worst], heap[index]) <= 0) return;
+    [heap[index], heap[worst]] = [heap[worst], heap[index]];
+    index = worst;
+  }
+}
+
+/** Selects the best `limit` surfaces and returns them in priority order. */
+export function selectHighestPriorityVisualSurfaces<T extends VisualSurfacePriorityInput>(
+  surfaces: Iterable<T>,
+  limit: number,
+): T[] {
+  if (!Number.isFinite(limit)) return [...surfaces].sort(compareVisualSurfacePriority);
+  const normalizedLimit = Math.max(0, Math.floor(limit));
+  if (normalizedLimit === 0) return [];
+  if (Array.isArray(surfaces) && surfaces.length <= normalizedLimit) {
+    return [...surfaces].sort(compareVisualSurfacePriority);
+  }
+
+  const heap: T[] = [];
+  for (const surface of surfaces) {
+    if (heap.length < normalizedLimit) {
+      heap.push(surface);
+      siftWorstUp(heap, heap.length - 1);
+    } else if (compareVisualSurfacePriority(surface, heap[0]) < 0) {
+      heap[0] = surface;
+      siftWorstDown(heap, 0);
+    }
+  }
+  return heap.sort(compareVisualSurfacePriority);
+}