1 files changed, 51 insertions, 13 deletions
diff --git a/runtime/queries/robot/highlights.scm b/runtime/queries/robot/highlights.scm
index 60f416b7..06e95798 100644
--- a/runtime/queries/robot/highlights.scm
+++ b/runtime/queries/robot/highlights.scm
@@ -1,21 +1,59 @@
-(comment) @comment
+[
+  (comment)
+  (extra_text)
+] @comment
+
+[
+  (section_header)
+  (setting_statement)
+  (keyword_setting)
+  (test_case_setting)
+] @keyword
+
+(variable_definition (variable_name) @variable)
+(keyword_definition (name) @function)
+(test_case_definition (name) @function)
+
+(keyword_invocation (keyword) @function)
 (ellipses) @punctuation.delimiter
 
-(section_header) @keyword
-(extra_text) @comment
+(text_chunk) @string
+(inline_python_expression) @string.special
+[
+  (scalar_variable)
+  (list_variable)
+  (dictionary_variable)
+] @variable
 
-(setting_statement) @keyword
+; Control structures
 
-(variable_definition (variable_name) @variable)
+"RETURN" @keyword.control.return
 
-(keyword_definition (name) @function)
-(keyword_definition (body (keyword_setting) @keyword))
+[
+  "FOR"
+  "IN"
+  "IN RANGE"
+  "IN ENUMERATE"
+  "IN ZIP"
+  (break_statement)
+  (continue_statement)
+] @keyword.control.repeat
+(for_statement "END" @keyword.control.repeat)
 
-(test_case_definition (name) @property)
+"WHILE" @keyword.control.repeat
+(while_statement "END" @keyword.control.repeat)
 
-(keyword_invocation (keyword) @function)
+[
+  "IF"
+  "ELSE IF"
+] @keyword.control.conditional
+(if_statement "END" @keyword.control.conditional)
+(if_statement (else_statement "ELSE" @keyword.control.conditional))
 
-(argument (text_chunk) @string)
-(argument (scalar_variable) @string.special)
-(argument (list_variable) @string.special)
-(argument (dictionary_variable) @string.special)
+[
+  "TRY"
+  "EXCEPT"
+  "FINALLY"
+] @keyword.control.exception
+(try_statement "END" @keyword.control.exception)
+(try_statement (else_statement "ELSE" @keyword.control.exception))