Updated queries and parser version

2 files changed, 119 insertions, 36 deletions

diff --git a/queries/tlaplus/highlights.scm b/queries/tlaplus/highlights.scm
index 98901ef61..a94a91f4e 100644
--- a/queries/tlaplus/highlights.scm
+++ b/queries/tlaplus/highlights.scm
@@ -1,6 +1,7 @@
 ; ; Intended for consumption by nvim-treesitter 
 ; ; Default capture names for nvim-treesitter found here:
 ; ; https://github.com/nvim-treesitter/nvim-treesitter/blob/e473630fe0872cb0ed97cd7085e724aa58bc1c84/lua/nvim-treesitter/highlight.lua#L14-L104
+; ; In this file, captures defined later take precedence over captures defined earlier
 
 ; Keywords
 [
@@ -74,6 +75,7 @@
   (temporal_exists)
   (temporal_forall)
 ] @keyword
+
 ;  Pluscal keywords
 [
   (pcal_algorithm_start)
@@ -136,41 +138,43 @@
 (escape_char) @string.escape
 (string_set) @type
 
-; Namespaces and includes
-(extends (identifier_ref) @include)
-(module name: (_) @namespace)
+; Namespaces
+(extends (identifier_ref) @namespace)
+(instance (identifier_ref) @namespace)
+(module name: (identifier) @namespace)
 (pcal_algorithm name: (identifier) @namespace)
 
-; 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)
+; 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: (identifier) @function)
-(operator_definition name: (_) @operator)
+(module_definition name: (_) @include)
+(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) @operator)
-(recursive_declaration (operator_declaration name: (_) @operator))
+(recursive_declaration (identifier) @function.macro)
+(recursive_declaration (operator_declaration name: (_) @function.macro))
 
 ; Constants and variables
-(constant_declaration (identifier) @constant.builtin)
-(constant_declaration (operator_declaration name: (_) @constant.builtin))
-(pcal_var_decl (identifier) @variable.builtin)
+(constant_declaration (identifier) @constant)
+(constant_declaration (operator_declaration name: (_) @constant))
+(pcal_var_decl (identifier) @variable)
 (pcal_with (identifier) @parameter)
 ((".") . (identifier) @attribute)
 (record_literal (identifier) @attribute)
 (set_of_records (identifier) @attribute)
-(variable_declaration (identifier) @variable.builtin)
+(variable_declaration (identifier) @variable)
 
 ; Parameters
-(quantifier_bound (identifier) @parameter)
-(quantifier_bound (tuple_of_identifiers (identifier) @parameter))
+(choose (identifier) @parameter)
+(choose (tuple_of_identifiers (identifier) @parameter))
 (lambda (identifier) @parameter)
 (module_definition (operator_declaration name: (_) @parameter))
 (module_definition parameter: (identifier) @parameter)
@@ -178,6 +182,9 @@
 (operator_definition parameter: (identifier) @parameter)
 (pcal_macro_decl parameter: (identifier) @parameter)
 (pcal_proc_var_decl (identifier) @parameter)
+(quantifier_bound (identifier) @parameter)
+(quantifier_bound (tuple_of_identifiers (identifier) @parameter))
+(unbounded_quantification (identifier) @parameter)
 
 ; Delimiters
 [
@@ -205,6 +212,10 @@
 ] @punctuation.delimiter
 
 ; Proofs
+(assume_prove (new (identifier) @parameter))
+(assume_prove (new (operator_declaration name: (_) @parameter)))
+(assumption name: (identifier) @constant)
+(pick_proof_step (identifier) @parameter)
 (proof_step_id "<" @punctuation.bracket)
 (proof_step_id (level) @label)
 (proof_step_id (name) @label)
@@ -213,6 +224,8 @@
 (proof_step_ref (level) @label)
 (proof_step_ref (name) @label)
 (proof_step_ref ">" @punctuation.bracket)
+(take_proof_step (identifier) @parameter)
+(theorem name: (identifier) @constant)
 
 ; Comments and tags
 (block_comment "(*" @comment)
@@ -226,7 +239,30 @@
 
 ; 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) @constant (#is? @constant constant))
+((identifier_ref) @function (#is? @function function))
+((identifier_ref) @function.macro (#is? @function.macro macro))
+((identifier_ref) @include (#is? @include import))
 ((identifier_ref) @parameter (#is? @parameter parameter))
-((identifier_ref) @variable.builtin (#is? @variable.builtin var))
+((identifier_ref) @variable (#is? @variable var))
+((prefix_op_symbol) @constant (#is? @constant constant))
+((prefix_op_symbol) @function.macro (#is? @function.macro macro))
+((prefix_op_symbol) @parameter (#is? @parameter parameter))
+((infix_op_symbol) @constant (#is? @constant constant))
+((infix_op_symbol) @function.macro (#is? @function.macro macro))
+((infix_op_symbol) @parameter (#is? @parameter parameter))
+((postfix_op_symbol) @constant (#is? @constant constant))
+((postfix_op_symbol) @function.macro (#is? @function.macro macro))
+((postfix_op_symbol) @parameter (#is? @parameter parameter))
+(bound_prefix_op symbol: (_) @constant (#is? @constant constant))
+(bound_prefix_op symbol: (_) @function.macro (#is? @function.macro macro))
+(bound_prefix_op symbol: (_) @parameter (#is? @parameter parameter))
+(bound_infix_op symbol: (_) @constant (#is? @constant constant))
+(bound_infix_op symbol: (_) @function.macro (#is? @function.macro macro))
+(bound_infix_op symbol: (_) @parameter (#is? @parameter parameter))
+(bound_postfix_op symbol: (_) @constant (#is? @constant constant))
+(bound_postfix_op symbol: (_) @function.macro (#is? @function.macro macro))
+(bound_postfix_op symbol: (_) @parameter (#is? @parameter parameter))
+(bound_nonfix_op symbol: (_) @constant (#is? @constant constant))
+(bound_nonfix_op symbol: (_) @function.macro (#is? @function.macro macro))
+(bound_nonfix_op symbol: (_) @parameter (#is? @parameter parameter))
diff --git a/queries/tlaplus/locals.scm b/queries/tlaplus/locals.scm
index 7687c5110..9a05eb7b7 100644
--- a/queries/tlaplus/locals.scm
+++ b/queries/tlaplus/locals.scm
@@ -1,32 +1,72 @@
-; Scopes
+; TLA+ scopes and definitions
 [
   (bounded_quantification)
+  (choose)
   (function_definition) 
+  (function_literal)
   (lambda) 
+  (let_in)
   (module) 
-  (module_definition) 
-  (pcal_algorithm)
-  (pcal_macro)
-  (pcal_procedure)
-  (pcal_with)
+  (module_definition)
+  (operator_definition)
+  (set_filter)
+  (set_map)
   (unbounded_quantification)
 ] @scope
 
-; Definitions
+(choose (identifier) @definition.parameter)
+(choose (tuple_of_identifiers (identifier) @definition.parameter))
 (constant_declaration (identifier) @definition.constant)
-(function_definition name: (identifier) @definition.function)
+(constant_declaration (operator_declaration name: (_) @definition.constant))
+(function_definition
+  name: (identifier) @definition.function
+  (#set! "definition.function.scope" "parent"))
 (lambda (identifier) @definition.parameter)
-(operator_definition name: (identifier) @definition.function)
+(module_definition
+  name: (_) @definition.import
+  (#set! "definition.import.scope" "parent"))
+(module_definition parameter: (identifier) @definition.parameter)
+(module_definition parameter: (operator_declaration name: (_) @definition.parameter))
+(operator_definition
+  name: (_) @definition.macro
+  (#set! "definition.macro.scope" "parent"))
 (operator_definition parameter: (identifier) @definition.parameter)
+(operator_definition parameter: (operator_declaration name: (_) @definition.parameter))
+(quantifier_bound (identifier) @definition.parameter)
+(quantifier_bound (tuple_of_identifiers (identifier) @definition.parameter))
+(unbounded_quantification (identifier) @definition.parameter)
+(variable_declaration (identifier) @definition.var)
+
+; Proof scopes and definitions
+[
+  (non_terminal_proof)
+  (suffices_proof_step)
+  (theorem)
+] @scope
+
+(assume_prove (new (identifier) @definition.parameter))
+(assume_prove (new (operator_declaration name: (_) @definition.parameter)))
+(assumption name: (identifier) @definition.constant)
+(pick_proof_step (identifier) @definition.parameter)
+(take_proof_step (identifier) @definition.parameter)
+(theorem
+  name: (identifier) @definition.constant
+  (#set! "definition.constant.scope" "parent"))
+
+; PlusCal scopes and definitions
+[
+  (pcal_algorithm)
+  (pcal_macro)
+  (pcal_procedure)
+  (pcal_with)
+] @scope
+
 (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
+; Built-in PlusCal variables
 (pcal_algorithm_body
   [
     (_ (identifier_ref) @definition.var)
@@ -40,3 +80,10 @@
 
 ; References
 (identifier_ref) @reference
+((prefix_op_symbol) @reference)
+(bound_prefix_op symbol: (_) @reference)
+((infix_op_symbol) @reference)
+(bound_infix_op symbol: (_) @reference)
+((postfix_op_symbol) @reference)
+(bound_postfix_op symbol: (_) @reference)
+(bound_nonfix_op symbol: (_) @reference)