feat!: drop modules, general refactor and cleanup

4 files changed, 557 insertions, 0 deletions

diff --git a/runtime/queries/tlaplus/folds.scm b/runtime/queries/tlaplus/folds.scm
new file mode 100644
index 000000000..2ca0168c0
--- /dev/null
+++ b/runtime/queries/tlaplus/folds.scm
@@ -0,0 +1,5 @@
+[
+  (extramodular_text)
+  (block_comment)
+  (non_terminal_proof)
+] @fold
diff --git a/runtime/queries/tlaplus/highlights.scm b/runtime/queries/tlaplus/highlights.scm
new file mode 100644
index 000000000..2eca927e7
--- /dev/null
+++ b/runtime/queries/tlaplus/highlights.scm
@@ -0,0 +1,380 @@
+; Keywords
+[
+  "ACTION"
+  "ASSUME"
+  "ASSUMPTION"
+  "AXIOM"
+  "BY"
+  "CASE"
+  "CHOOSE"
+  "CONSTANT"
+  "CONSTANTS"
+  "COROLLARY"
+  "DEF"
+  "DEFINE"
+  "DEFS"
+  "ELSE"
+  "EXCEPT"
+  "EXTENDS"
+  "HAVE"
+  "HIDE"
+  "IF"
+  "IN"
+  "INSTANCE"
+  "LAMBDA"
+  "LEMMA"
+  "LET"
+  "LOCAL"
+  "MODULE"
+  "NEW"
+  "OBVIOUS"
+  "OMITTED"
+  "ONLY"
+  "OTHER"
+  "PICK"
+  "PROOF"
+  "PROPOSITION"
+  "PROVE"
+  "QED"
+  "RECURSIVE"
+  "SF_"
+  "STATE"
+  "SUFFICES"
+  "TAKE"
+  "TEMPORAL"
+  "THEN"
+  "THEOREM"
+  "USE"
+  "VARIABLE"
+  "VARIABLES"
+  "WF_"
+  "WITH"
+  "WITNESS"
+  (address)
+  (all_map_to)
+  (assign)
+  (case_arrow)
+  (case_box)
+  (def_eq)
+  (exists)
+  (forall)
+  (gets)
+  (label_as)
+  (maps_to)
+  (set_in)
+  (temporal_exists)
+  (temporal_forall)
+] @keyword
+
+;  Pluscal keywords
+[
+  (pcal_algorithm_start)
+  "algorithm"
+  "assert"
+  "begin"
+  "call"
+  "define"
+  "end"
+  "fair"
+  "goto"
+  "macro"
+  "or"
+  "procedure"
+  "process"
+  (pcal_skip)
+  "variable"
+  "variables"
+  "when"
+  "with"
+] @keyword
+
+"await" @keyword.coroutine
+
+(pcal_with
+  "=" @keyword)
+
+(pcal_process
+  "=" @keyword)
+
+[
+  "if"
+  "then"
+  "else"
+  "elsif"
+  (pcal_end_if)
+  "either"
+  (pcal_end_either)
+] @keyword.conditional
+
+[
+  "while"
+  "do"
+  (pcal_end_while)
+  "with"
+  (pcal_end_with)
+] @keyword.repeat
+
+(pcal_return) @keyword.return
+
+"print" @function.macro
+
+; Literals
+(binary_number
+  (format) @keyword)
+
+(binary_number
+  (value) @number)
+
+(boolean) @boolean
+
+(boolean_set) @type
+
+(hex_number
+  (format) @keyword)
+
+(hex_number
+  (value) @number)
+
+(int_number_set) @type
+
+(nat_number) @number
+
+(nat_number_set) @type
+
+(octal_number
+  (format) @keyword)
+
+(octal_number
+  (value) @number)
+
+(real_number) @number
+
+(real_number_set) @type
+
+(string) @string
+
+(escape_char) @string.escape
+
+(string_set) @type
+
+; Namespaces
+(extends
+  (identifier_ref) @module)
+
+(instance
+  (identifier_ref) @module)
+
+(module
+  name: (identifier) @module)
+
+(pcal_algorithm
+  name: (identifier) @module)
+
+; Operators, functions, and macros
+(bound_infix_op
+  symbol: (_) @operator)
+
+(bound_nonfix_op
+  symbol: (_) @operator)
+
+(bound_postfix_op
+  symbol: (_) @operator)
+
+(bound_prefix_op
+  symbol: (_) @operator)
+
+(prefix_op_symbol) @operator
+
+(infix_op_symbol) @operator
+
+(postfix_op_symbol) @operator
+
+(function_definition
+  name: (identifier) @function)
+
+(module_definition
+  name: (_) @keyword.import)
+
+(operator_definition
+  name: (_) @function.macro)
+
+(pcal_macro_decl
+  name: (identifier) @function.macro)
+
+(pcal_macro_call
+  name: (identifier) @function.macro)
+
+(pcal_proc_decl
+  name: (identifier) @function.macro)
+
+(pcal_process
+  name: (identifier) @function)
+
+(recursive_declaration
+  (identifier) @function.macro)
+
+(recursive_declaration
+  (operator_declaration
+    name: (_) @function.macro))
+
+; Constants and variables
+(constant_declaration
+  (identifier) @constant)
+
+(constant_declaration
+  (operator_declaration
+    name: (_) @constant))
+
+(pcal_var_decl
+  (identifier) @variable)
+
+(pcal_with
+  (identifier) @variable.parameter)
+
+("."
+  .
+  (identifier) @attribute)
+
+(record_literal
+  (identifier) @attribute)
+
+(set_of_records
+  (identifier) @attribute)
+
+(variable_declaration
+  (identifier) @variable)
+
+; Parameters
+(choose
+  (identifier) @variable.parameter)
+
+(choose
+  (tuple_of_identifiers
+    (identifier) @variable.parameter))
+
+(lambda
+  (identifier) @variable.parameter)
+
+(module_definition
+  (operator_declaration
+    name: (_) @variable.parameter))
+
+(module_definition
+  parameter: (identifier) @variable.parameter)
+
+(operator_definition
+  (operator_declaration
+    name: (_) @variable.parameter))
+
+(operator_definition
+  parameter: (identifier) @variable.parameter)
+
+(pcal_macro_decl
+  parameter: (identifier) @variable.parameter)
+
+(pcal_proc_var_decl
+  (identifier) @variable.parameter)
+
+(quantifier_bound
+  (identifier) @variable.parameter)
+
+(quantifier_bound
+  (tuple_of_identifiers
+    (identifier) @variable.parameter))
+
+(unbounded_quantification
+  (identifier) @variable.parameter)
+
+; Delimiters
+[
+  (langle_bracket)
+  (rangle_bracket)
+  (rangle_bracket_sub)
+  "{"
+  "}"
+  "["
+  "]"
+  "]_"
+  "("
+  ")"
+] @punctuation.bracket
+
+[
+  ","
+  ":"
+  "."
+  "!"
+  ";"
+  (bullet_conj)
+  (bullet_disj)
+  (prev_func_val)
+  (placeholder)
+] @punctuation.delimiter
+
+; Proofs
+(assume_prove
+  (new
+    (identifier) @variable.parameter))
+
+(assume_prove
+  (new
+    (operator_declaration
+      name: (_) @variable.parameter)))
+
+(assumption
+  name: (identifier) @constant)
+
+(pick_proof_step
+  (identifier) @variable.parameter)
+
+(proof_step_id
+  "<" @punctuation.bracket)
+
+(proof_step_id
+  (level) @label)
+
+(proof_step_id
+  (name) @label)
+
+(proof_step_id
+  ">" @punctuation.bracket)
+
+(proof_step_ref
+  "<" @punctuation.bracket)
+
+(proof_step_ref
+  (level) @label)
+
+(proof_step_ref
+  (name) @label)
+
+(proof_step_ref
+  ">" @punctuation.bracket)
+
+(take_proof_step
+  (identifier) @variable.parameter)
+
+(theorem
+  name: (identifier) @constant)
+
+; Comments and tags
+(block_comment
+  "(*" @comment)
+
+(block_comment
+  "*)" @comment)
+
+(block_comment_text) @comment @spell
+
+(comment) @comment @spell
+
+(single_line) @comment
+
+(_
+  label: (identifier) @label)
+
+(label
+  name: (_) @label)
+
+(pcal_goto
+  statement: (identifier) @label)
diff --git a/runtime/queries/tlaplus/injections.scm b/runtime/queries/tlaplus/injections.scm
new file mode 100644
index 000000000..fbeff500b
--- /dev/null
+++ b/runtime/queries/tlaplus/injections.scm
@@ -0,0 +1,5 @@
+([
+  (comment)
+  (block_comment_text)
+] @injection.content
+  (#set! injection.language "comment"))
diff --git a/runtime/queries/tlaplus/locals.scm b/runtime/queries/tlaplus/locals.scm
new file mode 100644
index 000000000..e105351ef
--- /dev/null
+++ b/runtime/queries/tlaplus/locals.scm
@@ -0,0 +1,167 @@
+; TLA+ scopes and definitions
+[
+  (bounded_quantification)
+  (choose)
+  (function_definition)
+  (function_literal)
+  (lambda)
+  (let_in)
+  (module)
+  (module_definition)
+  (operator_definition)
+  (set_filter)
+  (set_map)
+  (unbounded_quantification)
+] @local.scope
+
+(choose
+  (identifier) @local.definition.parameter)
+
+(choose
+  (tuple_of_identifiers
+    (identifier) @local.definition.parameter))
+
+(constant_declaration
+  (identifier) @local.definition.constant)
+
+(constant_declaration
+  (operator_declaration
+    name: (_) @local.definition.constant))
+
+(function_definition
+  name: (identifier) @local.definition.function
+  (#set! definition.function.scope "parent"))
+
+(lambda
+  (identifier) @local.definition.parameter)
+
+(module_definition
+  name: (_) @local.definition.import
+  (#set! definition.import.scope "parent"))
+
+(module_definition
+  parameter: (identifier) @local.definition.parameter)
+
+(module_definition
+  parameter: (operator_declaration
+    name: (_) @local.definition.parameter))
+
+(operator_definition
+  name: (_) @local.definition.macro
+  (#set! definition.macro.scope "parent"))
+
+(operator_definition
+  parameter: (identifier) @local.definition.parameter)
+
+(operator_definition
+  parameter: (operator_declaration
+    name: (_) @local.definition.parameter))
+
+(quantifier_bound
+  (identifier) @local.definition.parameter)
+
+(quantifier_bound
+  (tuple_of_identifiers
+    (identifier) @local.definition.parameter))
+
+(unbounded_quantification
+  (identifier) @local.definition.parameter)
+
+(variable_declaration
+  (identifier) @local.definition.var)
+
+; Proof scopes and definitions
+[
+  (non_terminal_proof)
+  (suffices_proof_step)
+  (theorem)
+] @local.scope
+
+(assume_prove
+  (new
+    (identifier) @local.definition.parameter))
+
+(assume_prove
+  (new
+    (operator_declaration
+      name: (_) @local.definition.parameter)))
+
+(assumption
+  name: (identifier) @local.definition.constant)
+
+(pick_proof_step
+  (identifier) @local.definition.parameter)
+
+(take_proof_step
+  (identifier) @local.definition.parameter)
+
+(theorem
+  name: (identifier) @local.definition.constant
+  (#set! definition.constant.scope "parent"))
+
+; PlusCal scopes and definitions
+[
+  (pcal_algorithm)
+  (pcal_macro)
+  (pcal_procedure)
+  (pcal_with)
+] @local.scope
+
+(pcal_macro_decl
+  parameter: (identifier) @local.definition.parameter)
+
+(pcal_proc_var_decl
+  (identifier) @local.definition.parameter)
+
+(pcal_var_decl
+  (identifier) @local.definition.var)
+
+(pcal_with
+  (identifier) @local.definition.parameter)
+
+; Built-in PlusCal variables
+(pcal_algorithm_body
+  [
+    (_
+      (identifier_ref) @local.definition.var)
+    (_
+      (_
+        (identifier_ref) @local.definition.var))
+    (_
+      (_
+        (_
+          (identifier_ref) @local.definition.var)))
+    (_
+      (_
+        (_
+          (_
+            (identifier_ref) @local.definition.var))))
+    (_
+      (_
+        (_
+          (_
+            (_
+              (identifier_ref) @local.definition.var)))))
+  ]
+  (#any-of? @local.definition.var "self" "pc" "stack"))
+
+; References
+(identifier_ref) @local.reference
+
+(prefix_op_symbol) @local.reference
+
+(bound_prefix_op
+  symbol: (_) @local.reference)
+
+(infix_op_symbol) @local.reference
+
+(bound_infix_op
+  symbol: (_) @local.reference)
+
+(postfix_op_symbol) @local.reference
+
+(bound_postfix_op
+  symbol: (_) @local.reference)
+
+(bound_nonfix_op
+  symbol: (_) @local.reference)