update TLA+ revision, add queries for PlusCal (#2344)

4 files changed, 175 insertions, 64 deletions

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 6b40f6e3e..9ae975136 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -194,6 +194,7 @@ Used for xml-like tags
 
 ```
 @definition for various definitions
+@definition.constant
 @definition.function
 @definition.method
 @definition.var
diff --git a/queries/tlaplus/highlights.scm b/queries/tlaplus/highlights.scm
index c491ea181..01b302f23 100644
--- a/queries/tlaplus/highlights.scm
+++ b/queries/tlaplus/highlights.scm
@@ -1,4 +1,7 @@
-; highlights.scm
+; ; highlights.scm
+; ; Default capture names for tree-sitter highlight found here:
+; ; https://github.com/nvim-treesitter/nvim-treesitter/blob/e473630fe0872cb0ed97cd7085e724aa58bc1c84/lua/nvim-treesitter/highlight.lua#L14-L104
+
 
 ; Keywords
 [
@@ -57,47 +60,125 @@
   "WF_"
   "WITH"
   "WITNESS"
-  (forall)
+  (address)
+  (all_map_to)
+  (assign)
+  (case_arrow)
+  (case_box)
+  (def_eq)
   (exists)
-  (temporal_forall)
-  (temporal_exists)
-  (set_in)
+  (forall)
   (gets)
-  (case_box)
-  (case_arrow)
   (label_as)
+  (maps_to)
+  (set_in)
+  (temporal_exists)
+  (temporal_forall)
 ] @keyword
+;  Pluscal keywords
 [
-  (def_eq)
-  (maps_to)
-  (all_map_to)
-] @keyword.function
+  (pcal_algorithm_start)
+  "assert"
+  "await"
+  "begin"
+  "call"
+  "define"
+  "end"
+  "fair"
+  "goto"
+  "macro"
+  "or"
+  "procedure"
+  "process"
+  "skip"
+  "variable"
+  "variables"
+  "when"
+  "with"
+] @keyword
+(pcal_with ("=") @keyword)
+(pcal_process ("=") @keyword)
+[ 
+  "if" 
+  "then" 
+  "else" 
+  "elsif"
+  (pcal_end_if)
+  "either"
+  (pcal_end_either)
+] @conditional
+[ 
+  "while" 
+  "do" 
+  (pcal_end_while) 
+  "with" 
+  (pcal_end_with)
+] @repeat
+("return") @keyword.return
+("print") @function.macro
+
 
 ; Literals
-(nat_number) @number
-(real_number) @float
-(binary_number (format) @string.special)
-(octal_number (format) @string.special)
-(hex_number (format) @string.special)
+(binary_number (format) @keyword)
 (binary_number (value) @number)
-(octal_number (value) @number)
-(hex_number (value) @number)
-(string) @string
-(escape_char) @string.escape
 (boolean) @boolean
-(string_set) @type
 (boolean_set) @type
-(nat_number_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
+(string_set) @type
 
-; Comments
-(comment) @comment
-(block_comment) @comment
-(single_line) @comment
-(extramodular_text) @text
+; Namespaces and includes
+(extends (identifier_ref) @include)
+(module name: (_) @namespace)
+(pcal_algorithm name: (identifier) @namespace)
 
-; Punctuation and Delimiters
+; Operators and functions
+(bound_infix_op symbol: (_) @function.builtin)
+(bound_nonfix_op (infix_op_symbol) @operator)
+(bound_nonfix_op (postfix_op_symbol) @operator)
+(bound_nonfix_op (prefix_op_symbol) @operator)
+(bound_postfix_op symbol: (_) @function.builtin)
+(bound_prefix_op symbol: (_) @function.builtin)
+(function_definition name: (identifier) @function)
+(module_definition name: (identifier) @function)
+(operator_definition name: (_) @operator)
+(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) @operator)
+(recursive_declaration (operator_declaration name: (_) @operator))
+
+; Constants and variables
+(constant_declaration (identifier) @constant.builtin)
+(constant_declaration (operator_declaration name: (_) @constant.builtin))
+(pcal_var_decl (identifier) @variable.builtin)
+(pcal_with (identifier) @parameter)
+((".") . (identifier) @attribute)
+(record_literal (identifier) @attribute)
+(set_of_records (identifier) @attribute)
+(variable_declaration (identifier) @variable.builtin)
+
+; Parameters
+(quantifier_bound (identifier) @parameter)
+(quantifier_bound (tuple_of_identifiers (identifier) @parameter))
+(lambda (identifier) @parameter)
+(module_definition (operator_declaration name: (_) @parameter))
+(module_definition parameter: (identifier) @parameter)
+(operator_definition (operator_declaration name: (_) @parameter))
+(operator_definition parameter: (identifier) @parameter)
+(pcal_macro_decl parameter: (identifier) @parameter)
+(pcal_proc_var_decl (identifier) @parameter)
+
+; Delimiters
 [
   (langle_bracket)
   (rangle_bracket)
@@ -112,44 +193,15 @@
 ] @punctuation.bracket
 [
   ","
-  "."
   ":"
+  "."
   "!"
+  ";"
+  (bullet_conj)
+  (bullet_disj)
+  (prev_func_val)
+  (placeholder)
 ] @punctuation.delimiter
-(bullet_conj) @punctuation.special
-(bullet_disj) @punctuation.special
-(prev_func_val) @punctuation.special
-(placeholder) @punctuation.special
-
-; Namespaces
-(module name: (identifier) @namespace)
-(extends (identifier_ref) @namespace)
-(instance (identifier_ref) @namespace)
-(module_definition name: (identifier) @namespace)
-
-; Constants, Variables, and Operators
-(variable_declaration (identifier) @variable)
-(constant_declaration (identifier) @constant)
-(constant_declaration (operator_declaration name: (_) @constant))
-(recursive_declaration (identifier) @function.macro)
-(recursive_declaration (operator_declaration name: (_) @function.macro))
-(operator_definition name: (_) @function.macro)
-(function_definition name: (identifier) @function)
-(bound_prefix_op symbol: (_) @operator)
-(bound_nonfix_op (prefix_op_symbol) @operator)
-(bound_infix_op symbol: (_) @operator)
-(bound_nonfix_op (infix_op_symbol) @operator)
-(bound_postfix_op symbol: (_) @operator)
-(bound_nonfix_op (postfix_op_symbol) @operator)
-
-; Parameters
-(operator_definition parameter: (identifier) @parameter)
-(operator_definition (operator_declaration name: (_) @parameter))
-(module_definition parameter: (identifier) @parameter)
-(module_definition (operator_declaration name: (_) @parameter))
-(function_definition (quantifier_bound (identifier) @parameter))
-(function_definition (quantifier_bound (tuple_of_identifiers (identifier) @parameter)))
-(lambda (identifier) @parameter)
 
 ; Proofs
 (proof_step_id "<" @punctuation.bracket)
@@ -160,3 +212,19 @@
 (proof_step_ref (level) @number)
 (proof_step_ref (name) @constant)
 (proof_step_ref ">" @punctuation.bracket)
+
+; Comments and tags
+(block_comment "(*" @comment)
+(block_comment "*)" @comment)
+(block_comment_text) @comment
+(comment) @comment
+(single_line) @comment
+(_ label: (identifier) @tag)
+(pcal_goto statement: (identifier) @tag)
+
+; Reference highlighting with the same color as declarations.
+; `constant`, `operator`, and others are custom captures defined in locals.scm
+((identifier_ref) @constant.builtin (#is? @constant.builtin constant))
+((identifier_ref) @operator (#is? @operator function))
+((identifier_ref) @parameter (#is? @parameter parameter))
+((identifier_ref) @variable.builtin (#is? @variable.builtin var))
diff --git a/queries/tlaplus/injections.scm b/queries/tlaplus/injections.scm
index 24bda1cba..4d0833e1b 100644
--- a/queries/tlaplus/injections.scm
+++ b/queries/tlaplus/injections.scm
@@ -1,4 +1,4 @@
 [
   (comment)
-  (block_comment)
+  (block_comment_text)
 ] @comment
diff --git a/queries/tlaplus/locals.scm b/queries/tlaplus/locals.scm
new file mode 100644
index 000000000..7687c5110
--- /dev/null
+++ b/queries/tlaplus/locals.scm
@@ -0,0 +1,42 @@
+; Scopes
+[
+  (bounded_quantification)
+  (function_definition) 
+  (lambda) 
+  (module) 
+  (module_definition) 
+  (pcal_algorithm)
+  (pcal_macro)
+  (pcal_procedure)
+  (pcal_with)
+  (unbounded_quantification)
+] @scope
+
+; Definitions
+(constant_declaration (identifier) @definition.constant)
+(function_definition name: (identifier) @definition.function)
+(lambda (identifier) @definition.parameter)
+(operator_definition name: (identifier) @definition.function)
+(operator_definition parameter: (identifier) @definition.parameter)
+(pcal_macro_decl parameter: (identifier) @definition.parameter)
+(pcal_proc_var_decl (identifier) @definition.parameter)
+(pcal_var_decl (identifier) @definition.var)
+(pcal_with (identifier) @definition.parameter)
+(quantifier_bound (identifier) @definition.parameter)
+(quantifier_bound (tuple_of_identifiers (identifier) @definition.parameter))
+(variable_declaration (identifier) @definition.var)
+
+; Builtin variables
+(pcal_algorithm_body
+  [
+    (_ (identifier_ref) @definition.var)
+    (_ (_ (identifier_ref) @definition.var))
+    (_ (_ (_ (identifier_ref) @definition.var))) 
+    (_ (_ (_ (_ (identifier_ref) @definition.var))))
+    (_ (_ (_ (_ (_ (identifier_ref) @definition.var)))))
+  ]
+  (#vim-match? @definition.var "^(self|pc|stack)$")
+)
+
+; References
+(identifier_ref) @reference