WGSL syntax highlighting fix (#3996)

1 files changed, 37 insertions, 25 deletions

diff --git a/runtime/queries/wgsl/highlights.scm b/runtime/queries/wgsl/highlights.scm
index 7fbc87d8..baf9dd8f 100644
--- a/runtime/queries/wgsl/highlights.scm
+++ b/runtime/queries/wgsl/highlights.scm
@@ -1,23 +1,41 @@
-(const_literal) @constant.numeric
+(int_literal) @constant.numeric.integer
+(float_literal) @constant.numeric.float
+(bool_literal) @constant.builtin.boolean
 
-(type_declaration) @type
+(global_constant_declaration) @variable
+(global_variable_declaration) @variable
+(compound_statement) @variable
+(const_expression) @function
+
+(variable_identifier_declaration
+    (identifier) @variable
+    (type_declaration) @type)
 
 (function_declaration
-    (identifier) @function)
+    (identifier) @function
+    (function_return_type_declaration
+        (type_declaration) @type))
+
+(parameter
+    (variable_identifier_declaration
+        (identifier) @variable.parameter
+        (type_declaration) @type))
 
 (struct_declaration
     (identifier) @type)
+        
+(struct_declaration
+    (struct_member
+        (variable_identifier_declaration
+            (identifier) @variable.other.member
+            (type_declaration) @type)))
 
 (type_constructor_or_function_call_expression
     (type_declaration) @function)
 
-(parameter
-    (variable_identifier_declaration (identifier) @variable.parameter))
-
 [
     "struct"
     "bitcast"
-    ; "block"
     "discard"
     "enable"
     "fallthrough"
@@ -26,36 +44,28 @@
     "private"
     "read"
     "read_write"
-    "return"
     "storage"
     "type"
     "uniform"
     "var"
     "workgroup"
     "write"
+    "override"
     (texel_format)
-] @keyword ; TODO reserved keywords
+] @keyword
 
-[
-    (true)
-    (false)
-] @constant.builtin.boolean
+"fn" @keyword.function
 
-[ "," "." ":" ";" ] @punctuation.delimiter
+"return" @keyword.control.return
 
-;; brackets
-[
-    "("
-    ")"
-    "["
-    "]"
-    "{"
-    "}"
-] @punctuation.bracket
+["," "." ":" ";"] @punctuation.delimiter
+
+["(" ")" "[" "]" "{" "}"] @punctuation.bracket
 
 [
     "loop"
     "for"
+    "while"
     "break"
     "continue"
     "continuing"
@@ -64,7 +74,6 @@
 [
     "if"
     "else"
-    "elseif"
     "switch"
     "case"
     "default"
@@ -92,10 +101,13 @@
     "*"
     "~"
     "^"
+    "@"
+    "++"
+    "--"
 ] @operator
 
 (attribute
-    (identifier) @variable.other.member)
+    (identifier) @attribute)
 
 (comment) @comment