tla-ts 0.0.2 → 0.0.3-dev.a50fc89

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/src/parser.ts CHANGED
@@ -153,6 +153,8 @@ const EXPRESSION_PREFIXES = new Set([
153
153
  "<>",
154
154
  ]);
155
155
 
156
+ const CIRCLED_OPERATORS = new Set(["+", "-", "/", "\\X"]);
157
+
156
158
  const OPERATOR_ARGUMENTS = new Set([
157
159
  "SUBSET",
158
160
  "UNION",
@@ -213,6 +215,41 @@ const PROOF_STARTS = new Set(["BY", "PROOF", "OBVIOUS", "OMITTED"]);
213
215
 
214
216
  const PROVABLE_STEPS = new Set(["Qed", "Pick", "Case", "Assert"]);
215
217
 
218
+ const TLAPM_PUNCTUATION = new Set([
219
+ ",",
220
+ ".",
221
+ "_",
222
+ "(",
223
+ ")",
224
+ "[",
225
+ "]",
226
+ "{",
227
+ "}",
228
+ "<<",
229
+ ">>",
230
+ "]_",
231
+ ">>_",
232
+ "==",
233
+ "!",
234
+ "@",
235
+ ":",
236
+ "::",
237
+ ";",
238
+ "->",
239
+ "<-",
240
+ "|->",
241
+ "\\A",
242
+ "\\AA",
243
+ "\\E",
244
+ "\\EE",
245
+ "\\forall",
246
+ "\\exists",
247
+ "WF_",
248
+ "SF_",
249
+ "(*{",
250
+ "}*)",
251
+ ]);
252
+
216
253
  const UNNAMED_STEP_LABELS = new Set(["", "*", "-"]);
217
254
 
218
255
  const FACT_KEYWORDS = new Set([
@@ -264,7 +301,9 @@ export const extractModule = (source: string, name: string) => {
264
301
  return lines.slice(start, end === -1 ? undefined : end + 1).join("\n");
265
302
  };
266
303
 
267
- export const parseModule = (source: string): Module => {
304
+ // isTlapm parses as tlapm does for proof obligations, which accepts and rejects some modules that
305
+ // SANY does not.
306
+ const parse = (source: string, isTlapm: boolean): Module => {
268
307
  const tokens = tokenize(source);
269
308
  const cursor = { index: 0, offsideColumns: [] as number[] };
270
309
  // A RECURSIVE declaration covers the definitions of its own LET or module.
@@ -280,7 +319,7 @@ export const parseModule = (source: string): Module => {
280
319
 
281
320
  const getLocation = (token: Token): SourceLocation => ({
282
321
  line: token.line,
283
- column: token.column,
322
+ column: isTlapm && token.kind === "String" ? token.endColumn : token.column,
284
323
  range: undefined,
285
324
  });
286
325
 
@@ -324,10 +363,78 @@ export const parseModule = (source: string): Module => {
324
363
  const describe = (token: Token) =>
325
364
  token.kind === "End" ? "end of input" : JSON.stringify(token.text);
326
365
 
327
- const createError = (message: string, token: { line: number; column: number } = peekRaw()) =>
328
- new TlaSyntaxError(`${message} at line ${token.line}, column ${token.column}`);
366
+ const unitStart = { index: 0 };
367
+
368
+ const getDefinitionEnd = (fallback: { line: number; column: number }) => {
369
+ const equals = tokens
370
+ .slice(unitStart.index, cursor.index)
371
+ .find((token) => token.text === "==" && token.kind !== "String");
372
+ return equals ? { line: equals.line, column: equals.endColumn + 1 } : fallback;
373
+ };
374
+
375
+ // tlapm reads TRUE, FALSE, BOOLEAN, STRING and the other keywords only as keywords.
376
+ const isTlapmKeyword = (token: Token) =>
377
+ isTlapm && (token.kind === "Keyword" || BUILTIN_CONSTANTS.has(token.text));
329
378
 
330
- const isFieldName = (token: Token) => token.kind === "Identifier" || token.kind === "Keyword";
379
+ const getPreviousEnd = (position: { line: number; column: number }) => {
380
+ const previous = tokens.findLast(
381
+ (token) =>
382
+ token.kind !== "End" &&
383
+ (token.line < position.line ||
384
+ (token.line === position.line && token.column < position.column)),
385
+ );
386
+ return previous ? { line: previous.line, column: previous.endColumn + 1 } : position;
387
+ };
388
+
389
+ const getTlapmName = (token: Token) => {
390
+ if (/^-{4,}$/.test(token.text)) return "----";
391
+ if (/^={4,}$/.test(token.text)) return "====";
392
+ if (TLAPM_PUNCTUATION.has(token.text) || token.kind === "Number" || token.kind === "Decimal")
393
+ return token.text;
394
+ if (token.kind === "String") return JSON.stringify(token.text);
395
+ if (isTlapmKeyword(token)) return `keyword ${token.text}`;
396
+ return token.kind === "Identifier" ? `identifier ${token.text}` : `operator ${token.text}`;
397
+ };
398
+
399
+ const createError = (
400
+ message: string,
401
+ token: { line: number; column: number } = peekRaw(),
402
+ unexpected?: Token,
403
+ ) => {
404
+ const end = tokens[tokens.length - 1];
405
+ const isEnd = token.line === end.line && token.column === end.column;
406
+ return new TlaSyntaxError(
407
+ `${message} at line ${token.line}, column ${token.column}`,
408
+ undefined,
409
+ getFailure(message, token, isEnd, unexpected),
410
+ );
411
+ };
412
+
413
+ const getFailure = (
414
+ message: string,
415
+ token: { line: number; column: number },
416
+ isEnd: boolean,
417
+ unexpected: Token | undefined,
418
+ ) => {
419
+ if (isTlapm && isEnd) {
420
+ const end = tokens[tokens.length - 1];
421
+ return { reason: "Unexpected EOF", ...getDefinitionEnd(getPreviousEnd(end)), isEnd };
422
+ }
423
+ if (isTlapm && unexpected) {
424
+ return {
425
+ reason: `Unexpected ${getTlapmName(unexpected)}`,
426
+ ...getPreviousEnd(unexpected),
427
+ isEnd,
428
+ };
429
+ }
430
+ return { reason: message, line: token.line, column: token.column, isEnd };
431
+ };
432
+
433
+ const createUnexpected = (message: string, token: Token = peekRaw()) =>
434
+ createError(message, token, token);
435
+
436
+ const isFieldName = (token: Token) =>
437
+ (token.kind === "Identifier" || token.kind === "Keyword") && !isTlapmKeyword(token);
331
438
 
332
439
  const isAt = (text: string, offset = 0) => {
333
440
  const token = peek(offset);
@@ -349,15 +456,16 @@ export const parseModule = (source: string): Module => {
349
456
 
350
457
  const expect = (text: string) => {
351
458
  if (!isAt(text))
352
- throw createError(`Expected ${JSON.stringify(text)} but found ${describe(peek())}`);
459
+ throw createUnexpected(`Expected ${JSON.stringify(text)} but found ${describe(peek())}`);
353
460
  return advance();
354
461
  };
355
462
 
356
463
  const expectIdentifier = () => {
357
464
  const token = peek();
358
465
  if (token.kind !== "Identifier") {
359
- throw createError(`Expected an identifier but found ${describe(token)}`);
466
+ throw createUnexpected(`Expected an identifier but found ${describe(token)}`);
360
467
  }
468
+ if (isTlapmKeyword(token)) throw createUnexpected(`Unexpected keyword ${token.text}`);
361
469
  cursor.index++;
362
470
  return token.text;
363
471
  };
@@ -422,7 +530,7 @@ export const parseModule = (source: string): Module => {
422
530
  const token = peek();
423
531
  if (token.kind === "Identifier") {
424
532
  const start = cursor.index;
425
- const names = [expectIdentifier()];
533
+ const names = [advance().text];
426
534
  while (isAt("!") && peek(1).kind === "Identifier") {
427
535
  advance();
428
536
  names.push(expectIdentifier());
@@ -433,7 +541,7 @@ export const parseModule = (source: string): Module => {
433
541
  );
434
542
  }
435
543
  if (isAt("<<") || isAt("(")) return parsePrimary();
436
- throw createError(`Expected a subscript but found ${describe(token)}`);
544
+ throw createUnexpected(`Expected a subscript but found ${describe(token)}`);
437
545
  };
438
546
 
439
547
  const isDefinitionAhead = (operator: Token) => {
@@ -470,6 +578,23 @@ export const parseModule = (source: string): Module => {
470
578
  token.kind === "Operator" &&
471
579
  (INFIX_OPERATORS.has(token.text) || PREFIX_OPERATOR_REFERENCES.has(token.text)) &&
472
580
  isStandalone(token);
581
+ // tlapm reads a prefix operator as the start of an expression, and (+), (-), (/) and (\X) as
582
+ // the operators \oplus, \ominus, \oslash and \otimes.
583
+ const isTlapmIncomplete =
584
+ isTlapm &&
585
+ token.kind !== "String" &&
586
+ (PREFIX_OPERATOR_REFERENCES.has(token.text) ||
587
+ (OPERATOR_ARGUMENTS.has(token.text) &&
588
+ token.text !== "'" &&
589
+ !POSTFIX_OPERATORS.has(token.text)) ||
590
+ (CIRCLED_OPERATORS.has(token.text) &&
591
+ peek(-1).text === "(" &&
592
+ peek(-1).line === token.line &&
593
+ peek(-1).endColumn + 1 === token.column &&
594
+ peek(1).text === ")" &&
595
+ peek(1).line === token.line &&
596
+ peek(1).column === token.endColumn + 1));
597
+ if (isTlapmIncomplete && isStandalone(token)) throw createError("incomplete expression", token);
473
598
  if (isOperatorAlone) {
474
599
  cursor.index++;
475
600
  return withRange(
@@ -524,6 +649,7 @@ export const parseModule = (source: string): Module => {
524
649
  return {
525
650
  kind: bullet.text === "/\\" ? "And" : "Or",
526
651
  operands,
652
+ isBulleted: true,
527
653
  location: getLocation(bullet),
528
654
  };
529
655
  };
@@ -561,7 +687,7 @@ export const parseModule = (source: string): Module => {
561
687
  while (isAt(".") || isAt("[")) {
562
688
  if (accept(".")) {
563
689
  if (!isFieldName(peek()))
564
- throw createError(`Expected a field name but found ${describe(peek())}`);
690
+ throw createUnexpected(`Expected a field name but found ${describe(peek())}`);
565
691
  path.push({ kind: "Field", name: advance().text });
566
692
  continue;
567
693
  }
@@ -637,7 +763,14 @@ export const parseModule = (source: string): Module => {
637
763
  }
638
764
  const names = parseSeparated(expectIdentifier);
639
765
  expect(":");
640
- return { kind: "UnboundedQuantifier", names, body: parseExpression(), isTemporal, location };
766
+ return {
767
+ kind: "UnboundedQuantifier",
768
+ quantifier: start.text === "\\A" || start.text === "\\AA" ? "Forall" : "Exists",
769
+ names,
770
+ body: parseExpression(),
771
+ isTemporal,
772
+ location,
773
+ };
641
774
  };
642
775
 
643
776
  const parseCase = (start: Token): Expression => {
@@ -691,7 +824,13 @@ export const parseModule = (source: string): Module => {
691
824
  const isFactStart = cursor.index === factStart.index;
692
825
  const token = advance();
693
826
  const location = getLocation(token);
827
+ if (isTlapm && token.kind === "Keyword" && token.text === "LAMBDA") {
828
+ const params = parseSeparated(expectIdentifier);
829
+ expect(":");
830
+ return { kind: "Lambda", params, body: parseExpression(), location };
831
+ }
694
832
  if (token.kind === "Number") {
833
+ if (isTlapm && /^\\[BOH]/.test(token.text)) throw createError("incomplete expression", token);
695
834
  return { kind: "Number", value: getNumeralValue(token.text), text: token.text, location };
696
835
  }
697
836
  if (token.kind === "Decimal") return { kind: "Decimal", text: token.text, location };
@@ -721,7 +860,7 @@ export const parseModule = (source: string): Module => {
721
860
  const instance = openModules
722
861
  .flatMap((module) => module.instances)
723
862
  .find((candidate) => candidate.name === name);
724
- if (instance && !isFactStart && instance.params.length !== args.length) {
863
+ if (instance && !isFactStart && !isTlapm && instance.params.length !== args.length) {
725
864
  throw createError(
726
865
  `The operator ${name} requires ${instance.params.length} arguments`,
727
866
  token,
@@ -740,9 +879,13 @@ export const parseModule = (source: string): Module => {
740
879
  const qualifiedArgs = selectors.flatMap((selector) =>
741
880
  "args" in selector ? selector.args : [],
742
881
  );
743
- return qualifiedArgs.length > 0
744
- ? { kind: "Apply", operator, args: qualifiedArgs, location }
745
- : { kind: "Identifier", name: operator, location };
882
+ if (qualifiedArgs.length === 0) return { kind: "Identifier", name: operator, location };
883
+ const argumentCounts = selectors.map((selector) =>
884
+ "args" in selector ? selector.args.length : 0,
885
+ );
886
+ return argumentCounts.slice(0, -1).some((count) => count > 0)
887
+ ? { kind: "Apply", operator, args: qualifiedArgs, argumentCounts, location }
888
+ : { kind: "Apply", operator, args: qualifiedArgs, location };
746
889
  }
747
890
  if (token.kind === "Operator" && isOperatorName(token.text) && isAt("("))
748
891
  return { kind: "Apply", operator: token.text, args: parseArguments(), location };
@@ -779,6 +922,8 @@ export const parseModule = (source: string): Module => {
779
922
  const instances: Instance[] = [];
780
923
  recursiveScopes.push(new Set());
781
924
  while (!accept("IN")) {
925
+ if (isTlapm && isAt("RECURSIVE"))
926
+ throw createError("Unexpected keyword RECURSIVE", getDefinitionEnd(peekRaw()));
782
927
  if (accept("RECURSIVE")) {
783
928
  declareRecursive();
784
929
  continue;
@@ -816,7 +961,7 @@ export const parseModule = (source: string): Module => {
816
961
  };
817
962
  }
818
963
  }
819
- throw createError(`Unexpected ${describe(token)}`, token);
964
+ throw createUnexpected(`Unexpected ${describe(token)}`, token);
820
965
  };
821
966
 
822
967
  const parsePrefix = (): Expression => {
@@ -1004,9 +1149,13 @@ export const parseModule = (source: string): Module => {
1004
1149
  location,
1005
1150
  };
1006
1151
  }
1152
+ const definitionStart = cursor.index;
1007
1153
  const params = parseParams();
1008
1154
  expect("==");
1009
- if (isAt("INSTANCE")) return parseInstance(name, params, isLocal, location);
1155
+ if (isAt("INSTANCE")) {
1156
+ const instance = parseInstance(name, params, isLocal, location);
1157
+ return { ...instance, definitionRange: getRange(definitionStart) };
1158
+ }
1010
1159
  return parseOperator(name, params);
1011
1160
  };
1012
1161
 
@@ -1016,6 +1165,7 @@ export const parseModule = (source: string): Module => {
1016
1165
  isLocal: boolean,
1017
1166
  location: SourceLocation,
1018
1167
  ): Instance => {
1168
+ const start = cursor.index;
1019
1169
  expect("INSTANCE");
1020
1170
  const module = expectIdentifier();
1021
1171
  const substitutions = accept("WITH")
@@ -1025,7 +1175,16 @@ export const parseModule = (source: string): Module => {
1025
1175
  return { name: target, expression: parseArgument() };
1026
1176
  })
1027
1177
  : [];
1028
- return { kind: "Instance", name, params, module, substitutions, isLocal, location };
1178
+ return {
1179
+ kind: "Instance",
1180
+ name,
1181
+ params,
1182
+ module,
1183
+ substitutions,
1184
+ isLocal,
1185
+ location,
1186
+ instanceRange: getRange(start),
1187
+ };
1029
1188
  };
1030
1189
 
1031
1190
  const parseDeclaration = (): Declaration => {
@@ -1049,10 +1208,13 @@ export const parseModule = (source: string): Module => {
1049
1208
 
1050
1209
  const parseDeclarations = () => parseSeparated(parseDeclaration);
1051
1210
 
1052
- const declareRecursive = () =>
1053
- parseDeclarations().forEach((declaration) =>
1211
+ const declareRecursive = () => {
1212
+ const declarations = parseDeclarations();
1213
+ declarations.forEach((declaration) =>
1054
1214
  recursiveScopes[recursiveScopes.length - 1].add(declaration.name),
1055
1215
  );
1216
+ return declarations;
1217
+ };
1056
1218
 
1057
1219
  const getClosingOffset = (offset: number) => {
1058
1220
  const depth = { value: 0, offset };
@@ -1200,7 +1362,30 @@ export const parseModule = (source: string): Module => {
1200
1362
  module.variables.some((variable) => variable.name === name),
1201
1363
  );
1202
1364
 
1365
+ // tlapm's definable: names and operators joined by !.
1366
+ const parseDefinable = (): Fact => {
1367
+ const start = cursor.index;
1368
+ const token = peek();
1369
+ const parts = parseSeparated(() => {
1370
+ const part = peek();
1371
+ const isName = part.kind === "Identifier" && !isTlapmKeyword(part);
1372
+ if (!isName && part.kind !== "Operator") {
1373
+ throw createUnexpected(`Expected a definition name but found ${describe(part)}`, part);
1374
+ }
1375
+ cursor.index++;
1376
+ return part.text;
1377
+ }, "!");
1378
+ const location = getLocation(token);
1379
+ return withRange(
1380
+ parts.length === 1 && token.kind === "Operator"
1381
+ ? { kind: "OperatorReference", operator: token.text, location }
1382
+ : { kind: "Identifier", name: parts.join("!"), location },
1383
+ start,
1384
+ );
1385
+ };
1386
+
1203
1387
  const parseDefs = () => {
1388
+ if (isTlapm) return parseSeparated(parseDefinable);
1204
1389
  const defs = parseSeparated(parseFact);
1205
1390
  const invalid = defs.find((def) => def.kind === "Identifier" && isDeclaredSymbol(def.name));
1206
1391
  if (invalid) {
@@ -1218,13 +1403,16 @@ export const parseModule = (source: string): Module => {
1218
1403
  const start = cursor.index;
1219
1404
  const keyword = advance();
1220
1405
  const isOnly = keyword.text === "USE" && accept("ONLY");
1406
+ const keywordRange = getRange(start);
1221
1407
  const use: UseOrHide = {
1222
1408
  kind: keyword.text === "USE" ? "Use" : "Hide",
1223
1409
  isOnly,
1224
1410
  ...parseFacts(),
1411
+ keyword: keywordRange,
1225
1412
  location: getLocation(keyword),
1226
1413
  };
1227
- const invalid = use.kind === "Hide" ? use.facts.find((fact) => !isHideFact(fact)) : undefined;
1414
+ const invalid =
1415
+ use.kind === "Hide" && !isTlapm ? use.facts.find((fact) => !isHideFact(fact)) : undefined;
1228
1416
  if (invalid) {
1229
1417
  throw createError(
1230
1418
  "The only expression allowed as a fact in a HIDE is the name of a theorem, assumption, or step",
@@ -1240,7 +1428,7 @@ export const parseModule = (source: string): Module => {
1240
1428
  const isNew = accept("NEW");
1241
1429
  const level = peek().kind === "Keyword" ? NEW_SYMBOL_LEVELS.get(peek().text) : undefined;
1242
1430
  if (!isNew && level === undefined) {
1243
- throw createError(`Expected a declaration but found ${describe(peek())}`);
1431
+ throw createUnexpected(`Expected a declaration but found ${describe(peek())}`);
1244
1432
  }
1245
1433
  if (level !== undefined) advance();
1246
1434
  if (level === "Variable") {
@@ -1347,7 +1535,7 @@ export const parseModule = (source: string): Module => {
1347
1535
  const parseStepStart = (isQed: boolean) => {
1348
1536
  const step = peek();
1349
1537
  if (step.kind !== "Step") {
1350
- throw createError(`Expected a step number but found ${describe(step)}`);
1538
+ throw createUnexpected(`Expected a step number but found ${describe(step)}`, step);
1351
1539
  }
1352
1540
  if (!isCorrectLevel(step)) {
1353
1541
  throw createError(`${isQed ? "QED step" : "step"}'s number has bad level`, step);
@@ -1356,7 +1544,7 @@ export const parseModule = (source: string): Module => {
1356
1544
  throw createError("<*> and <+> cannot be used for a named step", step);
1357
1545
  }
1358
1546
  cursor.index++;
1359
- return { name: getStepName(step), location: getLocation(step) };
1547
+ return { name: getStepName(step), token: step.text, location: getLocation(step) };
1360
1548
  };
1361
1549
 
1362
1550
  // The names an assertion assumes stay in its own proof, and other steps declare theirs for the rest
@@ -1367,6 +1555,7 @@ export const parseModule = (source: string): Module => {
1367
1555
  const scope = new Set<string>();
1368
1556
  proofScopes.push(scope);
1369
1557
  const body: StepBody = accept("QED") ? { kind: "Qed" } : parseStepBody();
1558
+ const head = getRange(start);
1370
1559
  const hasProof = isProofStart(peek());
1371
1560
  if (hasProof && !PROVABLE_STEPS.has(body.kind)) {
1372
1561
  throw createError("Proof of a step that does not take a proof");
@@ -1374,7 +1563,7 @@ export const parseModule = (source: string): Module => {
1374
1563
  const proof = hasProof ? parseProof() : undefined;
1375
1564
  proofScopes.pop();
1376
1565
  if (body.kind !== "Assert" || body.isSuffices) scope.forEach(declareProofSymbol);
1377
- return withRange({ ...step, body, proof }, start);
1566
+ return withRange({ ...step, body, head, proof }, start);
1378
1567
  };
1379
1568
 
1380
1569
  const parseProofBody = (): Proof => {
@@ -1382,7 +1571,8 @@ export const parseModule = (source: string): Module => {
1382
1571
  if (isAt("BY") || (isAt("PROOF") && isAt("BY", 1))) {
1383
1572
  accept("PROOF");
1384
1573
  advance();
1385
- return { kind: "By", isOnly: accept("ONLY"), ...parseFacts(), location };
1574
+ const keyword = getRange(cursor.index - 1);
1575
+ return { kind: "By", isOnly: accept("ONLY"), ...parseFacts(), keyword, location };
1386
1576
  }
1387
1577
  const terminal = peek(isAt("PROOF") ? 1 : 0);
1388
1578
  if (
@@ -1391,7 +1581,8 @@ export const parseModule = (source: string): Module => {
1391
1581
  ) {
1392
1582
  accept("PROOF");
1393
1583
  advance();
1394
- return { kind: terminal.text === "OBVIOUS" ? "Obvious" : "Omitted", location };
1584
+ const keyword = getRange(cursor.index - 1);
1585
+ return { kind: terminal.text === "OBVIOUS" ? "Obvious" : "Omitted", keyword, location };
1395
1586
  }
1396
1587
  accept("PROOF");
1397
1588
  const steps = [parseStep()];
@@ -1421,9 +1612,12 @@ export const parseModule = (source: string): Module => {
1421
1612
  return { statement, proof: isProofStart(peek()) ? parseProof() : undefined };
1422
1613
  });
1423
1614
  const statement = parsed.statement;
1424
- module.theorems.push(
1425
- withRange({ name, statement, proof: parsed.proof, location: getLocation(keyword) }, start),
1615
+ const theorem = withRange(
1616
+ { name, statement, proof: parsed.proof, location: getLocation(keyword) },
1617
+ start,
1426
1618
  );
1619
+ module.theorems.push(theorem);
1620
+ module.units.push({ kind: "Theorem", theorem });
1427
1621
  if (name === undefined || statement.kind === "AssumeProve") return;
1428
1622
  module.definitions.push(
1429
1623
  withRange(
@@ -1467,7 +1661,8 @@ export const parseModule = (source: string): Module => {
1467
1661
  const module: Module = {
1468
1662
  name,
1469
1663
  modules: [],
1470
- extends: accept("EXTENDS") ? parseSeparated(expectIdentifier) : [],
1664
+ extends: [],
1665
+ extendsRanges: [],
1471
1666
  constants: [],
1472
1667
  variables: [],
1473
1668
  definitions: [],
@@ -1476,20 +1671,38 @@ export const parseModule = (source: string): Module => {
1476
1671
  location,
1477
1672
  theorems: [],
1478
1673
  uses: [],
1674
+ units: [],
1479
1675
  };
1480
1676
  openModules.push(module);
1481
1677
  recursiveScopes.push(new Set());
1678
+ if (accept("EXTENDS")) {
1679
+ parseSeparated(() => {
1680
+ const start = cursor.index;
1681
+ module.extends.push(expectIdentifier());
1682
+ module.extendsRanges.push(getRange(start));
1683
+ });
1684
+ }
1482
1685
  while (true) {
1483
1686
  const tokenIndex = cursor.index;
1687
+ unitStart.index = tokenIndex;
1484
1688
  const token = peek();
1485
1689
  if (token.kind === "ModuleEnd" || token.kind === "End") {
1690
+ if (token.kind === "ModuleEnd") {
1691
+ location.range = {
1692
+ module: name,
1693
+ start: { line: location.line, column: location.column },
1694
+ end: { line: token.line, column: token.endColumn },
1695
+ };
1696
+ }
1486
1697
  openModules.pop();
1487
1698
  recursiveScopes.pop();
1488
1699
  return module;
1489
1700
  }
1490
1701
  if (token.kind === "Separator") {
1491
1702
  if (isAt("MODULE", 1)) {
1492
- module.modules.push(parseModuleBody());
1703
+ const submodule = parseModuleBody();
1704
+ module.modules.push(submodule);
1705
+ module.units.push({ kind: "Module", module: submodule });
1493
1706
  expect("====");
1494
1707
  continue;
1495
1708
  }
@@ -1498,26 +1711,31 @@ export const parseModule = (source: string): Module => {
1498
1711
  }
1499
1712
  switch (token.text) {
1500
1713
  case "CONSTANT":
1501
- case "CONSTANTS":
1714
+ case "CONSTANTS": {
1502
1715
  advance();
1503
- module.constants.push(...parseDeclarations());
1716
+ const declarations = parseDeclarations();
1717
+ module.constants.push(...declarations);
1718
+ module.units.push({ kind: "Constants", declarations });
1504
1719
  continue;
1720
+ }
1505
1721
  case "VARIABLE":
1506
- case "VARIABLES":
1722
+ case "VARIABLES": {
1507
1723
  advance();
1508
- module.variables.push(
1509
- ...parseSeparated(() => {
1510
- const start = cursor.index;
1511
- const location = getLocation(peek());
1512
- return withRange({ name: expectIdentifier(), arity: 0, location }, start);
1513
- }),
1514
- );
1724
+ const declarations = parseSeparated(() => {
1725
+ const start = cursor.index;
1726
+ const location = getLocation(peek());
1727
+ return withRange({ name: expectIdentifier(), arity: 0, location }, start);
1728
+ });
1729
+ module.variables.push(...declarations);
1730
+ module.units.push({ kind: "Variables", declarations });
1515
1731
  continue;
1732
+ }
1516
1733
  case "ASSUME":
1517
1734
  case "ASSUMPTION":
1518
1735
  case "AXIOM": {
1519
1736
  const assumption = parseAssumption();
1520
1737
  module.assumptions.push(assumption);
1738
+ module.units.push({ kind: "Assumption", assumption });
1521
1739
  if (assumption.name !== undefined) {
1522
1740
  module.definitions.push({
1523
1741
  kind: "Operator",
@@ -1532,29 +1750,32 @@ export const parseModule = (source: string): Module => {
1532
1750
  continue;
1533
1751
  }
1534
1752
  case "USE":
1535
- case "HIDE":
1536
- module.uses.push(parseUseOrHide());
1753
+ case "HIDE": {
1754
+ const use = parseUseOrHide();
1755
+ module.uses.push(use);
1756
+ module.units.push({ kind: "Use", use });
1537
1757
  continue;
1758
+ }
1538
1759
  case "RECURSIVE":
1539
1760
  advance();
1540
- declareRecursive();
1761
+ module.units.push({ kind: "Recursive", declarations: declareRecursive() });
1541
1762
  continue;
1542
1763
  case "INSTANCE":
1543
- module.instances.push(
1764
+ addDefinitionUnit(
1765
+ module,
1544
1766
  withRange(parseInstance(undefined, [], false, getLocation(token)), tokenIndex),
1545
1767
  );
1546
1768
  continue;
1547
1769
  case "LOCAL": {
1548
1770
  advance();
1549
1771
  if (isAt("INSTANCE")) {
1550
- module.instances.push(
1772
+ addDefinitionUnit(
1773
+ module,
1551
1774
  withRange(parseInstance(undefined, [], true, getLocation(token)), tokenIndex),
1552
1775
  );
1553
1776
  continue;
1554
1777
  }
1555
- const unit = parseDefinition(true);
1556
- if (unit.kind === "Instance") module.instances.push(unit);
1557
- else module.definitions.push(unit);
1778
+ addDefinitionUnit(module, parseDefinition(true));
1558
1779
  continue;
1559
1780
  }
1560
1781
  }
@@ -1562,11 +1783,23 @@ export const parseModule = (source: string): Module => {
1562
1783
  parseTheorem(module);
1563
1784
  continue;
1564
1785
  }
1565
- const unit = parseDefinition(false);
1566
- if (unit.kind === "Instance") module.instances.push(unit);
1567
- else module.definitions.push(unit);
1786
+ addDefinitionUnit(module, parseDefinition(false));
1568
1787
  }
1569
1788
  };
1570
1789
 
1790
+ const addDefinitionUnit = (module: Module, unit: Definition | Instance) => {
1791
+ if (unit.kind === "Instance") {
1792
+ module.instances.push(unit);
1793
+ module.units.push({ kind: "Instance", instance: unit });
1794
+ return;
1795
+ }
1796
+ module.definitions.push(unit);
1797
+ module.units.push({ kind: "Definition", definition: unit });
1798
+ };
1799
+
1571
1800
  return parseModuleBody();
1572
1801
  };
1802
+
1803
+ export const parseModule = (source: string) => parse(source, false);
1804
+
1805
+ export const parseTlapmModule = (source: string) => parse(source, true);