tla-ts 0.0.2 → 0.0.3-dev.6979965

package/src/parallel-worker.ts

@@ -1,9 +1,20 @@
 import { createCanonicalizer } from "./canonical";
 import { createSearch } from "./check";
 import { createEnvironment } from "./environment";
-import { createFingerprintTable, getFingerprintOwner, writeStateFingerprint } from "./fingerprints";
-import type { LevelMessage, Port, StartMessage, WorkerMessage } from "./parallel";
-import { getSuccessors } from "./successors";
+import {
+  createFingerprintTable,
+  insertSharedFingerprint,
+  isSharedFingerprintsFull,
+  writeStateFingerprint,
+} from "./fingerprints";
+import type {
+  CoordinatorMessage,
+  LevelMessage,
+  Port,
+  StartMessage,
+  WorkerMessage,
+} from "./parallel";
+import { getSuccessors, type Successor } from "./successors";
 import { getTlcStringOrder, getTlcVariableOrder } from "./tlc-order";
 import { OUT_DEGREE_BUCKETS } from "./tlc-summary";
 import { createStateDecoder, createStateEncoder, type EncodedStates } from "./transfer";
@@ -11,21 +22,21 @@ import { getRankedStringCount, setStringOrder, type Value } from "./values";
 
 const WORKER_THREADS_MODULE = "node:worker_threads";
 
-interface Pending {
-  states: Value[][];
-  origins: number[];
-  fingerprints: number[];
-}
-
 interface PeerMessage {
-  phase: "query" | "answer" | "sizes" | "balance";
+  phase: "balance";
   level: number;
   from: number;
-  fingerprints?: Int32Array;
-  bits?: Uint8Array;
-  states?: EncodedStates;
-  unmarks?: Int32Array;
-  size?: number;
+  states: EncodedStates;
+}
+
+interface Expansion {
+  origin: number;
+  successors: Successor[] | undefined;
+  index: number;
+  generated: number;
+  newStates: Int32Array;
+  next: Value[][];
+  fingerprintStates: Value[][];
 }
 
 interface Transfer {
@@ -117,12 +128,13 @@ const run = (start: StartMessage, parent: Port) => {
   const nextAction = search?.spec.next;
   if (!search || nextAction === undefined)
     return parent.postMessage({ type: "abort" } satisfies WorkerMessage);
-  const count = start.count;
   const self = start.index;
   const canonicalize = createCanonicalizer();
   const decode = createStateDecoder(environment.variables.length);
-  const table = createFingerprintTable();
-  const known = createFingerprintTable();
+  const shared = { fingerprints: start.fingerprints };
+  // Under symmetry, a successor the worker generated before has the representative it had then,
+  // so the worker skips the search for one.
+  const generatedStates = search.isSymmetric ? createFingerprintTable() : undefined;
   const fingerprint = new Int32Array(2);
   const inbox = createInbox();
   const peers = start.ports.flatMap((port, index) => (port ? [{ port, index }] : []));
@@ -130,82 +142,44 @@ const run = (start: StartMessage, parent: Port) => {
     peer.port.addEventListener("message", (event) => inbox.deliver(event.data as PeerMessage));
     peer.port.start?.();
   });
-  for (let index = 0; index < start.seedFingerprints.length; index += 2) {
-    table.insert(start.seedFingerprints[index], start.seedFingerprints[index + 1]);
-  }
   const frontier = { states: decode(start.seed).map(canonicalize) };
   const send = (index: number, message: PeerMessage, transfer: unknown[]) =>
     start.ports[index]?.postMessage(message, transfer);
 
-  const expand = (states: Value[][], next: Value[][]) => {
-    const pending: Pending[] = Array.from({ length: count }, () => ({
-      states: [],
-      origins: [],
-      fingerprints: [],
-    }));
-    const statistics = { generated: 0 };
-    const newStates = new Int32Array(states.length);
-    for (let origin = 0; origin < states.length; origin++) {
-      const state = states[origin];
-      const successors = getSuccessors(nextAction, state, environment, search.spec.scope);
-      statistics.generated += successors.length;
-      if (successors.length === 0 && start.config.checkDeadlock) return undefined;
-      for (const successor of successors) {
+  // The worker pauses before a successor when the shared table needs to grow, and resumes there.
+  const expand = (expansion: Expansion) => {
+    const states = frontier.states;
+    for (; expansion.origin < states.length; expansion.origin++) {
+      const state = states[expansion.origin];
+      if (!expansion.successors) {
+        expansion.successors = getSuccessors(nextAction, state, environment, search.spec.scope);
+        expansion.generated += expansion.successors.length;
+        if (expansion.successors.length === 0 && start.config.checkDeadlock) return undefined;
+      }
+      const successors = expansion.successors;
+      for (; expansion.index < successors.length; expansion.index++) {
+        if (isSharedFingerprintsFull(shared.fingerprints)) return "paused";
+        const successor = successors[expansion.index];
         if (search.findViolatedActionProperty(state, successor.state)) return undefined;
         if (!search.satisfiesActionConstraints(state, successor.state)) continue;
-        canonicalize(successor.state);
-        writeStateFingerprint(search.getFingerprintState(successor.state), fingerprint);
-        const owner = getFingerprintOwner(fingerprint[1], count);
-        if (owner === self) {
-          if (!table.insert(fingerprint[0], fingerprint[1])) continue;
-          if (search.findViolatedInvariant(successor.state, state)) return undefined;
-          if (!search.satisfiesConstraints(successor.state, state)) {
-            table.remove(fingerprint[0], fingerprint[1]);
-            continue;
-          }
-          next.push(successor.state);
-          newStates[origin]++;
-          continue;
+        if (generatedStates) {
+          writeStateFingerprint(successor.state, fingerprint);
+          if (!generatedStates.insert(fingerprint[0], fingerprint[1])) continue;
         }
-        if (!known.insert(fingerprint[0], fingerprint[1])) continue;
-        pending[owner].states.push(successor.state);
-        pending[owner].origins.push(origin);
-        pending[owner].fingerprints.push(fingerprint[0], fingerprint[1]);
-      }
-    }
-    return { pending, generated: statistics.generated, newStates };
-  };
-
-  const answer = (fingerprints: Int32Array | number[]) => {
-    const bits = new Uint8Array(fingerprints.length / 2);
-    for (let index = 0; index < bits.length; index++) {
-      bits[index] = table.insert(fingerprints[index * 2], fingerprints[index * 2 + 1]) ? 1 : 0;
-    }
-    return bits;
-  };
-
-  const commit = (pending: Pending, bits: Uint8Array, next: Value[][], newStates: Int32Array) => {
-    const unmarks: number[] = [];
-    for (let index = 0; index < bits.length; index++) {
-      if (bits[index] === 0) continue;
-      const state = pending.states[index];
-      const parentState = frontier.states[pending.origins[index]];
-      if (search.findViolatedInvariant(state, parentState)) return undefined;
-      if (search.satisfiesConstraints(state, parentState)) {
-        next.push(state);
-        newStates[pending.origins[index]]++;
-        continue;
+        canonicalize(successor.state);
+        const fingerprintState = search.getFingerprintState(successor.state);
+        writeStateFingerprint(fingerprintState, fingerprint);
+        if (!insertSharedFingerprint(shared.fingerprints, fingerprint[0], fingerprint[1])) continue;
+        if (search.findViolatedInvariant(successor.state, state)) return undefined;
+        if (!search.satisfiesConstraints(successor.state, state)) continue;
+        expansion.next.push(successor.state);
+        expansion.fingerprintStates.push(fingerprintState);
+        expansion.newStates[expansion.origin]++;
       }
-      known.remove(pending.fingerprints[index * 2], pending.fingerprints[index * 2 + 1]);
-      unmarks.push(pending.fingerprints[index * 2], pending.fingerprints[index * 2 + 1]);
-    }
-    return Int32Array.from(unmarks);
-  };
-
-  const removeAll = (fingerprints: Int32Array) => {
-    for (let index = 0; index < fingerprints.length; index += 2) {
-      table.remove(fingerprints[index], fingerprints[index + 1]);
+      expansion.successors = undefined;
+      expansion.index = 0;
     }
+    return "done";
   };
 
   const balance = async (level: number, next: Value[][], sizes: number[]) => {
@@ -222,60 +196,50 @@ const run = (start: StartMessage, parent: Port) => {
       send(index, { phase: "balance", level, from: self, states }, [states.words.buffer]);
     });
     const incoming = await inbox.receive("balance", level, peers.length);
-    incoming.forEach((message) => {
-      if (message.states) next.push(...decode(message.states).map(canonicalize));
-    });
+    incoming.forEach((message) => next.push(...decode(message.states).map(canonicalize)));
     return true;
   };
 
+  const waiters = { resume: () => {}, balance: (_sizes: number[]) => {} };
+  const report = (message: WorkerMessage) =>
+    parent.postMessage(message, message.type === "level" ? [message.fingerprints.buffer] : []);
+
   const runLevel = async (level: number): Promise<LevelMessage | undefined> => {
-    const next: Value[][] = [];
-    const expanded = expand(frontier.states, next);
-    if (!expanded || getRankedStringCount() !== startupStrings) return undefined;
-    peers.forEach((peer) => {
-      const fingerprints = Int32Array.from(expanded.pending[peer.index].fingerprints);
-      send(peer.index, { phase: "query", level, from: self, fingerprints }, [fingerprints.buffer]);
-    });
-    const queries = await inbox.receive("query", level, peers.length);
-    queries.forEach((query) => {
-      const bits = answer(query.fingerprints ?? []);
-      send(query.from, { phase: "answer", level, from: self, bits }, [bits.buffer]);
-    });
-    const answers = await inbox.receive("answer", level, peers.length);
-    for (const reply of answers) {
-      const unmarks = commit(
-        expanded.pending[reply.from],
-        reply.bits ?? new Uint8Array(),
-        next,
-        expanded.newStates,
-      );
-      if (!unmarks) return undefined;
-      send(reply.from, { phase: "sizes", level, from: self, unmarks }, [unmarks.buffer]);
+    const expansion: Expansion = {
+      origin: 0,
+      successors: undefined,
+      index: 0,
+      generated: 0,
+      newStates: new Int32Array(frontier.states.length),
+      next: [],
+      fingerprintStates: [],
+    };
+    for (let status = expand(expansion); status !== "done"; status = expand(expansion)) {
+      if (status === undefined) return undefined;
+      const resumed = new Promise<void>((resolve) => (waiters.resume = resolve));
+      report({ type: "paused" });
+      await resumed;
     }
+    if (getRankedStringCount() !== startupStrings) return undefined;
+    const next = expansion.next;
     const added = next.length;
     const outDegrees = Array.from({ length: OUT_DEGREE_BUCKETS }, () => 0);
-    expanded.newStates.forEach((degree) => outDegrees[Math.min(degree, OUT_DEGREE_BUCKETS - 1)]++);
-    const fingerprints = BigInt64Array.from(next, (state) =>
-      environment.runtime.getFingerprint(state),
-    );
-    const sizes = Array.from({ length: count }, () => 0);
-    sizes[self] = added;
-    peers.forEach((peer) =>
-      send(peer.index, { phase: "sizes", level, from: self, size: added }, []),
-    );
-    const received = await inbox.receive("sizes", level, peers.length * 2);
-    received.forEach((message) => {
-      if (message.unmarks) removeAll(message.unmarks);
-      else sizes[message.from] = message.size ?? 0;
-    });
-    if (getRankedStringCount() !== startupStrings) return undefined;
-    if (!(await balance(level, next, sizes))) return undefined;
+    expansion.newStates.forEach((degree) => outDegrees[Math.min(degree, OUT_DEGREE_BUCKETS - 1)]++);
+    const fingerprints = search.getStoredFingerprints(next, expansion.fingerprintStates);
+    const sizes = new Promise<number[]>((resolve) => (waiters.balance = resolve));
+    report({ type: "expanded", from: self, added });
+    if (!(await balance(level, next, await sizes))) return undefined;
     frontier.states = next;
-    return { type: "level", level, generated: expanded.generated, added, outDegrees, fingerprints };
+    return {
+      type: "level",
+      level,
+      generated: expansion.generated,
+      added,
+      outDegrees,
+      fingerprints,
+    };
   };
 
-  const report = (message: WorkerMessage) =>
-    parent.postMessage(message, message.type === "level" ? [message.fingerprints.buffer] : []);
   const step = (level: number) => {
     runLevel(level).then(
       (message) => report(message ?? { type: "abort" }),
@@ -283,8 +247,13 @@ const run = (start: StartMessage, parent: Port) => {
     );
   };
   parent.addEventListener("message", (event) => {
-    const message = event.data as { type: string; level: number };
+    const message = event.data as CoordinatorMessage;
     if (message.type === "level") step(message.level);
+    if (message.type === "resize") {
+      shared.fingerprints = message.fingerprints;
+      waiters.resume();
+    }
+    if (message.type === "balance") waiters.balance(message.sizes);
   });
   report({ type: "ready" });
 };

package/src/parallel.ts

@@ -11,7 +11,14 @@ import { parseConfig, type ModelConfig } from "./config";
 import { createEnvironment } from "./environment";
 import { TlaEvaluationError } from "./errors";
 import { isImpure } from "./evaluate";
-import { createFingerprintTable, getFingerprintOwner, writeStateFingerprint } from "./fingerprints";
+import {
+  createSharedFingerprints,
+  growSharedFingerprints,
+  insertSharedFingerprint,
+  isSharedFingerprintsFull,
+  writeStateFingerprint,
+  type SharedFingerprints,
+} from "./fingerprints";
 import { getInitialStates } from "./successors";
 import { getTlcStringOrder, getTlcVariableOrder } from "./tlc-order";
 import { getFingerprintDistance, OUT_DEGREE_BUCKETS } from "./tlc-summary";
@@ -34,10 +41,9 @@ export interface StartMessage {
   modules: [string, string | undefined][];
   classes: [string, Uint8Array | undefined][];
   index: number;
-  count: number;
   ports: (Port | undefined)[];
   seed: EncodedStates;
-  seedFingerprints: Int32Array;
+  fingerprints: SharedFingerprints;
 }
 
 export interface LevelMessage {
@@ -49,7 +55,17 @@ export interface LevelMessage {
   fingerprints: BigInt64Array;
 }
 
-export type WorkerMessage = { type: "ready" } | { type: "abort" } | LevelMessage;
+export type WorkerMessage =
+  | { type: "ready" }
+  | { type: "abort" }
+  | { type: "paused" }
+  | { type: "expanded"; from: number; added: number }
+  | LevelMessage;
+
+export type CoordinatorMessage =
+  | { type: "level"; level: number }
+  | { type: "resize"; fingerprints: SharedFingerprints }
+  | { type: "balance"; sizes: number[] };
 
 interface WorkerHandle {
   postMessage: (message: unknown, transfer?: unknown[]) => void;
@@ -90,6 +106,7 @@ export const check = async (options: CheckOptions): Promise<CheckResult> => {
   const count = options.workers ?? 1;
   const isSequential =
     count < 2 ||
+    typeof SharedArrayBuffer !== "function" ||
     config.view !== undefined ||
     config.postconditions.length > 0 ||
     config.possible.length > 0;
@@ -131,12 +148,9 @@ export const check = async (options: CheckOptions): Promise<CheckResult> => {
   }
 
   const canonicalize = createCanonicalizer();
-  const table = createFingerprintTable();
+  const shared = { fingerprints: createSharedFingerprints() };
   const fingerprint = new Int32Array(2);
-  const seeds = Array.from({ length: count }, () => ({
-    encoder: createStateEncoder(),
-    fingerprints: [] as number[],
-  }));
+  const seeds = Array.from({ length: count }, createStateEncoder);
   const statistics = {
     generatedStates: 0,
     distinctStates: 0,
@@ -169,17 +183,15 @@ export const check = async (options: CheckOptions): Promise<CheckResult> => {
       canonicalize(state);
       statistics.generatedStates++;
       writeStateFingerprint(search.getFingerprintState(state), fingerprint);
-      if (!table.insert(fingerprint[0], fingerprint[1])) continue;
+      if (isSharedFingerprintsFull(shared.fingerprints)) {
+        shared.fingerprints = growSharedFingerprints(shared.fingerprints);
+      }
+      if (!insertSharedFingerprint(shared.fingerprints, fingerprint[0], fingerprint[1])) continue;
       if (search.findViolatedInvariant(state) || search.findFailedInitialPredicate(state)) {
         return false;
       }
-      if (!search.satisfiesConstraints(state)) {
-        table.remove(fingerprint[0], fingerprint[1]);
-        continue;
-      }
-      const seed = seeds[getFingerprintOwner(fingerprint[1], count)];
-      if (!seed.encoder.encode(state)) return false;
-      seed.fingerprints.push(fingerprint[0], fingerprint[1]);
+      if (!search.satisfiesConstraints(state)) continue;
+      if (!seeds[seeded.length % count].encode(state)) return false;
       seeded.push(state);
       statistics.distinctStates++;
     }
@@ -200,18 +212,19 @@ export const check = async (options: CheckOptions): Promise<CheckResult> => {
   if (statistics.distinctStates === 0) return createSuccess();
 
   const reports: WorkerMessage[] = [];
-  const pending = { resolve: (_reports: WorkerMessage[]) => {} };
+  const pending = { count, resolve: (_reports: WorkerMessage[]) => {} };
   const onMessage = (message: WorkerMessage) => {
     reports.push(message);
-    if (message.type === "abort" || reports.length === count) {
+    if (message.type === "abort" || reports.length === pending.count) {
       pending.resolve(reports.splice(0));
     }
   };
-  const waitForReports = () =>
+  const waitForReports = (expected: number) =>
     new Promise<WorkerMessage[]>((resolve) => {
+      pending.count = expected;
       pending.resolve = resolve;
     });
-  const ready = waitForReports();
+  const ready = waitForReports(count);
   const workers = await Promise.all(Array.from({ length: count }, () => spawnWorker(onMessage)));
   const terminate = () => workers.forEach((worker) => worker.terminate());
   const channels = new Map<string, MessageChannel>();
@@ -225,8 +238,7 @@ export const check = async (options: CheckOptions): Promise<CheckResult> => {
     const ports = Array.from({ length: count }, (_, peer) =>
       peer === index ? undefined : getPort(index, peer),
     );
-    const seed = seeds[index].encoder.finish();
-    const seedFingerprints = Int32Array.from(seeds[index].fingerprints);
+    const seed = seeds[index].finish();
     const message: StartMessage = {
       type: "start",
       source: options.source,
@@ -234,26 +246,44 @@ export const check = async (options: CheckOptions): Promise<CheckResult> => {
       modules: [...modules],
       classes: [...classes],
       index,
-      count,
       ports,
       seed,
-      seedFingerprints,
+      fingerprints: shared.fingerprints,
     };
-    worker.postMessage(message, [
-      ...ports.filter((port) => port !== undefined),
-      seed.words.buffer,
-      seedFingerprints.buffer,
-    ]);
+    worker.postMessage(message, [...ports.filter((port) => port !== undefined), seed.words.buffer]);
   });
   const isAborted = (messages: WorkerMessage[]) =>
     messages.some((message) => message.type === "abort");
+  const send = (message: CoordinatorMessage) =>
+    workers.forEach((worker) => worker.postMessage(message));
+  // Workers pause when the shared table needs to grow, and resume once it has.
+  const collectExpansions = async (messages: WorkerMessage[]): Promise<WorkerMessage[]> => {
+    const paused = messages.filter((message) => message.type === "paused").length;
+    if (paused === 0 || isAborted(messages)) return messages;
+    shared.fingerprints = growSharedFingerprints(shared.fingerprints);
+    const resumed = waitForReports(paused);
+    send({ type: "resize", fingerprints: shared.fingerprints });
+    const expanded = messages.filter((message) => message.type !== "paused");
+    return [...expanded, ...(await collectExpansions(await resumed))];
+  };
   if (isAborted(await ready)) {
     terminate();
     return fallBack();
   }
   for (let level = 1; ; level++) {
-    const reported = waitForReports();
-    workers.forEach((worker) => worker.postMessage({ type: "level", level }));
+    const expanding = waitForReports(count);
+    send({ type: "level", level });
+    const expansions = await collectExpansions(await expanding);
+    if (isAborted(expansions)) {
+      terminate();
+      return fallBack();
+    }
+    const sizes = Array.from({ length: count }, () => 0);
+    expansions.forEach((message) => {
+      if (message.type === "expanded") sizes[message.from] = message.added;
+    });
+    const reported = waitForReports(count);
+    send({ type: "balance", sizes });
     const messages = await reported;
     if (isAborted(messages)) {
       terminate();