1 files changed, 59 insertions, 11 deletions
diff --git a/queries/rego/highlights.scm b/queries/rego/highlights.scm
index 70bf849a..42d276cd 100644
--- a/queries/rego/highlights.scm
+++ b/queries/rego/highlights.scm
@@ -1,16 +1,64 @@
 ; highlights.scm
-"import" @include
-"package" @include
+[
+  (import) 
+  (package)
+] @include
 
+[
+  (with)
+  (as)
+  (every)
+  (some)
+  (in)
+  (not)
+  (if)
+  (contains)
+  (else)
+  (default)
+  "null"
+] @keyword
+
+[
+  "true"
+  "false"
+] @boolean
+
+[
+  (assignment_operator)
+  (bool_operator)
+  (arith_operator)
+  (bin_operator)
+] @operator
+
+[
+  (string)
+  (raw_string)
+] @string
+
+(term (ref (var))) @variable
 
-(reserved_keywords) @keyword
 (comment) @comment
-(rego_block rego_rule_name: (identifier) @function)
-(builtin_function function_name: (function_name) @function.builtin)
-(opening_parameter) @punctuation.bracket
-(closing_parameter) @punctuation.bracket
-(string_definition) @string
+
 (number) @number
-(operator) @operator
-(true) @boolean
-(false) @boolean
+
+(expr_call func_name: (fn_name (var) @function .))
+
+(expr_call func_arguments: (fn_args (expr) @parameter))
+
+(rule_args (term) @parameter)
+
+[
+  (open_paren)
+  (close_paren)
+  (open_bracket)
+  (close_bracket)
+  (open_curly)
+  (close_curly)
+] @punctuation.bracket
+
+(rule (rule_head (var) @method))
+
+(rule 
+  (rule_head (term (ref (var) @namespace)))
+  (rule_body (query (literal (expr (expr_infix (expr (term (ref (var)) @_output)))))) (#eq? @_output @namespace))
+)