tla-ts 0.0.3-dev.a50fc89 → 0.0.3-dev.b74f28c

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tla-ts",
-  "version": "0.0.3-dev.a50fc89",
+  "version": "0.0.3-dev.b74f28c",
   "description": "A TLA+ interpreter and model checker in TypeScript, 2× faster than TLC",
   "keywords": [
     "formal verification",

package/src/fingerprints.ts

@@ -1,12 +1,10 @@
 import { getValueFingerprint, getValueHash, type Value } from "./values";
 
 const EMPTY = 0;
-const TOMBSTONE = 1;
 const INITIAL_CAPACITY = 1 << 12;
 
 interface FingerprintTable {
   insert: (first: number, second: number) => boolean;
-  remove: (first: number, second: number) => void;
 }
 
 const mix = (hash: number) => {
@@ -26,15 +24,13 @@ export const writeStateFingerprint = (state: Value[], fingerprint: Int32Array) =
   fingerprint[1] = mix(second);
 };
 
-export const getFingerprintOwner = (second: number, count: number) => (second >>> 0) % count;
-
 export const createFingerprintTable = (): FingerprintTable => {
   const table = { slots: new Int32Array(INITIAL_CAPACITY * 2), used: 0 };
   const findSlot = (slots: Int32Array, first: number, second: number) => {
     const mask = slots.length / 2 - 1;
     for (let index = first & mask; ; index = (index + 1) & mask) {
       const stored = slots[index * 2];
-      if (stored === EMPTY && slots[index * 2 + 1] === EMPTY) return index;
+      if (stored === EMPTY) return index;
       if (stored === first && slots[index * 2 + 1] === second) return index;
     }
   };
@@ -59,11 +55,50 @@ export const createFingerprintTable = (): FingerprintTable => {
       table.used++;
       return true;
     },
-    remove: (first, second) => {
-      const slot = findSlot(table.slots, first, second);
-      if (table.slots[slot * 2] !== first) return;
-      table.slots[slot * 2] = EMPTY;
-      table.slots[slot * 2 + 1] = TOMBSTONE;
-    },
   };
 };
+
+// Workers share one table of 64-bit fingerprints. The first word is never 0, so 0 marks an empty slot.
+export interface SharedFingerprints {
+  slots: BigInt64Array;
+  used: Int32Array;
+}
+
+const INITIAL_SHARED_CAPACITY = 1 << 10;
+
+const insertSharedKey = (slots: BigInt64Array, key: bigint) => {
+  const mask = slots.length - 1;
+  for (let index = Number(key >> 32n) & mask; ; index = (index + 1) & mask) {
+    const stored = Atomics.load(slots, index);
+    if (stored === key) return false;
+    if (stored !== 0n) continue;
+    const previous = Atomics.compareExchange(slots, index, 0n, key);
+    if (previous === 0n) return true;
+    if (previous === key) return false;
+  }
+};
+
+export const createSharedFingerprints = (): SharedFingerprints => ({
+  slots: new BigInt64Array(new SharedArrayBuffer(INITIAL_SHARED_CAPACITY * 8)),
+  used: new Int32Array(new SharedArrayBuffer(4)),
+});
+
+export const insertSharedFingerprint = (
+  table: SharedFingerprints,
+  first: number,
+  second: number,
+) => {
+  const key = BigInt.asIntN(64, (BigInt(first) << 32n) | BigInt(second >>> 0));
+  if (!insertSharedKey(table.slots, key)) return false;
+  Atomics.add(table.used, 0, 1);
+  return true;
+};
+
+export const isSharedFingerprintsFull = (table: SharedFingerprints) =>
+  Atomics.load(table.used, 0) * 2 >= table.slots.length;
+
+export const growSharedFingerprints = (table: SharedFingerprints): SharedFingerprints => {
+  const slots = new BigInt64Array(new SharedArrayBuffer(table.slots.byteLength * 2));
+  for (const key of table.slots) if (key !== 0n) insertSharedKey(slots, key);
+  return { slots, used: table.used };
+};

package/src/parallel-worker.ts

@@ -1,65 +1,41 @@
 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";
 import { getRankedStringCount, setStringOrder, type Value } from "./values";
 
 const WORKER_THREADS_MODULE = "node:worker_threads";
+const LOCKED = -1;
 
-interface Pending {
+interface Expansion {
   states: Value[][];
-  origins: number[];
-  fingerprints: number[];
+  origin: number;
+  end: number;
+  successors: Successor[] | undefined;
+  index: number;
+  added: number;
+  generated: number;
+  outDegrees: number[];
+  next: Value[][];
+  fingerprintStates: Value[][];
 }
 
-interface PeerMessage {
-  phase: "query" | "answer" | "sizes" | "balance";
-  level: number;
-  from: number;
-  fingerprints?: Int32Array;
-  bits?: Uint8Array;
-  states?: EncodedStates;
-  unmarks?: Int32Array;
-  size?: number;
-}
-
-interface Transfer {
-  from: number;
-  to: number;
-  count: number;
-}
-
-const getBalancePlan = (sizes: number[]) => {
-  const total = sizes.reduce((sum, size) => sum + size, 0);
-  const average = total / sizes.length;
-  if (Math.max(...sizes) - Math.min(...sizes) <= Math.max(64, average / 4)) return [];
-  const targets = sizes.map(
-    (_, index) => Math.floor(average) + (index < total % sizes.length ? 1 : 0),
-  );
-  const donors = sizes.flatMap((size, index) =>
-    size > targets[index] ? [{ index, amount: size - targets[index] }] : [],
-  );
-  const receivers = sizes.flatMap((size, index) =>
-    size < targets[index] ? [{ index, amount: targets[index] - size }] : [],
-  );
-  const transfers: Transfer[] = [];
-  for (const donor of donors) {
-    for (const receiver of receivers) {
-      const moved = Math.min(donor.amount, receiver.amount);
-      if (moved === 0) continue;
-      transfers.push({ from: donor.index, to: receiver.index, count: moved });
-      donor.amount -= moved;
-      receiver.amount -= moved;
-    }
-  }
-  return transfers;
-};
-
 const getParentPort = async () => {
   if (typeof (globalThis as { postMessage?: unknown }).postMessage === "function") {
     return globalThis as unknown as Port;
@@ -70,33 +46,6 @@ const getParentPort = async () => {
   return threads.parentPort as unknown as Port;
 };
 
-const createInbox = () => {
-  const messages = new Map<string, PeerMessage[]>();
-  const waiters = new Map<string, { count: number; resolve: (messages: PeerMessage[]) => void }>();
-  const getKey = (phase: PeerMessage["phase"], level: number) => `${phase}:${level}`;
-  const settle = (key: string) => {
-    const waiter = waiters.get(key);
-    const received = messages.get(key) ?? [];
-    if (!waiter || received.length < waiter.count) return;
-    waiters.delete(key);
-    messages.delete(key);
-    waiter.resolve(received);
-  };
-  return {
-    deliver: (message: PeerMessage) => {
-      const key = getKey(message.phase, message.level);
-      messages.set(key, [...(messages.get(key) ?? []), message]);
-      settle(key);
-    },
-    receive: (phase: PeerMessage["phase"], level: number, count: number) =>
-      new Promise<PeerMessage[]>((resolve) => {
-        const key = getKey(phase, level);
-        waiters.set(key, { count, resolve });
-        settle(key);
-      }),
-  };
-};
-
 const run = (start: StartMessage, parent: Port) => {
   const modules = new Map(start.modules);
   const loadModule = (name: string) => modules.get(name);
@@ -117,165 +66,163 @@ 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 };
+  const requests = start.requests;
+  const remaining = start.remaining;
+  // 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 }] : []));
-  peers.forEach((peer) => {
-    peer.port.addEventListener("message", (event) => inbox.deliver(event.data as PeerMessage));
-    peer.port.start?.();
+  const waiters = { resume: () => {}, loot: (_states: EncodedStates) => {} };
+  start.ports.forEach((port) => {
+    port?.addEventListener("message", (event) => waiters.loot(event.data as EncodedStates));
+    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) {
+  // A worker that runs out of states asks a peer for some by writing its index + 1 to the peer's
+  // request word. The peer sends half of the states it hasn't started, keeps expanding, and wakes
+  // the workers waiting to ask it next.
+  const answer = (expansion: Expansion, request: number) => {
+    const thief = Atomics.load(requests, self) - 1;
+    const given = Math.floor((expansion.end - expansion.origin - 1) / 2);
+    const encoder = createStateEncoder();
+    const states = expansion.states.slice(expansion.end - given, expansion.end);
+    const isEncoded = given > 0 && states.every(encoder.encode);
+    if (isEncoded) expansion.end -= given;
+    const loot = isEncoded ? encoder.finish() : createStateEncoder().finish();
+    start.ports[thief]?.postMessage(loot, [loot.words.buffer]);
+    Atomics.store(requests, self, request);
+    Atomics.notify(requests, self);
+  };
+
+  // A worker answers any request before it pauses or runs out of states, and takes no more until
+  // it resumes.
+  const lock = (expansion: Expansion) => {
+    Atomics.store(remaining, self, 0);
+    if (Atomics.compareExchange(requests, self, 0, LOCKED) === 0) return;
+    answer(expansion, LOCKED);
+  };
+
+  const unlock = (expansion: Expansion) => {
+    Atomics.store(remaining, self, expansion.end - expansion.origin);
+    Atomics.store(requests, self, 0);
+  };
+
+  // The worker pauses before a successor when the shared table needs to grow, and resumes there.
+  const expand = (expansion: Expansion) => {
+    for (; expansion.origin < expansion.end; expansion.origin++) {
+      const state = expansion.states[expansion.origin];
+      if (!expansion.successors) {
+        if (Atomics.load(requests, self) > 0) answer(expansion, 0);
+        Atomics.store(remaining, self, expansion.end - expansion.origin - 1);
+        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]);
+        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.added++;
       }
+      expansion.outDegrees[Math.min(expansion.added, OUT_DEGREE_BUCKETS - 1)]++;
+      expansion.successors = undefined;
+      expansion.index = 0;
+      expansion.added = 0;
     }
-    return { pending, generated: statistics.generated, newStates };
+    return "done";
   };
 
-  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]]++;
+  // Each attempt asks the peer with the most states left. While another worker's request to that
+  // peer is open, the worker waits for the peer to answer it. The worker stops once no peer that
+  // takes requests has two states left.
+  const steal = async (expansion: Expansion) => {
+    const skipped = new Set<number>();
+    for (;;) {
+      const victim = start.ports.reduce<number | undefined>((best, port, index) => {
+        if (!port || skipped.has(index) || Atomics.load(remaining, index) < 2) return best;
+        return best === undefined || Atomics.load(remaining, index) > Atomics.load(remaining, best)
+          ? index
+          : best;
+      }, undefined);
+      if (victim === undefined) return false;
+      const request = Atomics.compareExchange(requests, victim, 0, self + 1);
+      if (request === LOCKED) skipped.add(victim);
+      if (request > 0) Atomics.wait(requests, victim, request);
+      if (request !== 0) continue;
+      const loot = await new Promise<EncodedStates>((resolve) => (waiters.loot = resolve));
+      if (loot.count === 0) {
+        skipped.add(victim);
         continue;
       }
-      known.remove(pending.fingerprints[index * 2], pending.fingerprints[index * 2 + 1]);
-      unmarks.push(pending.fingerprints[index * 2], pending.fingerprints[index * 2 + 1]);
+      expansion.states = decode(loot).map(canonicalize);
+      expansion.origin = 0;
+      expansion.end = expansion.states.length;
+      unlock(expansion);
+      return true;
     }
-    return Int32Array.from(unmarks);
   };
 
-  const removeAll = (fingerprints: Int32Array) => {
-    for (let index = 0; index < fingerprints.length; index += 2) {
-      table.remove(fingerprints[index], fingerprints[index + 1]);
-    }
-  };
-
-  const balance = async (level: number, next: Value[][], sizes: number[]) => {
-    const transfers = getBalancePlan(sizes);
-    if (transfers.length === 0) return true;
-    const outgoing = new Map(peers.map((peer) => [peer.index, createStateEncoder()]));
-    for (const transfer of transfers.filter((candidate) => candidate.from === self)) {
-      const encoder = outgoing.get(transfer.to);
-      const states = next.splice(next.length - transfer.count, transfer.count);
-      if (!encoder || !states.every(encoder.encode)) return false;
-    }
-    outgoing.forEach((encoder, index) => {
-      const states = encoder.finish();
-      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));
-    });
-    return true;
-  };
+  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 = {
+      states: frontier.states,
+      origin: 0,
+      end: frontier.states.length,
+      successors: undefined,
+      index: 0,
+      added: 0,
+      generated: 0,
+      outDegrees: Array.from({ length: OUT_DEGREE_BUCKETS }, () => 0),
+      next: [],
+      fingerprintStates: [],
+    };
+    for (;;) {
+      const status = expand(expansion);
+      if (status === undefined) return undefined;
+      lock(expansion);
+      if (status === "paused") {
+        const resumed = new Promise<void>((resolve) => (waiters.resume = resolve));
+        report({ type: "paused" });
+        await resumed;
+        unlock(expansion);
+        continue;
+      }
+      if (!(await steal(expansion))) break;
     }
-    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;
-    frontier.states = next;
-    return { type: "level", level, generated: expanded.generated, added, outDegrees, fingerprints };
+    frontier.states = expansion.next;
+    Atomics.store(remaining, self, frontier.states.length);
+    return {
+      type: "level",
+      level,
+      generated: expansion.generated,
+      added: expansion.next.length,
+      outDegrees: expansion.outDegrees,
+      fingerprints: search.getStoredFingerprints(expansion.next, expansion.fingerprintStates),
+    };
   };
 
-  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,9 +230,14 @@ 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();
+    }
   });
+  Atomics.store(remaining, self, frontier.states.length);
   report({ type: "ready" });
 };