tla-ts 0.0.1 → 0.0.3-dev.a50fc89

package/src/standard-modules.ts

@@ -1,5 +1,6 @@
-import { TlaEvaluationError } from "./errors";
+import { TlaEvaluationError, TlaModuleError } from "./errors";
 import {
+  createOverflowError,
   defineArithmetic,
   defineBinary,
   defineComparison,
@@ -7,12 +8,14 @@ import {
   defineHigherOrderOperator,
   defineOperator,
   getBoolean,
-  getElements,
   getNumber,
   getOperator,
-  getSequence,
   getSet,
+  getIntegersSignature,
+  getJavaInteger,
   getValue,
+  isOutOfRange,
+  overrideWithJava,
 } from "./builtins";
 import type { Builtin, BuiltinArgument } from "./evaluate";
 import type { CheckerRuntime } from "./runtime";
@@ -21,10 +24,10 @@ import { createJsonOperators, type FileSystem } from "./json";
 import { createRandomizationOperators } from "./randomization";
 import { createPossibleOperators, createTlcExtOperators } from "./tlc-ext";
 import { createTlcOperators } from "./tlc-module";
+import { TLC_JAVA_METHODS } from "./tlc-modules";
 import {
-  checkInteger,
   createSet,
-  createSortedSet,
+  describeKind,
   createTuple,
   differenceSorted,
   formatValue,
@@ -32,16 +35,24 @@ import {
   getDomain,
   intersectSorted,
   isElementOf,
+  getSequenceElements,
   isEmptySet,
   isEnumerable,
+  isFiniteSet,
   isFunctionValue,
   isSetValue,
   toFiniteSet,
   unionSorted,
-  valuesEqual,
+  isEqualLikeTlc,
+  createEnumeratedSet,
+  EMPTY_SET,
+  isEnumeratedSet,
+  type IntervalValue,
+  type MixedSetValue,
   type SetOperationValue,
   type SetValue,
   type Value,
+  rankString,
 } from "./values";
 
 interface StandardModule {
@@ -53,30 +64,67 @@ const isSubsetOf = (left: Value[], right: Value) =>
   left.every((element) => isElementOf(element, right));
 
 const BOOLEAN_SET = createSet([false, true]);
-const EMPTY_SET = createSortedSet([]);
-const MERGED_UNION_LIMIT = 8;
+
+// Merging the sets pairwise in a balanced tree compares each element about log2 of their count
+// times.
+const unionInTree = (sets: SetValue[], start: number, end: number): SetValue => {
+  if (end - start === 1) return sets[start];
+  const middle = (start + end) >> 1;
+  return unionSorted(unionInTree(sets, start, middle), unionInTree(sets, middle, end));
+};
 
 const SET_OPERATORS: SetOperationValue["operator"][] = ["\\cup", "\\cap", "\\"];
 
+type ReducibleSet = SetValue | MixedSetValue | IntervalValue;
+
+const isReducible = (value: Value): value is ReducibleSet =>
+  typeof value === "object" &&
+  (value.kind === "Set" || value.kind === "MixedSet" || value.kind === "Interval");
+
+const isMixedSet = (value: Value) => typeof value === "object" && value.kind === "MixedSet";
+
+const getReducibleElements = (set: ReducibleSet) =>
+  set.kind === "Interval" ? toFiniteSet(set).elements : set.elements;
+
+// TLC combines a set enumeration or an interval with another set element by element, testing each
+// element for membership in the other set.
+const combineReducible = (
+  operator: SetOperationValue["operator"],
+  set: ReducibleSet,
+  other: Value,
+): Value => {
+  const elements = getReducibleElements(set);
+  if (operator !== "\\cup") {
+    const isKept = operator === "\\cap";
+    return createEnumeratedSet(
+      elements.filter((element) => isElementOf(element, other) === isKept),
+    );
+  }
+  if (elements.length === 0) return other;
+  if (!isReducible(other)) return { kind: "SetOperation", operator, left: set, right: other };
+  const added = toFiniteSet(other).elements.filter((element) => !isElementOf(element, set));
+  return createEnumeratedSet([...elements, ...added]);
+};
+
 export const CORE_OPERATORS: Builtin[] = [
   defineConstant("TRUE", true),
   defineConstant("FALSE", false),
   defineConstant("BOOLEAN", BOOLEAN_SET),
   defineConstant("STRING", { kind: "InfiniteSet", name: "STRING" }),
-  defineBinary("=", valuesEqual),
-  defineBinary("#", (left, right) => !valuesEqual(left, right)),
+  defineBinary("=", isEqualLikeTlc),
+  defineBinary("#", (left, right) => !isEqualLikeTlc(left, right)),
   defineOperator("~", 1, (args) => !getBoolean(args, 0, "~")),
   defineOperator("<=>", 2, (args) => getBoolean(args, 0, "<=>") === getBoolean(args, 1, "<=>")),
   defineOperator("/\\", 2, (args) => getBoolean(args, 0, "/\\") && getBoolean(args, 1, "/\\")),
   defineOperator("\\/", 2, (args) => getBoolean(args, 0, "\\/") || getBoolean(args, 1, "\\/")),
   defineOperator("=>", 2, (args) => !getBoolean(args, 0, "=>") || getBoolean(args, 1, "=>")),
   defineOperator("\\in", 2, (args) =>
-    isElementOf(getValue(args, 0, "\\in"), getSet(args, 1, "\\in")),
+    isElementOf(getValue(args, 0, "\\in"), getValue(args, 1, "\\in")),
   ),
   defineOperator(
     "\\notin",
     2,
-    (args) => !isElementOf(getValue(args, 0, "\\notin"), getSet(args, 1, "\\notin")),
+    (args) => !isElementOf(getValue(args, 0, "\\notin"), getValue(args, 1, "\\notin")),
   ),
   defineOperator("\\subseteq", 2, (args) => {
     const left = getSet(args, 0, "\\subseteq");
@@ -88,13 +136,14 @@ export const CORE_OPERATORS: Builtin[] = [
       right.kind === "Subset" &&
       isEnumerable(left.base);
     if (isPowerSets) return isSubsetOf(toFiniteSet(left.base).elements, right.base);
-    return isSubsetOf(getElements(args, 0, "\\subseteq"), right);
+    return isSubsetOf(toFiniteSet(left).elements, right);
   }),
   ...SET_OPERATORS.map((operator) =>
     defineOperator(operator, 2, (args) => {
       const left = getSet(args, 0, operator);
       const right = getSet(args, 1, operator);
-      if (isEnumerable(left) && isEnumerable(right)) {
+      if (isReducible(left) && isReducible(right) && !isMixedSet(left) && !isMixedSet(right)) {
+        if (operator === "\\cup" && isEmptySet(left)) return right;
         const merge =
           operator === "\\cup"
             ? unionSorted
@@ -103,35 +152,28 @@ export const CORE_OPERATORS: Builtin[] = [
               : differenceSorted;
         return merge(toFiniteSet(left), toFiniteSet(right));
       }
-      const isReducible = (value: Value) =>
-        typeof value === "object" && (value.kind === "Set" || value.kind === "Interval");
-      const isSwapped = operator === "\\cup" && !isReducible(left) && isReducible(right);
-      const set: Value = isSwapped
-        ? { kind: "SetOperation", operator, left: right, right: left }
-        : { kind: "SetOperation", operator, left, right };
-      return isEnumerable(set) ? toFiniteSet(set) : set;
+      if (isReducible(left)) return combineReducible(operator, left, right);
+      if (operator !== "\\" && isReducible(right)) return combineReducible(operator, right, left);
+      return { kind: "SetOperation", operator, left, right };
     }),
   ),
   defineOperator("SUBSET", 1, (args) => ({ kind: "Subset", base: getSet(args, 0, "SUBSET") })),
   defineOperator("UNION", 1, (args) => {
-    const sets = getElements(args, 0, "UNION");
-    sets.forEach((element) => {
-      if (!isSetValue(element)) {
-        throw new TlaEvaluationError(
-          `UNION expects a set of sets but found ${formatValue(element)}`,
-        );
-      }
-    });
-    if (sets.every((element) => typeof element === "object" && element.kind === "Set")) {
-      if (sets.length > MERGED_UNION_LIMIT) {
-        return createSet(sets.flatMap((element) => toFiniteSet(element).elements));
-      }
-      return sets.reduce<SetValue>(
-        (union, element) => unionSorted(union, toFiniteSet(element)),
-        EMPTY_SET,
+    const outer = getSet(args, 0, "UNION");
+    const sets = isEnumeratedSet(outer) ? outer.elements.filter(isEnumeratedSet) : [];
+    if (!isEnumeratedSet(outer) || sets.length !== outer.elements.length) {
+      return { kind: "Union", base: outer };
+    }
+    if (isMixedSet(outer) || sets.some(isMixedSet)) {
+      const union: MixedSetValue = { kind: "MixedSet", elements: [] };
+      sets.forEach((set) =>
+        set.elements.forEach((member) => {
+          if (!isElementOf(member, union)) union.elements.push(member);
+        }),
       );
+      return createEnumeratedSet(union.elements);
     }
-    return sets.reduce((left, right) => ({ kind: "SetOperation", operator: "\\cup", left, right }));
+    return sets.length === 0 ? EMPTY_SET : unionInTree(sets.map(toFiniteSet), 0, sets.length);
   }),
   defineOperator("DOMAIN", 1, (args) => {
     const value = getValue(args, 0, "DOMAIN");
@@ -142,24 +184,36 @@ export const CORE_OPERATORS: Builtin[] = [
   }),
 ];
 
-const NATURALS_OPERATORS: Builtin[] = [
+const DOT_DOT_SIGNATURE = getIntegersSignature("DotDot", "IntervalValue", 2);
+
+const NEG_SIGNATURE = getIntegersSignature("Neg", "IntValue", 1);
+
+const NATURALS_OPERATORS = [
   defineConstant("Nat", { kind: "InfiniteSet", name: "Nat" }),
-  defineArithmetic("+", (left, right) => left + right),
-  defineArithmetic("-", (left, right) => left - right),
-  defineArithmetic("*", (left, right) => left * right),
-  defineArithmetic("^", (base, exponent) => {
-    if (exponent < 0 || (base === 0 && exponent === 0)) {
-      throw new TlaEvaluationError(`${base}^${exponent} is undefined`);
+  defineArithmetic("+", "Plus", (left, right) => left + right),
+  defineArithmetic("-", "Minus", (left, right) => left - right),
+  defineArithmetic("*", "Times", (left, right) => left * right),
+  defineArithmetic("^", "Expt", (base, exponent) => {
+    if (exponent < 0) {
+      throw new TlaModuleError(
+        `The second argument of ^ should be a natural number, but instead it is:\n${exponent}`,
+      );
     }
+    if (base === 0 && exponent === 0) throw new TlaModuleError("0^0 is undefined.");
     return base ** exponent;
   }),
-  defineArithmetic("\\div", (left, right) => {
-    if (right === 0) throw new TlaEvaluationError(`Division by zero in ${left} \\div 0`);
-    return Math.floor(left / right);
+  defineArithmetic("\\div", "Divide", (left, right) => {
+    if (right === 0) throw new TlaModuleError("The second argument of \\div is 0.");
+    const quotient = Math.floor(left / right);
+    if (isOutOfRange(quotient)) throw createOverflowError(`${left} \\div ${right}`);
+    return quotient;
   }),
-  defineArithmetic("%", (left, right) => {
-    if (right <= 0)
-      throw new TlaEvaluationError(`The second argument of % must be positive, but was ${right}`);
+  defineArithmetic("%", "Mod", (left, right) => {
+    if (right <= 0) {
+      throw new TlaModuleError(
+        `The second argument of % should be a positive number, but instead it is:\n${right}`,
+      );
+    }
     return ((left % right) + right) % right;
   }),
   defineComparison("<", (left, right) => left < right),
@@ -168,85 +222,162 @@ const NATURALS_OPERATORS: Builtin[] = [
   defineComparison(">=", (left, right) => left >= right),
   defineOperator("..", 2, (args) => ({
     kind: "Interval",
-    low: getNumber(args, 0, ".."),
-    high: getNumber(args, 1, ".."),
+    low: getJavaInteger(getValue(args, 0, ".."), DOT_DOT_SIGNATURE),
+    high: getJavaInteger(getValue(args, 1, ".."), DOT_DOT_SIGNATURE),
   })),
 ];
 
-const INTEGERS_OPERATORS: Builtin[] = [
+const INTEGERS_OPERATORS = [
   defineConstant("Int", { kind: "InfiniteSet", name: "Int" }),
-  defineOperator("-.", 1, (args) => checkInteger(-getNumber(args, 0, "-"))),
+  defineOperator("-.", 1, (args) => {
+    const value = getJavaInteger(getValue(args, 0, "-"), NEG_SIGNATURE);
+    if (isOutOfRange(-value)) throw createOverflowError(`-${value}`);
+    return -value;
+  }),
 ];
 
-const getSequenceOrString = (args: BuiltinArgument[], index: number, name: string) => {
-  const value = getValue(args, index, name);
-  return typeof value === "string" ? value : getSequence(args, index, name);
+const getSequenceArgument = (args: BuiltinArgument[], index: number, message: string) => {
+  const value = getValue(args, index, message);
+  const elements = getSequenceElements(value);
+  if (!elements) throw new TlaModuleError(`${message}\n${formatValue(value)}`);
+  return elements;
 };
 
-const SEQUENCES_OPERATORS: Builtin[] = [
-  defineOperator("Seq", 1, (args) => ({ kind: "SequenceSet", base: getSet(args, 0, "Seq") })),
-  defineOperator("Len", 1, (args) => getSequenceOrString(args, 0, "Len").length),
+const getSequenceOrString = (args: BuiltinArgument[], index: number, message: string) => {
+  const value = getValue(args, index, message);
+  return typeof value === "string" ? value : getSequenceArgument(args, index, message);
+};
+
+const sliceSequence = (sequence: string | Value[], start: number, end?: number) =>
+  typeof sequence === "string"
+    ? rankString(sequence.slice(start, end))
+    : createTuple(sequence.slice(start, end));
+
+const getSequenceMessage = (position: string, name: string) =>
+  `The ${position}argument of ${name} should be a sequence, but instead it is:`;
+
+const getEvaluatingMessage = (form: string, kind: string) =>
+  `Evaluating an expression of the form ${form} when s is not a ${kind}:`;
+
+const withJavaMethods = (module: string, operators: Builtin[]) => {
+  const methods = TLC_JAVA_METHODS.get(module);
+  if (!methods) return operators;
+  return operators.map((operator) => {
+    const signature = methods.get(operator.name);
+    return signature ? overrideWithJava(operator, signature) : operator;
+  });
+};
+
+const SEQUENCES_OPERATORS = withJavaMethods("Sequences", [
+  defineOperator("Seq", 1, (args) => ({ kind: "SequenceSet", base: getValue(args, 0, "Seq") })),
+  defineOperator(
+    "Len",
+    1,
+    (args) => getSequenceOrString(args, 0, getSequenceMessage("", "Len")).length,
+  ),
   defineOperator("\\o", 2, (args) => {
-    const left = getSequenceOrString(args, 0, "\\o");
-    const right = getSequenceOrString(args, 1, "\\o");
-    if (typeof left === "string" && typeof right === "string") return left + right;
-    if (typeof left === "string" || typeof right === "string") {
-      throw new TlaEvaluationError("\\o cannot concatenate a string with a sequence");
+    const left = getSequenceOrString(args, 0, getEvaluatingMessage("s \\o t", "sequence"));
+    if (typeof left !== "string") {
+      const right = getSequenceArgument(args, 1, getEvaluatingMessage("t \\o s", "sequence"));
+      return createTuple([...left, ...right]);
     }
-    return createTuple([...left, ...right]);
+    const right = getValue(args, 1, "\\o");
+    if (typeof right !== "string") {
+      throw new TlaModuleError(
+        `${getEvaluatingMessage("t \\o s", "string")}\n${formatValue(right)}`,
+      );
+    }
+    return rankString(left + right);
   }),
   defineOperator("Append", 2, (args) =>
-    createTuple([...getSequence(args, 0, "Append"), getValue(args, 1, "Append")]),
+    createTuple([
+      ...getSequenceArgument(args, 0, getEvaluatingMessage("Append(s, v)", "sequence")),
+      getValue(args, 1, "Append"),
+    ]),
   ),
   defineOperator("Head", 1, (args) => {
-    const elements = getSequence(args, 0, "Head");
+    const elements = getSequenceArgument(args, 0, getSequenceMessage("", "Head"));
     if (elements.length === 0)
-      throw new TlaEvaluationError("Attempted to apply Head to the empty sequence");
+      throw new TlaModuleError("Attempted to apply Head to the empty sequence.");
     return elements[0];
   }),
   defineOperator("Tail", 1, (args) => {
-    const sequence = getSequenceOrString(args, 0, "Tail");
+    const sequence = getSequenceOrString(args, 0, getSequenceMessage("", "Tail"));
     if (sequence.length === 0)
-      throw new TlaEvaluationError("Attempted to apply Tail to the empty sequence");
-    return typeof sequence === "string" ? sequence.slice(1) : createTuple(sequence.slice(1));
+      throw new TlaModuleError("Attempted to apply Tail to the empty sequence.");
+    return sliceSequence(sequence, 1);
   }),
   defineOperator("SubSeq", 3, (args) => {
-    const sequence = getSequenceOrString(args, 0, "SubSeq");
+    const sequence = getSequenceOrString(args, 0, getSequenceMessage("first ", "SubSeq"));
     const start = getNumber(args, 1, "SubSeq");
     const end = getNumber(args, 2, "SubSeq");
-    if (start > end) return typeof sequence === "string" ? "" : createTuple([]);
-    if (start < 1 || end > sequence.length) {
-      throw new TlaEvaluationError(
-        `SubSeq bounds ${start}..${end} are outside 1..${sequence.length}`,
+    if (start > end) return sliceSequence(sequence, 0, 0);
+    const outOfDomain =
+      start < 1
+        ? { position: "second", value: start }
+        : end > sequence.length
+          ? { position: "third", value: end }
+          : undefined;
+    if (outOfDomain) {
+      throw new TlaModuleError(
+        `The ${outOfDomain.position} argument of SubSeq must be in the domain of its first argument:\n${formatValue(getValue(args, 0, "SubSeq"))}\n, but instead it is\n${outOfDomain.value}`,
       );
     }
-    return typeof sequence === "string"
-      ? sequence.slice(start - 1, end)
-      : createTuple(sequence.slice(start - 1, end));
+    return sliceSequence(sequence, start - 1, end);
   }),
   defineHigherOrderOperator("SelectSeq", 2, (args) => {
     const test = getOperator(args, 1, "SelectSeq");
     return createTuple(
-      getSequence(args, 0, "SelectSeq").filter((element) => {
+      getSequenceArgument(args, 0, getSequenceMessage("first ", "SelectSeq")).filter((element) => {
         const result = test.call([element]);
-        if (typeof result !== "boolean")
-          throw new TlaEvaluationError("The test of SelectSeq must return a boolean");
+        if (typeof result !== "boolean") {
+          throw new TlaModuleError(
+            `The second argument of SelectSeq should be a operator, but instead it is:\n${formatValue(result)}`,
+          );
+        }
         return result;
       }),
     );
   }),
-];
+]);
 
-const FINITE_SETS_OPERATORS: Builtin[] = [
-  defineOperator("IsFiniteSet", 1, (args) => isEnumerable(getSet(args, 0, "IsFiniteSet"))),
+const FINITE_SETS_OPERATORS = withJavaMethods("FiniteSets", [
+  defineOperator("IsFiniteSet", 1, (args) => {
+    const value = getValue(args, 0, "IsFiniteSet");
+    if (isSetValue(value)) return isFiniteSet(value);
+    if (typeof value === "object" && value.kind === "Function") return true;
+    throw new TlaEvaluationError(
+      typeof value === "object" && value.kind === "LazyFunction"
+        ? `Attempted to check if the function:\n${formatValue(value)}\nis a finite set.`
+        : `Attempted to check if the ${describeKind(value)} ${formatValue(value)} is a finite set.`,
+    );
+  }),
   defineOperator("Cardinality", 1, (args) => {
-    const set = getSet(args, 0, "Cardinality");
-    if (isEmptySet(set)) return 0;
-    if (!isEnumerable(set))
-      throw new TlaEvaluationError(`Cardinality of the infinite set ${formatValue(set)}`);
+    const set = getValue(args, 0, "Cardinality");
+    if (
+      !isSetValue(set) ||
+      (typeof set === "object" && (set.kind === "InfiniteSet" || set.kind === "SequenceSet"))
+    ) {
+      throw new TlaModuleError(
+        `Attempted to compute cardinality of the value\n${formatValue(set)}`,
+      );
+    }
     return getCardinality(set);
   }),
-];
+]);
+
+// The modules that each standard module loads through EXTENDS and LOCAL INSTANCE in TLC.
+export const STANDARD_DEPENDENCIES = new Map([
+  ["Integers", ["Naturals"]],
+  ["Sequences", ["Naturals"]],
+  ["FiniteSets", ["Naturals", "Sequences"]],
+  ["TLC", ["Naturals", "Sequences", "FiniteSets"]],
+  ["Randomization", ["Naturals", "FiniteSets", "Sequences"]],
+  ["_Possible", ["TLC", "TLCExt", "Naturals", "Sequences", "FiniteSets", "Integers"]],
+  ["Json", ["Sequences", "FiniteSets", "TLC", "Naturals"]],
+  ["Bags", ["Naturals"]],
+  ["TLCExt", ["TLC", "Integers", "Naturals", "Sequences", "FiniteSets"]],
+]);
 
 export const getStandardModule = (
   name: string,
@@ -264,7 +395,7 @@ export const getStandardModule = (
     case "FiniteSets":
       return { extends: [], operators: FINITE_SETS_OPERATORS };
     case "TLC":
-      return { extends: [], operators: createTlcOperators(print, runtime) };
+      return { extends: [], operators: withJavaMethods("TLC", createTlcOperators(print, runtime)) };
     case "Randomization":
       return { extends: [], operators: createRandomizationOperators(runtime) };
     case "_Possible":

package/src/subexpressions.ts

@@ -1,6 +1,5 @@
 import type {
   Bound,
-  BoundPattern,
   Definition,
   Expression,
   Instance,
@@ -11,7 +10,8 @@ import type {
 } from "./ast";
 import { TlaSyntaxError } from "./errors";
 import { getChildren } from "./evaluate";
-import { mapChildren, mapDefinition } from "./instances";
+import { getSourceName, mapChildren, mapDefinition } from "./instances";
+import { getArity, getBoundNames, getPatternNames, LOCAL_RESOLUTION } from "./resolve";
 
 type Lookup = (name: string) => Definition | undefined;
 
@@ -27,21 +27,10 @@ interface Selection {
   location: SourceLocation;
 }
 
-const LOCAL_RESOLUTION = { kind: "Local" } as const;
-
 const SELECTION_NAME = "$subexpression";
 
-const describeLocation = (location: SourceLocation) =>
-  `line ${location.line}, column ${location.column}`;
-
 const createError = (message: string, location: SourceLocation) =>
-  new TlaSyntaxError(`${message} at ${describeLocation(location)}`);
-
-const getPatternNames = (pattern: BoundPattern) =>
-  pattern.kind === "Name" ? [pattern.name] : pattern.names;
-
-const getBoundNames = (bounds: Bound[]) =>
-  bounds.flatMap((bound) => getPatternNames(bound.pattern));
+  new TlaSyntaxError(`${message} at line ${location.line}, column ${location.column}`);
 
 const getBoundSets = (bounds: Bound[]) =>
   bounds
@@ -164,11 +153,8 @@ const findLabel = (
   return undefined;
 };
 
-const takeArguments = (
-  selector: Selector & { kind: "Name" | "Bound" },
-  count: number,
-  pool: Expression[],
-) => (selector.args.length === 0 && pool.length > 0 ? pool.splice(0, count) : selector.args);
+const takeArguments = (args: Expression[], count: number, pool: Expression[]) =>
+  args.length === 0 && pool.length > 0 ? pool.splice(0, count) : args;
 
 const wrapLets = (body: Expression, lets: Definition[][], location: SourceLocation) =>
   lets.reduceRight<Expression>(
@@ -221,7 +207,7 @@ const createApplication = (selection: Selection, body: Expression): Expression =
   };
 };
 
-const applyDefinition = (definition: Definition, selection: Selection): Expression => {
+const applyDefinition = (definition: Definition, selection: Selection) => {
   const location = selection.location;
   const params =
     definition.kind === "Operator"
@@ -267,11 +253,11 @@ const selectSubexpression = (
       names.push(selector.name);
       args.push(...selector.args);
       const name = names.join("!");
-      const definition = local ? local.find((candidate) => candidate.name === name) : lookup(name);
+      const definition = local
+        ? local.find((candidate) => getSourceName(candidate) === name)
+        : lookup(name);
       if (definition) {
-        const arity = definition.kind === "Operator" ? definition.params.length : 0;
-        const given = args.length === 0 && pool.length > 0 ? pool.splice(0, arity) : args;
-        return { definition, args: given };
+        return { definition, args: takeArguments(args, getArity(definition), pool) };
       }
       state.index++;
     }
@@ -281,7 +267,7 @@ const selectSubexpression = (
   const enterDefinition = (local: Definition[] | undefined) => {
     const found = findDefinition(local);
     const definition = found.definition;
-    const arity = definition.kind === "Operator" ? definition.params.length : 0;
+    const arity = getArity(definition);
     const isLast = state.index === selectors.length - 1;
     if (found.args.length !== arity && found.args.length !== 0) {
       throw createError(`The operator ${definition.name} requires ${arity} arguments`, location);
@@ -298,13 +284,14 @@ const selectSubexpression = (
   };
 
   const followLabels = () => {
-    while (state.index < selectors.length && selectors[state.index].kind === "Name") {
-      const selector = selectors[state.index] as Selector & { kind: "Name" };
+    while (state.index < selectors.length) {
+      const selector = selectors[state.index];
+      if (selector.kind !== "Name") return;
       const found = state.current && findLabel(state.current, selector.name, []);
       if (!found) throw createError(`Cannot find label ${selector.name}`, location);
       selection.lets.push(...found.lets);
       selection.params.push(...found.label.params.map((name) => ({ name, arity: 0 })));
-      selection.args.push(...takeArguments(selector, found.label.params.length, pool));
+      selection.args.push(...takeArguments(selector.args, found.label.params.length, pool));
       state.current = found.label.body;
       state.index++;
     }
@@ -331,7 +318,7 @@ const selectSubexpression = (
     if (binder && (selector.kind === "Bound" || selector.kind === "At")) {
       selection.params.push(...binder.names.map((name) => ({ name, arity: 0 })));
       if (selector.kind === "Bound") {
-        selection.args.push(...takeArguments(selector, binder.names.length, pool));
+        selection.args.push(...takeArguments(selector.args, binder.names.length, pool));
       }
       state.current = binder.body;
       return;
@@ -372,8 +359,7 @@ const selectSubexpression = (
         throw createError("A name selector here must be from a LET clause", location);
       }
       selection.lets.push(state.current.definitions);
-      const definitions = state.current.definitions;
-      const last = enterDefinition(definitions);
+      const last = enterDefinition(state.current.definitions);
       if (last) {
         return applyDefinition(last.definition, {
           ...selection,
@@ -407,10 +393,10 @@ const containsReference = (expression: Expression): boolean =>
   (expression.kind === "Apply" && expression.operator.includes("!")) ||
   getChildren(expression).some(containsReference);
 
-export const resolveSubexpressions = (expression: Expression, lookup: Lookup): Expression => {
+const resolveSubexpressions = (expression: Expression, lookup: Lookup): Expression => {
   if (!containsReference(expression)) return expression;
   if (expression.kind === "Let") {
-    const local: Lookup = (name: string) =>
+    const local: Lookup = (name) =>
       expression.definitions.find((definition) => definition.name === name) ?? lookup(name);
     return mapChildren(expression, (child) => resolveSubexpressions(child, local));
   }
@@ -446,10 +432,10 @@ export const resolveInstanceSubexpressions = (instance: Instance, module: Module
 
 export const resolveModuleSubexpressions = (module: Module, available: Definition[]): Module => {
   const lookup = createLookup([...available, ...module.definitions]);
-  const isResolved =
+  const hasReferences =
     module.definitions.some((definition) => containsReference(definition.body)) ||
     module.assumptions.some((assumption) => containsReference(assumption.expression));
-  if (!isResolved) return module;
+  if (!hasReferences) return module;
   return {
     ...module,
     definitions: module.definitions.map((definition) =>