tla-ts 0.0.1 → 0.0.3-dev.3cf2ad0

package/src/parallel.ts

@@ -1,7 +1,7 @@
 import type { Expression } from "./ast";
 import { createCanonicalizer } from "./canonical";
 import {
-  check,
+  checkSequentially,
   createSearch,
   findAssumptionViolation,
   type CheckOptions,
@@ -9,11 +9,21 @@ import {
 } from "./check";
 import { parseConfig, type ModelConfig } from "./config";
 import { createEnvironment } from "./environment";
-import { getChildren, isImpureBuiltin, type Environment } from "./evaluate";
-import { createFingerprintTable, getFingerprintOwner, getStateFingerprint } from "./fingerprints";
-import { parseModule } from "./parser";
+import { TlaEvaluationError } from "./errors";
+import { isImpure } from "./evaluate";
+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";
 import { createStateEncoder, type EncodedStates } from "./transfer";
+import { getRankedStringCount, setStringOrder, type Value } from "./values";
 
 const WORKER_THREADS_MODULE = "node:worker_threads";
 
@@ -29,11 +39,13 @@ export interface StartMessage {
   source: string;
   config: ModelConfig;
   modules: [string, string | undefined][];
+  classes: [string, Uint8Array | undefined][];
   index: number;
-  count: number;
   ports: (Port | undefined)[];
   seed: EncodedStates;
-  seedFingerprints: Int32Array;
+  fingerprints: SharedFingerprints;
+  requests: Int32Array;
+  remaining: Int32Array;
 }
 
 export interface LevelMessage {
@@ -41,13 +53,19 @@ export interface LevelMessage {
   level: number;
   generated: number;
   added: number;
+  outDegrees: number[];
+  fingerprints: BigInt64Array;
 }
 
-export type WorkerMessage = { type: "ready" } | { type: "abort" } | LevelMessage;
+export type WorkerMessage =
+  | { type: "ready" }
+  | { type: "abort" }
+  | { type: "paused" }
+  | LevelMessage;
 
-export interface ParallelCheckOptions extends CheckOptions {
-  workers: number;
-}
+export type CoordinatorMessage =
+  | { type: "level"; level: number }
+  | { type: "resize"; fingerprints: SharedFingerprints };
 
 interface WorkerHandle {
   postMessage: (message: unknown, transfer?: unknown[]) => void;
@@ -73,7 +91,7 @@ const spawnWorker = async (onMessage: (message: WorkerMessage) => void): Promise
     /* @vite-ignore */ /* webpackIgnore: true */ WORKER_THREADS_MODULE
   );
   const worker = new threads.Worker(new URL("./parallel-worker.js", import.meta.url));
-  worker.on("message", (message: WorkerMessage) => onMessage(message));
+  worker.on("message", onMessage);
   worker.on("error", () => onMessage({ type: "abort" }));
   return {
     postMessage: (message, transfer) =>
@@ -82,93 +100,137 @@ const spawnWorker = async (onMessage: (message: WorkerMessage) => void): Promise
   };
 };
 
-const usesImpureBuiltin = (expression: Expression, environment: Environment): boolean =>
-  ((expression.kind === "Apply" || expression.kind === "OperatorReference") &&
-    isImpureBuiltin(expression.operator, environment)) ||
-  (expression.kind === "Identifier" && isImpureBuiltin(expression.name, environment)) ||
-  getChildren(expression).some((child) => usesImpureBuiltin(child, environment));
-
-const isImpure = (environment: Environment) =>
-  [...environment.definitions.values()].some(
-    (definition) =>
-      usesImpureBuiltin(definition.body, environment) ||
-      (definition.kind === "Function" &&
-        definition.bounds.some((bound) => usesImpureBuiltin(bound.set, environment))),
-  );
-
-export const checkInParallel = async (options: ParallelCheckOptions): Promise<CheckResult> => {
+export const check = async (options: CheckOptions): Promise<CheckResult> => {
   const config = typeof options.config === "string" ? parseConfig(options.config) : options.config;
-  const checkSequentially = () => check({ ...options, config });
+  const fallBack = () => checkSequentially({ ...options, config });
+  const count = options.workers ?? 1;
   const isSequential =
-    options.workers < 2 ||
+    count < 2 ||
+    typeof SharedArrayBuffer !== "function" ||
     config.view !== undefined ||
     config.postconditions.length > 0 ||
     config.possible.length > 0;
-  if (isSequential) return checkSequentially();
+  if (isSequential) return fallBack();
   const modules = new Map<string, string | undefined>();
-  const specification = createEnvironment(parseModule(options.source), config, {
-    loadModule: (name) => {
-      const source = options.loadModule?.(name);
-      modules.set(name, source);
-      return source;
-    },
+  const loadModule = (name: string) => {
+    const source = options.loadModule?.(name);
+    modules.set(name, source);
+    return source;
+  };
+  const classes = new Map<string, Uint8Array | undefined>();
+  const loadClass = (name: string) => {
+    const bytes = options.loadClass?.(name);
+    classes.set(name, bytes);
+    return bytes;
+  };
+  setStringOrder(getTlcStringOrder(options.source, config, loadModule));
+  const startupStrings = getRankedStringCount();
+  const specification = createEnvironment(config, {
+    source: options.source,
+    loadModule,
+    loadClass,
     print: options.onPrint ?? console.log,
     files: { readFile: options.readFile, writeFile: options.writeFile },
   });
   const environment = specification.environment;
-  if (findAssumptionViolation(specification.assumptions, environment)) return checkSequentially();
-  const search = createSearch(config, environment);
-  if (!search || search.properties.liveness.length > 0 || isImpure(environment)) {
-    return checkSequentially();
+  const search = createSearch(config, environment, () =>
+    getTlcVariableOrder(options.source, loadModule),
+  );
+  if (findAssumptionViolation(specification.assumptions, environment)) return fallBack();
+  if (
+    !search ||
+    search.spec.init === undefined ||
+    search.spec.next === undefined ||
+    search.properties.liveness.length > 0 ||
+    isImpure(environment)
+  ) {
+    return fallBack();
   }
 
-  const count = options.workers;
   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 statistics = { generatedStates: 0, distinctStates: 0, depth: 0 };
-  for (const state of getInitialStates(search.spec.init, environment)) {
-    canonicalize(state);
-    statistics.generatedStates++;
-    getStateFingerprint(search.getFingerprintState(state), fingerprint);
-    if (!table.insert(fingerprint[0], fingerprint[1])) continue;
-    if (search.findViolatedInvariant(state) || search.findFailedInitialPredicate(state)) {
-      return checkSequentially();
-    }
-    if (!search.satisfiesConstraints(state)) {
-      table.remove(fingerprint[0], fingerprint[1]);
-      continue;
+  const seeds = Array.from({ length: count }, createStateEncoder);
+  const statistics = {
+    generatedStates: 0,
+    distinctStates: 0,
+    queuedStates: 0,
+    depth: 0,
+    outDegrees: Array.from({ length: OUT_DEGREE_BUCKETS }, () => 0),
+  };
+  const tlcFingerprints: BigInt64Array[] = [];
+  const createSuccess = (): CheckResult => ({
+    status: "success",
+    violation: undefined,
+    failure: undefined,
+    ...statistics,
+    fingerprintDistance: getFingerprintDistance(
+      statistics.distinctStates,
+      statistics.generatedStates,
+      () => {
+        const fingerprints = new BigInt64Array(statistics.distinctStates);
+        tlcFingerprints.reduce((offset, level) => {
+          fingerprints.set(level, offset);
+          return offset + level.length;
+        }, 0);
+        return fingerprints;
+      },
+    ),
+  });
+  const seedInitialStates = (init: Expression) => {
+    const seeded: Value[][] = [];
+    for (const state of getInitialStates(init, environment, search.spec.scope)) {
+      canonicalize(state);
+      statistics.generatedStates++;
+      writeStateFingerprint(search.getFingerprintState(state), fingerprint);
+      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)) continue;
+      if (!seeds[seeded.length % count].encode(state)) return false;
+      seeded.push(state);
+      statistics.distinctStates++;
     }
-    const seed = seeds[getFingerprintOwner(fingerprint[1], count)];
-    if (!seed.encoder.encode(state)) return checkSequentially();
-    seed.fingerprints.push(fingerprint[0], fingerprint[1]);
-    statistics.distinctStates++;
+    tlcFingerprints.push(
+      BigInt64Array.from(seeded, (state) => environment.runtime.getFingerprint(state)),
+    );
+    return true;
+  };
+  // The sequential search reports a violation or a failed evaluation in the initial states.
+  try {
+    if (!seedInitialStates(search.spec.init)) return fallBack();
+  } catch (error) {
+    if (!(error instanceof TlaEvaluationError)) throw error;
+    return fallBack();
   }
-  if (statistics.distinctStates === 0)
-    return { status: "success", violation: undefined, ...statistics };
-  statistics.depth = 1;
+  if (getRankedStringCount() !== startupStrings) return fallBack();
+  statistics.depth = statistics.generatedStates === 0 ? 0 : 1;
+  if (statistics.distinctStates === 0) return createSuccess();
 
   const reports: WorkerMessage[] = [];
-  const pending = { resolve: (_: 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 requests = new Int32Array(new SharedArrayBuffer(count * 4));
+  const remaining = new Int32Array(new SharedArrayBuffer(count * 4));
   const workers = await Promise.all(Array.from({ length: count }, () => spawnWorker(onMessage)));
   const terminate = () => workers.forEach((worker) => worker.terminate());
   const channels = new Map<string, MessageChannel>();
-  const getPort = (from: number, to: number): Port => {
+  const getPort = (from: number, to: number) => {
     const key = from < to ? `${from}:${to}` : `${to}:${from}`;
     const channel = channels.get(key) ?? new MessageChannel();
     channels.set(key, channel);
@@ -178,46 +240,63 @@ export const checkInParallel = async (options: ParallelCheckOptions): Promise<Ch
     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,
       config,
       modules: [...modules],
+      classes: [...classes],
       index,
-      count,
       ports,
       seed,
-      seedFingerprints,
+      fingerprints: shared.fingerprints,
+      requests,
+      remaining,
     };
-    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 collectLevels = 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 reported = messages.filter((message) => message.type !== "paused");
+    return [...reported, ...(await collectLevels(await resumed))];
+  };
   if (isAborted(await ready)) {
     terminate();
-    return checkSequentially();
+    return fallBack();
   }
   for (let level = 1; ; level++) {
-    const reported = waitForReports();
-    workers.forEach((worker) => worker.postMessage({ type: "level", level }));
-    const messages = await reported;
+    const reported = waitForReports(count);
+    // Every worker is idle, so peers can ask any of them for states once the level starts.
+    requests.forEach((_, index) => Atomics.store(requests, index, 0));
+    send({ type: "level", level });
+    const messages = await collectLevels(await reported);
     if (isAborted(messages)) {
       terminate();
-      return checkSequentially();
+      return fallBack();
     }
     const levels = messages.flatMap((message) => (message.type === "level" ? [message] : []));
     const added = levels.reduce((sum, message) => sum + message.added, 0);
     statistics.generatedStates += levels.reduce((sum, message) => sum + message.generated, 0);
     statistics.distinctStates += added;
+    levels.forEach((message) => {
+      message.outDegrees.forEach(
+        (degreeCount, degree) => (statistics.outDegrees[degree] += degreeCount),
+      );
+      tlcFingerprints.push(message.fingerprints);
+    });
     if (added === 0) {
       terminate();
-      return { status: "success", violation: undefined, ...statistics };
+      return createSuccess();
     }
     statistics.depth++;
   }