tla-ts 0.0.2 → 0.0.3-dev.a50fc89

package/src/run-cli.ts

@@ -1,10 +1,29 @@
-import { existsSync, mkdirSync, readFileSync, readSync, writeFileSync } from "node:fs";
-import { basename, dirname, join, resolve } from "node:path";
+import {
+  existsSync,
+  globSync,
+  mkdirSync,
+  readFileSync,
+  readSync,
+  realpathSync,
+  writeFileSync,
+} from "node:fs";
+import { spawnSync } from "node:child_process";
+import { basename, delimiter, dirname, join, resolve } from "node:path";
 import { availableParallelism } from "node:os";
 import type { CheckResult, CheckStatus, TraceStep } from "./check";
+import { defineCommand, type ParsedArgs } from "./cli-options";
 import { parseConfig } from "./config";
-import { getErrorExitCode, TlaConfigError, TlaError, TlaEvaluationError } from "./errors";
+import {
+  getErrorExitCode,
+  TlaConfigError,
+  TlaError,
+  TlaEvaluationError,
+  TlaProofError,
+} from "./errors";
 import { EXIT_STATUS } from "./exit-codes";
+import type { LineRange } from "./obligation-generate";
+import { getAtPrefix, getLocusText } from "./obligation-expression";
+import { formatSummary, getObligations, WHOLE_MODULE, type ModuleObligations } from "./obligations";
 import { check } from "./parallel";
 import { prettyPrint } from "./pretty-print";
 import {
@@ -26,17 +45,54 @@ const EXIT_CODES: Record<CheckStatus, number> = {
   error: EXIT_STATUS.error,
 };
 
-const USAGE =
-  "Usage: tla-ts <Spec.tla> [-config <Spec.cfg>] [-deadlock] [-lncheck <final|default>] [-seed <num>] [-workers <count|auto>]";
+const CHECK_COMMAND = defineCommand("tla-ts <Spec.tla|Spec.cfg|glob>... {options}", [
+  { flag: "-config <Spec.cfg>" },
+  { flag: "-deadlock" },
+  { flag: "-lncheck <final|default>" },
+  { flag: "-seed <num>" },
+  { flag: "-workers <count|auto>" },
+]);
+
+const PROVE_COMMAND = defineCommand("tla-ts prove {options} <Module.tla|glob>...", [
+  { flag: "-N", isRequired: true },
+  { flag: "--toolbox <start> <end>" },
+  { flag: "--line <num>" },
+  { flag: "--printallobs" },
+  { flag: "--summary" },
+  { flag: "--strict" },
+  { flag: "-I <directory>", isRepeatable: true },
+]);
+
+const USAGE = [
+  `Usage: ${CHECK_COMMAND.usage}`,
+  `       ${PROVE_COMMAND.usage}`,
+  "       tla-ts prove [<tlapm option>]... <Module.tla|glob>...",
+].join("\n");
+
+const GLOB_CHARACTERS = /[*?[{]/;
+
+interface InputPath {
+  path: string;
+  isGlobbed: boolean;
+}
+
+const getInputPaths = (paths: string[]): InputPath[] =>
+  paths.flatMap((arg): InputPath[] =>
+    GLOB_CHARACTERS.test(arg)
+      ? globSync(arg)
+          .sort()
+          .map((path) => ({ path, isGlobbed: true }))
+      : [{ path: arg, isGlobbed: false }],
+  );
 
 const isJavaLong = (text: string) =>
   /^[+-]?\d+$/.test(text) && BigInt.asIntN(64, BigInt(text)) === BigInt(text);
 
-const getWorkerCount = (args: string[]) => {
-  const flag = args.indexOf("-workers");
-  if (flag === -1) return 1;
-  const value = args[flag + 1];
-  return value === "auto" ? availableParallelism() : Math.max(1, Number.parseInt(value, 10) || 1);
+const getWorkerCount = (values: string[] | undefined) => {
+  if (!values) return 1;
+  return values[0] === "auto"
+    ? availableParallelism()
+    : Math.max(1, Number.parseInt(values[0], 10) || 1);
 };
 
 const printState = (state: Record<string, Value>) => {
@@ -130,26 +186,21 @@ const printResult = (result: CheckResult) => {
   if (outDegree) console.log(outDegree);
 };
 
-const run = async (args: string[]) => {
-  const specPath = args.find((arg, index) => arg.endsWith(".tla") && args[index - 1] !== "-config");
-  if (!specPath) {
-    console.error(USAGE);
-    return 1;
-  }
+const checkSpec = async (specPath: string, args: ParsedArgs) => {
   if (!existsSync(specPath)) {
     console.error(`Cannot find the specification file ${specPath}`);
     return EXIT_STATUS.configParseError;
   }
-  const configFlag = args.indexOf("-config");
-  const configPath = configFlag === -1 ? specPath.replace(/\.tla$/, ".cfg") : args[configFlag + 1];
+  const configPath =
+    args.get("-config")?.[0] ??
+    (args.has("-config") ? undefined : specPath.replace(/\.tla$/, ".cfg"));
   if (!configPath || !existsSync(configPath)) {
     console.error(`Cannot find the configuration file ${configPath ?? ""}`);
     return EXIT_STATUS.configParseError;
   }
   const config = parseConfig(readFileSync(configPath, "utf8"));
-  if (args.includes("-deadlock")) config.checkDeadlock = false;
-  const seedFlag = args.indexOf("-seed");
-  const seed = seedFlag === -1 ? "0" : args[seedFlag + 1];
+  if (args.has("-deadlock")) config.checkDeadlock = false;
+  const seed = args.has("-seed") ? args.get("-seed")?.[0] : "0";
   if (seed === undefined) {
     console.error("Error: seed required.");
     return 1;
@@ -158,8 +209,7 @@ const run = async (args: string[]) => {
     console.error(`Error: An integer for seed required. But encountered ${seed}`);
     return 1;
   }
-  const livenessFlag = args.indexOf("-lncheck");
-  const livenessCheck = livenessFlag === -1 ? "default" : args[livenessFlag + 1]?.toLowerCase();
+  const livenessCheck = args.has("-lncheck") ? args.get("-lncheck")?.[0]?.toLowerCase() : "default";
   if (livenessCheck === undefined) {
     console.error("Error: expect a strategy such as final for -lncheck option.");
     return 1;
@@ -167,7 +217,7 @@ const run = async (args: string[]) => {
   const directory = dirname(specPath);
   console.log(`Checking ${basename(specPath)} with ${basename(configPath)}`);
   const result = await check({
-    workers: getWorkerCount(args),
+    workers: getWorkerCount(args.get("-workers")),
     livenessCheck: ["final", "seqfinal", "off"].includes(livenessCheck) ? "final" : "default",
     seed: BigInt(seed),
     source: readFileSync(specPath, "utf8"),
@@ -199,12 +249,230 @@ const printError = (error: TlaError) => {
   console.error(`${error.name}: ${error.message}`);
 };
 
-export const runCli = async () => {
+// Prints a TlaError and returns its exit code.
+const reportErrors = async (run: () => number | Promise<number>) => {
   try {
-    process.exitCode = await run(process.argv.slice(2));
+    return await run();
   } catch (error) {
     if (!(error instanceof TlaError)) throw error;
     printError(error);
-    process.exitCode = getErrorExitCode(error);
+    return getErrorExitCode(error);
+  }
+};
+
+// Runs each path in turn, so an error in one leaves the others to report. The exit code is the
+// first that is not success.
+const runEach = async (paths: string[], runPath: (path: string) => number | Promise<number>) => {
+  const exitCodes: number[] = [];
+  for (const [index, path] of paths.entries()) {
+    if (index > 0) console.log("");
+    exitCodes.push(await reportErrors(() => runPath(path)));
+  }
+  return exitCodes.find((exitCode) => exitCode !== EXIT_STATUS.success) ?? EXIT_STATUS.success;
+};
+
+// tlapm installs its standard library, with modules such as TLAPS, at lib/tlapm/stdlib beside
+// bin/tlapm.
+const getTlapmPaths = () =>
+  (process.env.PATH ?? "")
+    .split(delimiter)
+    .map((directory) => join(directory, "tlapm"))
+    .filter((path) => existsSync(path));
+
+const getTlapmLibrary = () =>
+  getTlapmPaths()
+    .map((path) => join(dirname(dirname(realpathSync(path))), "lib/tlapm/stdlib"))
+    .find((path) => existsSync(path));
+
+const expandGlobs = (paths: string[]) =>
+  paths.flatMap((path) => (GLOB_CHARACTERS.test(path) ? globSync(path).sort() : [path]));
+
+const proveWithTlapm = (args: string[]) => {
+  const tlapm = getTlapmPaths()[0];
+  if (tlapm === undefined) {
+    console.error(
+      "Error: tla-ts prove runs tlapm to prove obligations, but tlapm isn't on your PATH. Install tlapm, or pass -N to only generate the obligations.",
+    );
+    return 1;
+  }
+  return spawnSync(tlapm, expandGlobs(args), { stdio: "inherit" }).status ?? 1;
+};
+
+const readModule = (path: string) =>
+  existsSync(path) ? { file: path, source: readFileSync(path, "utf8") } : undefined;
+
+// A message of tlapm's toolbox protocol, which it writes to stderr.
+const printToolboxMessage = (fields: string[]) =>
+  console.error(`@!!BEGIN\n${fields.map((field) => `@!!${field}\n`).join("")}@!!END\n`);
+
+const TLAPM_VERSION = "5f5037f";
+
+const quoteArgument = (arg: string) =>
+  /^[+,\-./0-9:=@A-Z_a-z]+$/.test(arg) ? arg : `'${arg.replaceAll("'", "'\\''")}'`;
+
+const printToolboxHeader = (args: string[]) => {
+  const now = new Date();
+  const pad = (value: number) => String(value).padStart(2, "0");
+  const date = `${now.getFullYear()}-${pad(now.getMonth() + 1)}-${pad(now.getDate())}`;
+  const time = `${pad(now.getHours())}:${pad(now.getMinutes())}:${pad(now.getSeconds())}`;
+  console.log(
+    `\n\\* TLAPM version ${TLAPM_VERSION}\n\\* launched at ${date} ${time} with command line:\n\\* ${["tla-ts", "prove", ...args].map(quoteArgument).join(" ")}\n`,
+  );
+};
+
+const STRICT_EXIT_STATUS = { incomplete: 11, emptyTarget: 12 };
+
+const checkStrictly = (result: ModuleObligations, isTargeted: boolean, isSummary: boolean) => {
+  const summary = result.summary;
+  if (!isSummary) {
+    summary.absent.forEach((locus) => console.error(`Missing proof at ${getLocusText(locus)}`));
+    summary.omitted.forEach((locus) => console.error(`Omitted proof at ${getLocusText(locus)}`));
+  }
+  const isIncomplete = summary.absent.length + summary.omitted.length > 0;
+  if (isIncomplete) {
+    console.error(
+      `${getAtPrefix(result.location)}[ERROR]: Proof incomplete in module ${JSON.stringify(result.name)}: ${summary.absent.length} missing, ${summary.omitted.length} omitted proof step(s).`,
+    );
   }
+  const isEmpty = isTargeted && result.obligations.length === 0;
+  if (isEmpty) {
+    console.error(
+      `${getAtPrefix(result.location)}[ERROR]: No proof obligation found for the selected target.`,
+    );
+  }
+  return [
+    ...(isIncomplete ? [STRICT_EXIT_STATUS.incomplete] : []),
+    ...(isEmpty ? [STRICT_EXIT_STATUS.emptyTarget] : []),
+  ];
+};
+
+const TLAPM_FAILURE_STATUS = 3;
+
+const generateObligations = (generate: () => ModuleObligations[]) => {
+  try {
+    return generate();
+  } catch (error) {
+    if (!(error instanceof TlaProofError)) throw error;
+    console.error(` tlapm ending abnormally with ${error.exception}\n`);
+    return undefined;
+  }
+};
+
+const prove = (args: string[]) => {
+  if (!args.includes("-N") && !args.includes("--summary")) return proveWithTlapm(args);
+  const parsed = PROVE_COMMAND.parse(args);
+  const unsupported = parsed.unknown[0];
+  if (unsupported !== undefined) {
+    console.error(`Error: tla-ts prove does not support the tlapm option ${unsupported}.`);
+    return 1;
+  }
+  if (parsed.incomplete) {
+    console.error(`Error: option '${parsed.incomplete.name}' needs an argument.`);
+    return 1;
+  }
+  const nonInteger = parsed.options.find(
+    (option) => option.name !== "-I" && option.values.some((value) => !/^-?\d+$/.test(value)),
+  );
+  if (nonInteger) {
+    console.error(`Error: option '${nonInteger.name}' expects an integer.`);
+    return 1;
+  }
+  const files = expandGlobs(parsed.positionals);
+  const modulePaths = files.map((path) => (path.endsWith(".tla") ? path : `${path}.tla`));
+  if (modulePaths.length === 0) {
+    console.error(USAGE);
+    return 1;
+  }
+  const isToolbox = parsed.has("--toolbox") || parsed.has("--line");
+  if (isToolbox) printToolboxHeader(args);
+  // Like tlapm, the last of --toolbox and --line wins, and a 0 means the whole module.
+  const lines = parsed.options.reduce<LineRange>((range, option) => {
+    const [start, end] = option.values.map(Number);
+    if (option.name === "--toolbox") return { start, end: end === 0 ? Infinity : end };
+    if (option.name === "--line") return start === 0 ? WHOLE_MODULE : { start, end: start };
+    return range;
+  }, WHOLE_MODULE);
+  // tlapm searches the current directory and each module's, then its library, then each -I
+  // directory, the last given first.
+  const tlapmLibrary = getTlapmLibrary();
+  const includes = parsed.options.flatMap((option) => (option.name === "-I" ? option.values : []));
+  const directories = modulePaths.reduce(
+    (found, path) => (found.includes(dirname(path)) ? found : [dirname(path), ...found]),
+    [".", ...(tlapmLibrary ? [tlapmLibrary] : []), ...includes.toReversed()],
+  );
+  // tlapm names each file it reads by joining a search directory and the name, as in ./Module.tla.
+  const results = generateObligations(() => {
+    const missingIndex = modulePaths.findIndex((path) => !existsSync(path));
+    if (missingIndex !== -1) {
+      const missing = files[missingIndex];
+      console.error(`File "<unknown>":\nError: File ${JSON.stringify(missing)} not found`);
+      throw new TlaProofError(`File ${missing} not found`, { failure: "Module.Parser.parse_file" });
+    }
+    return getObligations(
+      modulePaths.map((path) => ({ file: `./${path}`, source: readFileSync(path, "utf8") })),
+      (name) =>
+        directories.map((directory) => readModule(`${directory}/${name}.tla`)).find(Boolean),
+      {
+        lines,
+        onWarning: (warning) => {
+          if (warning.channel === "stdout") return console.log(warning.text);
+          const isToolboxMessage =
+            warning.channel === "toolbox" || warning.channel === "toolboxOnly";
+          if (isToolbox && isToolboxMessage)
+            return printToolboxMessage(["type:warning", `msg:${warning.text}`]);
+          if (warning.channel !== "toolboxOnly") console.error(warning.text);
+        },
+      },
+    );
+  });
+  if (!results) return TLAPM_FAILURE_STATUS;
+  const failures = results.flatMap((result) => {
+    if (parsed.has("--summary")) console.log(formatSummary(result));
+    if (isToolbox) {
+      result.obligations.forEach((obligation) =>
+        printToolboxMessage([
+          "type:obligation",
+          `id:${obligation.id}`,
+          `loc:${obligation.location}`,
+          "status:to be proved",
+          ...(parsed.has("--printallobs") ? [`obl:${obligation.text}`] : []),
+        ]),
+      );
+      printToolboxMessage(["type:obligationsnumber", `count:${result.obligations.length}`]);
+    }
+    if (!parsed.has("--strict")) return [];
+    return checkStrictly(result, lines.start > 0 || lines.end < Infinity, parsed.has("--summary"));
+  });
+  return failures.includes(STRICT_EXIT_STATUS.incomplete)
+    ? STRICT_EXIT_STATUS.incomplete
+    : (failures[0] ?? EXIT_STATUS.success);
+};
+
+const run = async (args: string[]) => {
+  if (args[0] === "prove") return reportErrors(() => prove(args.slice(1)));
+  const parsed = CHECK_COMMAND.parse(args);
+  const inputs = getInputPaths(parsed.positionals);
+  const isConfigured = parsed.has("-config");
+  // A glob also matches modules that only other specs import, which have no configuration.
+  const specPaths = inputs
+    .map((input) => ({ ...input, path: input.path.replace(/\.cfg$/, ".tla") }))
+    .filter(
+      (input) =>
+        input.path.endsWith(".tla") &&
+        (!input.isGlobbed || isConfigured || existsSync(input.path.replace(/\.tla$/, ".cfg"))),
+    )
+    .map((input) => input.path);
+  if (specPaths.length === 0) {
+    console.error(USAGE);
+    return 1;
+  }
+  if (isConfigured && specPaths.length > 1) {
+    console.error("Error: -config applies to one specification, but several were given.");
+    return 1;
+  }
+  return runEach([...new Set(specPaths)], (path) => checkSpec(path, parsed));
+};
+
+export const runCli = async () => {
+  process.exitCode = await run(process.argv.slice(2));
 };

package/src/standard-levels.ts

@@ -28,7 +28,7 @@ const SIGNATURES: Signature[] = [
   ["Head", 0, [0], [2], [1]],
   ["Tail", 0, [0], [2], [1]],
   ["SubSeq", 0, [0, 0, 0], [2, 2, 2], [1, 1, 1]],
-  ["SelectSeq", 0, [0, 1], [2, 2], [1, 1], [], ["1, 0, 0"]],
+  ["SelectSeq", 0, [0, 1], [2, 2], [1, 1], [], ["1,0,0"]],
   ["IsFiniteSet", 0, [0], [2], [1]],
   ["Cardinality", 0, [0], [2], [1]],
   ["IsABag", 0, [0], [2], [1]],
@@ -53,7 +53,7 @@ const SIGNATURES: Signature[] = [
   [":>", 0, [0, 0], [2, 3], [1, 1]],
   ["@@", 0, [0, 0], [2, 2], [1, 1]],
   ["Permutations", 0, [0], [2], [1]],
-  ["SortSeq", 0, [0, 2], [2, 2], [1, 1], [], ["1, 0, 0", "1, 0, 1"]],
+  ["SortSeq", 0, [0, 2], [2, 2], [1, 1], [], ["1,0,0", "1,0,1"]],
   ["RandomElement", 0, [0], [3], [1]],
   ["Any", 0, [], [], []],
   ["ToString", 0, [0], [2], [1]],
@@ -61,7 +61,7 @@ const SIGNATURES: Signature[] = [
   ["AssertEq", 0, [0, 0], [2, 2], [1, 1]],
   ["AssertError", 0, [0, 0], [2, 3], [1, 0]],
   ["TLCGetOrDefault", 0, [0, 0], [3, 3], [0, 0]],
-  ["TLCGetAndSet", 0, [0, 2, 0, 0], [3, 3, 3, 3], [0, 0, 0, 0], [], ["1, 2, 1"]],
+  ["TLCGetAndSet", 0, [0, 2, 0, 0], [3, 3, 3, 3], [0, 0, 0, 0], [], ["1,2,1"]],
   ["Trace", 0, [], [], []],
   ["CounterExample", 0, [], [], []],
   ["ToTrace", 0, [0], [3], [0]],

package/src/transfer.ts

@@ -1,4 +1,11 @@
-import { createModelValue, createSortedFunction, createSortedSet, type Value } from "./values";
+import {
+  createModelValue,
+  createSortedFunction,
+  createSortedSet,
+  getDomainInterval,
+  getFunctionDomain,
+  type Value,
+} from "./values";
 
 const FALSE = 0;
 const TRUE = 1;
@@ -7,6 +14,7 @@ const STRING = 3;
 const MODEL_VALUE = 4;
 const SET = 5;
 const FUNCTION = 6;
+const INTERVAL_FUNCTION = 7;
 
 export interface EncodedStates {
   words: Int32Array;
@@ -65,6 +73,13 @@ export const createStateEncoder = (): StateEncoder => {
       return value.elements.every(encodeValue);
     }
     if (value.kind === "Function") {
+      const interval = getDomainInterval(value.domain);
+      if (interval) {
+        push(INTERVAL_FUNCTION);
+        push(interval.low);
+        push(value.range.length);
+        return value.range.every(encodeValue);
+      }
       push(FUNCTION);
       push(value.domain.length);
       return value.domain.every(encodeValue) && value.range.every(encodeValue);
@@ -102,6 +117,12 @@ export const createStateDecoder =
       if (tag === INTEGER) return words[cursor.index++];
       if (tag === STRING) return encoded.strings[words[cursor.index++]];
       if (tag === MODEL_VALUE) return createModelValue(encoded.strings[words[cursor.index++]]);
+      if (tag === INTERVAL_FUNCTION) {
+        const low = words[cursor.index++];
+        const length = words[cursor.index++];
+        const domain = getFunctionDomain({ kind: "Interval", low, high: low + length - 1 });
+        return createSortedFunction(domain, decodeElements(length));
+      }
       const length = words[cursor.index++];
       if (tag === SET) return createSortedSet(decodeElements(length));
       const domain = decodeElements(length);

package/src/values.ts

@@ -4,6 +4,7 @@ import { prettyPrint } from "./pretty-print";
 export interface ModelValue {
   kind: "ModelValue";
   name: string;
+  type: string | undefined;
 }
 
 export interface SetValue {
@@ -233,8 +234,7 @@ const compareSets = (left: SetValue, right: SetValue) => {
   return 0;
 };
 
-const getModelValueType = (value: Value) =>
-  isModelValue(value) && value.name.length > 2 && value.name[1] === "_" ? value.name[0] : undefined;
+const getModelValueType = (value: Value) => (isModelValue(value) ? value.type : undefined);
 
 const checkModelValueTypes = (left: Value, right: Value, isEquality: boolean) => {
   const leftType = getModelValueType(left);
@@ -349,11 +349,14 @@ export const compareValues = (left: Value, right: Value): number => {
       const otherSize = Math.max(0, right.high - right.low + 1);
       return size !== otherSize || size === 0 ? size - otherSize : left.low - right.low;
     }
+    if (leftKind === "ModelValue" && rightKind === "ModelValue") {
+      if (left.type !== undefined && right.type !== undefined) {
+        checkModelValueTypes(left, right, false);
+      }
+      return compareStrings(left.name, right.name);
+    }
     if (leftKind === "ModelValue" || rightKind === "ModelValue") {
       checkModelValueTypes(left, right, false);
-      if (leftKind === "ModelValue" && rightKind === "ModelValue") {
-        return compareStrings(left.name, right.name);
-      }
       return leftKind === "ModelValue" ? -1 : 1;
     }
     if (left === ANY_SET || right === ANY_SET) return compareAny(left === ANY_SET ? right : left);
@@ -386,7 +389,11 @@ export const createModelValue = (name: string): ModelValue => {
   rankString(name);
   const known = modelValues.get(name);
   if (known) return known;
-  const value: ModelValue = { kind: "ModelValue", name };
+  const value: ModelValue = {
+    kind: "ModelValue",
+    name,
+    type: name.length > 2 && name[1] === "_" ? name[0] : undefined,
+  };
   modelValues.set(name, value);
   return value;
 };
@@ -1667,6 +1674,25 @@ export const isElementOf = (element: Value, set: Value): boolean => {
   }
 };
 
+// A function that agrees with a member of a function set outside some entries is a member when
+// those entries are.
+export const isUpdatedElementOf = (element: Value, member: Value, set: Value): boolean => {
+  if (element === member) return true;
+  if (
+    typeof element !== "object" ||
+    typeof member !== "object" ||
+    typeof set !== "object" ||
+    element.kind !== "Function" ||
+    member.kind !== "Function" ||
+    set.kind !== "FunctionSet" ||
+    element.domain !== member.domain
+  )
+    return isElementOf(element, set);
+  return element.range.every((item, index) =>
+    isUpdatedElementOf(item, member.range[index], set.range),
+  );
+};
+
 export const toFiniteFunction = (value: Value): FunctionValue => {
   if (typeof value === "object" && value.kind === "Function") return value;
   if (typeof value === "object" && value.kind === "LazyFunction") {
@@ -1694,6 +1720,19 @@ const describeMissingArgument = (functionValue: Value, argument: Value) =>
     ? `In applying the function\n${prettyPrint(formatValue(functionValue))},\nthe first argument is:\n${prettyPrint(formatValue(argument))}\nwhich is not in its domain.\n`
     : `Attempted to apply function ${formatValue(functionValue)} to ${formatValue(argument)}, which is not in its domain`;
 
+const findKeyIndex = (domain: Value[], argument: Value) => {
+  const first = domain[0];
+  if (typeof argument === "number" && typeof first === "number") {
+    const offset = argument - first;
+    if (domain[offset] === argument) return offset;
+  }
+  if (domain.length <= 8) {
+    for (let index = 0; index < domain.length; index++)
+      if (domain[index] === argument) return index;
+  }
+  return findIndex(domain, argument);
+};
+
 export const applyFunction = (
   functionValue: Value,
   argument: Value,
@@ -1714,17 +1753,7 @@ export const applyFunction = (
       `Attempted to apply ${formatValue(functionValue)}, which is not a function`,
     );
   }
-  const domain = functionValue.domain;
-  const first = domain[0];
-  if (typeof argument === "number" && typeof first === "number") {
-    const offset = argument - first;
-    if (domain[offset] === argument) return functionValue.range[offset];
-  }
-  if (domain.length <= 8) {
-    for (let index = 0; index < domain.length; index++)
-      if (domain[index] === argument) return functionValue.range[index];
-  }
-  const position = findIndex(domain, argument);
+  const position = findKeyIndex(functionValue.domain, argument);
   if (position === -1) {
     throw new TlaEvaluationError(
       describeMissing?.(functionValue) ?? describeMissingArgument(functionValue, argument),
@@ -1733,12 +1762,16 @@ export const applyFunction = (
   return functionValue.range[position];
 };
 
+export const getDomainInterval = (domain: Value[]) =>
+  isIntervalDomain(domain)
+    ? (intervalArrays.get(domain) ?? getIntervalValue(domain.length))
+    : undefined;
+
 export const getDomain = (functionValue: Value): Value => {
   if (typeof functionValue === "object" && functionValue.kind === "LazyFunction")
     return functionValue.domain;
   const domain = toFiniteFunction(functionValue).domain;
-  if (!isIntervalDomain(domain)) return createSortedSet(domain);
-  return intervalArrays.get(domain) ?? getIntervalValue(domain.length);
+  return getDomainInterval(domain) ?? createSortedSet(domain);
 };
 
 export const updateFunction = (
@@ -1762,7 +1795,7 @@ export const updateFunction = (
     };
   }
   const finite = toFiniteFunction(functionValue);
-  const position = findIndex(finite.domain, argument);
+  const position = findKeyIndex(finite.domain, argument);
   if (position === -1) return finite;
   const range = [...finite.range];
   const previous = finite.range[position];