Improve queries

.gitignore

@@ -65,3 +65,4 @@ setup.py
 bindings/
 src/*
 !src/scanner.c
+!src/grammar.json

grammar.js

@@ -15,10 +15,10 @@ module.exports = grammar({
     name: 'wgsl',
 
     externals: $ => [
-        $._block_comment,
+        $.block_comment,
         $._disambiguate_template,
-        $._template_args_start,
-        $._template_args_end,
+        $.template_args_start,
+        $.template_args_end,
         $._less_than,
         $._less_than_equal,
         $._shift_left,
@@ -31,14 +31,13 @@ module.exports = grammar({
     ],
 
     extras: $ => [
-        $._comment,
-        $._block_comment,
+        $.line_comment,
+        $.block_comment,
         $._blankspace,
     ],
 
     inline: $ => [
         $.global_decl,
-        $._reserved,
     ],
 
     // WGSL has no parsing conflicts.
@@ -81,7 +80,7 @@ module.exports = grammar({
 
         diagnostic_rule_name: $ => choice($.diagnostic_name_token, seq($.diagnostic_name_token, '.', $.diagnostic_name_token)),
 
-        template_list: $ => seq($._template_args_start, $.template_arg_comma_list, $._template_args_end),
+        template_list: $ => seq($.template_args_start, $.template_arg_comma_list, $.template_args_end),
 
         template_arg_comma_list: $ => seq($.template_arg_expression, repeat(seq(',', $.template_arg_expression)), optional(',')),
 
@@ -183,9 +182,11 @@ module.exports = grammar({
 
         additive_operator: $ => choice('+', '-'),
 
-        shift_expression: $ => choice($.additive_expression, seq($.unary_expression, $._shift_left, $.unary_expression), seq($.unary_expression, $._shift_right, $.unary_expression)),
+        shift_expression: $ => choice($.additive_expression, seq($.unary_expression, $.shift_operator, $.unary_expression)),
 
-        relational_expression: $ => choice($.shift_expression, seq($.shift_expression, $._less_than, $.shift_expression), seq($.shift_expression, $._greater_than, $.shift_expression), seq($.shift_expression, $._less_than_equal, $.shift_expression), seq($.shift_expression, $._greater_than_equal, $.shift_expression), seq($.shift_expression, '==', $.shift_expression), seq($.shift_expression, '!=', $.shift_expression)),
+        shift_operator: $ => choice($._shift_left, $._shift_right),
+
+        relational_expression: $ => choice($.shift_expression, seq($.shift_expression, $.comparison_operator, $.shift_expression)),
 
         short_circuit_and_expression: $ => choice($.relational_expression, seq($.short_circuit_and_expression, '&&', $.relational_expression)),
 
@@ -207,6 +208,8 @@ module.exports = grammar({
 
         compound_assignment_operator: $ => choice('+=', '-=', '*=', '/=', '%=', '&=', '|=', '^=', $._shift_right_assign, $._shift_left_assign),
 
+        comparison_operator: $ => choice('==', '!=', $._less_than, $._less_than_equal, $._greater_than, $._greater_than_equal),
+
         increment_statement: $ => seq($.lhs_expression, '++'),
 
         decrement_statement: $ => seq($.lhs_expression, '--'),
@@ -269,7 +272,7 @@ module.exports = grammar({
 
         function_decl: $ => seq(repeat($.attribute), $.function_header, $.compound_statement),
 
-        function_header: $ => seq('fn', $.ident, '(', optional($.param_list), ')', optional(seq('->', repeat($.attribute), $.template_elaborated_ident))),
+        function_header: $ => seq('fn', $.ident, '(', optional($.param_list), ')', optional(seq('->', repeat($.attribute), $.type_specifier))),
 
         param_list: $ => seq($.param, repeat(seq(',', $.param)), optional(',')),
 
@@ -295,8 +298,10 @@ module.exports = grammar({
 
         ident: $ => seq($.ident_pattern_token, $._disambiguate_template),
 
-        _comment: $ => /\/\/.*/,
+        line_comment: $ => /\/\/.*/,
+
+        _blankspace: $ => /[\u0020\u0009\u000a\u000b\u000c\u000d\u0085\u200e\u200f\u2028\u2029]/u,
 
-        _blankspace: $ => /[\u0020\u0009\u000a\u000b\u000c\u000d\u0085\u200e\u200f\u2028\u2029]/u
+        reserved: $ => choice('NULL', 'Self', 'abstract', 'active', 'alignas', 'alignof', 'as', 'asm', 'asm_fragment', 'async', 'attribute', 'auto', 'await', 'become', 'cast', 'catch', 'class', 'co_await', 'co_return', 'co_yield', 'coherent', 'column_major', 'common', 'compile', 'compile_fragment', 'concept', 'const_cast', 'consteval', 'constexpr', 'constinit', 'crate', 'debugger', 'decltype', 'delete', 'demote', 'demote_to_helper', 'do', 'dynamic_cast', 'enum', 'explicit', 'export', 'extends', 'extern', 'external', 'fallthrough', 'filter', 'final', 'finally', 'friend', 'from', 'fxgroup', 'get', 'goto', 'groupshared', 'highp', 'impl', 'implements', 'import', 'inline', 'instanceof', 'interface', 'layout', 'lowp', 'macro', 'macro_rules', 'match', 'mediump', 'meta', 'mod', 'module', 'move', 'mut', 'mutable', 'namespace', 'new', 'nil', 'noexcept', 'noinline', 'nointerpolation', 'non_coherent', 'noncoherent', 'noperspective', 'null', 'nullptr', 'of', 'operator', 'package', 'packoffset', 'partition', 'pass', 'patch', 'pixelfragment', 'precise', 'precision', 'premerge', 'priv', 'protected', 'pub', 'public', 'readonly', 'ref', 'regardless', 'register', 'reinterpret_cast', 'require', 'resource', 'restrict', 'self', 'set', 'shared', 'sizeof', 'smooth', 'snorm', 'static', 'static_assert', 'static_cast', 'std', 'subroutine', 'super', 'target', 'template', 'this', 'thread_local', 'throw', 'trait', 'try', 'type', 'typedef', 'typeid', 'typename', 'typeof', 'union', 'unless', 'unorm', 'unsafe', 'unsized', 'use', 'using', 'varying', 'virtual', 'volatile', 'wgsl', 'where', 'with', 'writeonly', 'yield'),
     }
 })

queries/folds.scm

@@ -0,0 +1,12 @@
+; Fold compound statements
+[
+  (compound_statement)
+  (continuing_compound_statement)
+  (switch_body)
+] @fold
+
+; Fold struct declarations
+(struct_body_decl) @fold
+
+; Fold function declarations
+(function_decl) @fold

queries/highlights.scm

@@ -1,78 +1,142 @@
-; Types
-(type_specifier) @type
-(template_elaborated_ident) @type
+((ident) @variable
+  (#set! priority 99))
 
-; Functions
-(function_decl
-  (function_header
-    (ident) @function))
-(call_expression
-  (call_phrase
-    (template_elaborated_ident) @function))
+(bool_literal) @boolean
 
-; Variables and parameters
-(param
-  (ident) @variable.parameter)
-(variable_decl
-  (optionally_typed_ident (ident) @variable))
-(assert_statement) @variable
+(int_literal) @number
+
+(float_literal) @number.float
+
+(line_comment) @comment
+
+(block_comment) @comment
+
+(type_specifier) @type
 
-; Struct and struct members
 (struct_decl
   (ident) @type)
-(struct_member
-  (member_ident) @property)
 
-; Attributes
+(type_alias_decl
+  (ident) @type)
+
+(template_arg_expression) @type
+
+(enable_extension_name) @module
+
+(language_extension_name) @module
+
+(severity_control_name) @constant
+
+(diagnostic_rule_name) @module
+
+(member_ident) @variable.member
+
+(component_or_swizzle_specifier) @variable.member
+
+(function_header
+  (ident) @function)
+
+(param
+  (ident)) @variable.parameter
+
+(call_expression
+  (call_phrase
+    (template_elaborated_ident
+      (ident) @function.call)))
+
+(func_call_statement
+  (call_phrase
+    (template_elaborated_ident
+      (ident) @function.call)))
+
 (attribute) @attribute
 
-; Literals
-(bool_literal) @constant.builtin
-(int_literal) @number
-(float_literal) @float
-
-; Keywords and storage classes
-(global_variable_decl
-  (variable_decl
-    (optionally_typed_ident
-      (ident) @variable)))
-
-; Control flow
-"if" @keyword.control
-"else" @keyword.control
-"loop" @keyword.control
-"for" @keyword.control
-"while" @keyword.control
-"switch" @keyword.control
-"case" @keyword.control
-"default" @keyword.control
-"break" @keyword.control
-"return" @keyword.control
-"discard" @keyword.control
-
-; Declarations
-(global_variable_decl) @keyword
-(global_value_decl) @keyword
-(variable_decl) @keyword
-
-; Punctuation
-"(" @punctuation.bracket
-")" @punctuation.bracket
-"[" @punctuation.bracket
-"]" @punctuation.bracket
-"{" @punctuation.bracket
-"}" @punctuation.bracket
-"," @punctuation.delimiter
-"." @punctuation.delimiter
-":" @punctuation.delimiter
-";" @punctuation.delimiter
-
-; Builtin functions and types
-(builtin_value_name) @function.builtin
-
-; Additional keywords
-"var" @keyword
-"let" @keyword
-"const" @keyword
-"fn" @keyword
-"struct" @keyword
+[
+  "const"
+  "const_assert"
+  "let"
+  "override"
+  "var"
+] @keyword
+
+[
+  "alias"
+  "struct"
+] @keyword.type
+
+[
+  "discard"
+  "return"
+] @keyword.return
+
+"fn" @keyword.function
+
+[
+  "diagnostic"
+  "enable"
+  "requires"
+] @keyword.directive
+
+[
+  "loop"
+  "for"
+  "while"
+  "continuing"
+  (break_statement)
+  (continue_statement)
+] @keyword.repeat
+
+[
+  "if"
+  "else"
+  "switch"
+  "case"
+  "default"
+] @keyword.conditional
+
+[
+  "+"
+  "-"
+  "*"
+  "/"
+  "%"
+  "="
+  "&&"
+  "||"
+  "&"
+  "|"
+  "^"
+  "~"
+  "++"
+  "--"
+  "!"
+  (compound_assignment_operator)
+  (comparison_operator)
+  (shift_operator)
+] @operator
+
+[
+  "."
+  ","
+  ";"
+  ":"
+  "->"
+] @punctuation.delimiter
+
+[
+  "("
+  ")"
+  "["
+  "]"
+  "{"
+  "}"
+] @punctuation.bracket
+
+[
+  (template_args_start)
+  (template_args_end)
+] @punctuation.special
+
+((ident) @keyword.modifier
+  (#match? @keyword.modifier
+    "^(function|private|workgroup|uniform|storage|read|write|read_write)$"))

queries/indents.scm

@@ -0,0 +1,35 @@
+; Indent blocks
+(compound_statement) @indent.begin
+
+(compound_statement
+  "}" @indent.end)
+
+(continuing_compound_statement) @indent.begin
+
+(continuing_compound_statement
+  "}" @indent.end)
+
+(switch_body) @indent.begin
+
+(switch_body
+  "}" @indent.end)
+
+(struct_body_decl) @indent.begin
+
+(struct_body_decl
+  "}" @indent.end)
+
+(function_header) @indent.begin
+
+(function_header
+  ")" @indent.end)
+
+(argument_expression_list) @indent.begin
+
+(argument_expression_list
+  ")" @indent.end)
+
+[
+  "}"
+  ")"
+] @indent.branch