From fdaf12a35d06602d470b0c45ed8815eca7df6c03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?VESSE=20L=C3=A9o?= <leovesse@gmail.com>
Date: Thu, 13 Mar 2025 13:59:17 +0000
Subject: [PATCH] feat(tlaplus) : added `tlaplus` config + grammar (#13081)

---
 book/src/generated/lang-support.md     |   1 +
 languages.toml                         |  15 ++
 runtime/queries/tlaplus/highlights.scm | 225 +++++++++++++++++++++++++
 runtime/queries/tlaplus/locals.scm     |  70 ++++++++
 4 files changed, 311 insertions(+)
 create mode 100644 runtime/queries/tlaplus/highlights.scm
 create mode 100644 runtime/queries/tlaplus/locals.scm

diff --git a/book/src/generated/lang-support.md b/book/src/generated/lang-support.md
index 0d4bc7f8..b44faf9d 100644
--- a/book/src/generated/lang-support.md
+++ b/book/src/generated/lang-support.md
@@ -225,6 +225,7 @@
 | textproto | ✓ | ✓ | ✓ |  |
 | tfvars | ✓ |  | ✓ | `terraform-ls` |
 | thrift | ✓ |  |  |  |
+| tlaplus | ✓ |  |  |  |
 | todotxt | ✓ |  |  |  |
 | toml | ✓ | ✓ |  | `taplo` |
 | tsq | ✓ |  |  | `ts_query_ls` |
diff --git a/languages.toml b/languages.toml
index 3e790d76..8eae5a51 100644
--- a/languages.toml
+++ b/languages.toml
@@ -4226,3 +4226,18 @@ language-servers = ["sourcepawn-studio"]
 [[grammar]]
 name = "sourcepawn"
 source = { git = "https://github.com/nilshelmig/tree-sitter-sourcepawn", rev = "f2af8d0dc14c6790130cceb2a20027eb41a8297c" }
+
+[[language]]
+name = "tlaplus"
+scope = "scope.tlaplus"
+injection-regex = "tla"
+file-types = ["tla"]
+comment-tokens = "\\*"
+block-comment-tokens = {start = "(*", end="*)"}
+indent = {tab-width = 4, unit = " "}
+formatter = {command = "tlafmt", args = ["--stdin"]}
+
+[[grammar]]
+name = "tlaplus"
+source = { git = "https://github.com/tlaplus-community/tree-sitter-tlaplus", rev = "4ba91b07b97741a67f61221d0d50e6d962e4987e"}
+
diff --git a/runtime/queries/tlaplus/highlights.scm b/runtime/queries/tlaplus/highlights.scm
new file mode 100644
index 00000000..f9eaa79d
--- /dev/null
+++ b/runtime/queries/tlaplus/highlights.scm
@@ -0,0 +1,225 @@
+; ; Intended for consumption by GitHub and the tree-sitter highlight command
+; ; Default capture names found here:
+; ; https://github.com/tree-sitter/tree-sitter/blob/f5d1c0b8609f8697861eab352ead44916c068c74/cli/src/highlight.rs#L150-L171
+; ; In this file, captures defined earlier take precedence over captures defined later.
+
+; TLA⁺ 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
+[
+  "algorithm"
+  "assert"
+  "await"
+  "begin"
+  "call"
+  "define"
+  "either"
+  "else"
+  "elsif"
+  "end"
+  "fair"
+  "goto"
+  "if"
+  "macro"
+  "or"
+  "print"
+  "procedure"
+  "process"
+  "variable"
+  "variables"
+  "when"
+  "with"
+  "then"
+  (pcal_algorithm_start)
+  (pcal_end_either)
+  (pcal_end_if)
+  (pcal_return)
+  (pcal_skip)
+  (pcal_process ("="))
+  (pcal_with ("="))
+] @keyword
+
+; Literals
+(binary_number (format) @keyword)
+(binary_number (value) @constant.numeric)
+(boolean) @constant.builtin.boolean
+(boolean_set) @type
+(hex_number (format) @keyword)
+(hex_number (value) @constant.numeric)
+(int_number_set) @type
+(nat_number) @constant.numeric.integer
+(nat_number_set) @type
+(octal_number (format) @keyword)
+(octal_number (value) @constant.numeric)
+(real_number) @constant.numeric.integer
+(real_number_set) @type
+(string) @string
+(escape_char) @string.special.symbol
+(string_set) @type
+
+; Namespaces and includes
+(extends (identifier_ref) @namespace)
+(instance (identifier_ref) @namespace)
+(module name: (_) @namespace)
+(pcal_algorithm name: (identifier) @namespace)
+
+; Constants and variables
+(constant_declaration (identifier) @constant)
+(constant_declaration (operator_declaration name: (_) @constant))
+(pcal_var_decl (identifier) @variable.builtin)
+(pcal_with (identifier) @variable.parameter)
+((".") . (identifier) @attribute)
+(record_literal (identifier) @attribute)
+(set_of_records (identifier) @attribute)
+(variable_declaration (identifier) @variable.builtin)
+
+; 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)
+
+; Operators, functions, and macros
+(function_definition name: (identifier) @function)
+(module_definition name: (_) @namespace)
+(operator_definition name: (_) @operator)
+(pcal_macro_decl name: (identifier) @function)
+(pcal_macro_call name: (identifier) @function)
+(pcal_proc_decl name: (identifier) @function)
+(pcal_process name: (identifier) @function)
+(recursive_declaration (identifier) @operator)
+(recursive_declaration (operator_declaration name: (_) @operator))
+
+; 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) @tag)
+(proof_step_id (name) @tag)
+(proof_step_id ">" @punctuation.bracket)
+(proof_step_ref "<" @punctuation.bracket)
+(proof_step_ref (level) @tag)
+(proof_step_ref (name) @tag)
+(proof_step_ref ">" @punctuation.bracket)
+(take_proof_step (identifier) @variable.parameter)
+(theorem name: (identifier) @constant)
+
+; Comments and tags
+(block_comment "(*" @comment.block)
+(block_comment "*)" @comment.block)
+(block_comment_text) @comment.block
+(comment) @comment
+(single_line) @comment
+(_ label: (identifier) @tag)
+(label name: (_) @tag)
+(pcal_goto statement: (identifier) @tag)
+
+; Put these last so they are overridden by everything else
+(bound_infix_op symbol: (_) @function.builtin)
+(bound_nonfix_op symbol: (_) @function.builtin)
+(bound_postfix_op symbol: (_) @function.builtin)
+(bound_prefix_op symbol: (_) @function.builtin)
+((prefix_op_symbol) @function.builtin)
+((infix_op_symbol) @function.builtin)
+((postfix_op_symbol) @function.builtin)
diff --git a/runtime/queries/tlaplus/locals.scm b/runtime/queries/tlaplus/locals.scm
new file mode 100644
index 00000000..aee8d612
--- /dev/null
+++ b/runtime/queries/tlaplus/locals.scm
@@ -0,0 +1,70 @@
+; 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
+
+; Definitions
+(choose (identifier) @local.definition)
+(choose (tuple_of_identifiers (identifier) @local.definition))
+(constant_declaration (identifier) @local.definition)
+(constant_declaration (operator_declaration name: (_) @local.definition))
+(function_definition name: (identifier) @local.definition)
+(lambda (identifier) @local.definition)
+(module_definition name: (_) @local.definition)
+(module_definition parameter: (identifier) @local.definition)
+(module_definition parameter: (operator_declaration name: (_) @local.definition))
+(operator_definition name: (_) @local.definition)
+(operator_definition parameter: (identifier) @local.definition)
+(operator_definition parameter: (operator_declaration name: (_) @local.definition))
+(quantifier_bound (identifier) @local.definition)
+(quantifier_bound (tuple_of_identifiers (identifier) @local.definition))
+(unbounded_quantification (identifier) @local.definition)
+(variable_declaration (identifier) @local.definition)
+
+; Proof scopes and definitions
+[
+  (non_terminal_proof)
+  (suffices_proof_step)
+  (theorem)
+] @local.scope
+
+(assume_prove (new (identifier) @local.definition))
+(assume_prove (new (operator_declaration name: (_) @local.definition)))
+(assumption name: (identifier) @local.definition)
+(pick_proof_step (identifier) @local.definition)
+(take_proof_step (identifier) @local.definition)
+(theorem name: (identifier) @local.definition)
+
+;PlusCal scopes and definitions
+[
+  (pcal_algorithm)
+  (pcal_macro)
+  (pcal_procedure)
+  (pcal_with)
+] @local.scope
+
+(pcal_macro_decl parameter: (identifier) @local.definition)
+(pcal_proc_var_decl (identifier) @local.definition)
+(pcal_var_decl (identifier) @local.definition)
+(pcal_with (identifier) @local.definition)
+
+; 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)