1 files changed, 89 insertions, 86 deletions
diff --git a/runtime/queries/wgsl/highlights.scm b/runtime/queries/wgsl/highlights.scm
index baf9dd8f..2817ea3a 100644
--- a/runtime/queries/wgsl/highlights.scm
+++ b/runtime/queries/wgsl/highlights.scm
@@ -2,57 +2,26 @@
 (float_literal) @constant.numeric.float
 (bool_literal) @constant.builtin.boolean
 
-(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
-    (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)))
+[
+  "bitcast"
+  "discard"
+  "enable"
+  "fallthrough"
+] @keyword
 
-(type_constructor_or_function_call_expression
-    (type_declaration) @function)
+[
+  "let"
+  "override"
+  "struct"
+  "type"
+  "var"
+  (texel_format)
+] @keyword.storage.type
 
 [
-    "struct"
-    "bitcast"
-    "discard"
-    "enable"
-    "fallthrough"
-    "fn"
-    "let"
-    "private"
-    "read"
-    "read_write"
-    "storage"
-    "type"
-    "uniform"
-    "var"
-    "workgroup"
-    "write"
-    "override"
-    (texel_format)
-] @keyword
+  (access_mode)
+  (address_space)
+] @keyword.storage.modifier
 
 "fn" @keyword.function
 
@@ -62,53 +31,87 @@
 
 ["(" ")" "[" "]" "{" "}"] @punctuation.bracket
 
+(type_declaration ["<" ">"] @punctuation.bracket)
+
 [
-    "loop"
-    "for"
-    "while"
-    "break"
-    "continue"
-    "continuing"
+  "break"
+  "continue"
+  "continuing"
+] @keyword.control
+
+[
+  "loop"
+  "for"
+  "while"
 ] @keyword.control.repeat
 
 [
-    "if"
-    "else"
-    "switch"
-    "case"
-    "default"
+  "if"
+  "else"
+  "switch"
+  "case"
+  "default"
 ] @keyword.control.conditional
 
 [
-    "&"
-    "&&"
-    "/"
-    "!"
-    "="
-    "=="
-    "!="
-    ">"
-    ">="
-    ">>"
-    "<"
-    "<="
-    "<<"
-    "%"
-    "-"
-    "+"
-    "|"
-    "||"
-    "*"
-    "~"
-    "^"
-    "@"
-    "++"
-    "--"
+  "!"
+  "!="
+  "%"
+  "%="
+  "&"
+  "&&"
+  "&="
+  "*"
+  "*="
+  "+"
+  "++"
+  "+="
+  "-"
+  "--"
+  "-="
+  "->"
+  "/"
+  "/="
+  "<"
+  "<<"
+  "<="
+  "="
+  "=="
+  ">"
+  ">="
+  ">>"
+  "@"
+  "^"
+  "^="
+  "|"
+  "|="
+  "||"
+  "~"
 ] @operator
 
+(function_declaration
+  (identifier) @function)
+
+(parameter
+  (variable_identifier_declaration
+    (identifier) @variable.parameter))
+
+(struct_declaration
+  (identifier) @type)
+
+(struct_declaration
+  (struct_member
+    (variable_identifier_declaration
+      (identifier) @variable.other.member)))
+
+(type_constructor_or_function_call_expression
+  (type_declaration (identifier) @function))
+
+(type_declaration _ @type)
+
 (attribute
-    (identifier) @attribute)
+  (identifier) @attribute)
 
-(comment) @comment
+(identifier) @variable
 
-(ERROR) @error
+(comment) @comment