diff --git a/code/custom/4coder_tree_sitter.cpp b/code/custom/4coder_tree_sitter.cpp
index 76f1bfce..77e6c200 100644
--- a/code/custom/4coder_tree_sitter.cpp
+++ b/code/custom/4coder_tree_sitter.cpp
@@ -395,7 +395,7 @@ async_task_cancel_nowait(Application_Links* app, Async_System* async_system, Asy
 // Token Highlighting
 ////////////////////////////////////////////////////////////////////
 
-function void
+function bool
 tree_sitter_highlight_node(
   Application_Links* app,
   TSQuery* query,
@@ -403,6 +403,7 @@ tree_sitter_highlight_node(
   TSQueryCursor* query_cursor,
   Text_Layout_ID text_layout_id
 ){
+  bool highlighted_any_code = false;
   ts_query_cursor_exec(query_cursor, query, top_node);
 
   TSQueryMatch query_match;
@@ -416,13 +417,33 @@ tree_sitter_highlight_node(
     const char* tmp = ts_query_capture_name_for_id(query, capture.index, &length);
     String_Const_u8 capture_name = SCu8((char*)tmp, length);
 
+    // TODO(PS): let each language provide a custom override
+    String_Const_u8 color_name = SCu8("defcolor_text_default");
+    if (string_match(capture_name, SCu8("function"))) {
+      color_name = SCu8("defcolor_function");
+    } else if (string_match(capture_name, SCu8("type"))) {
+      color_name = SCu8("defcolor_type");
+    } else if (string_match(capture_name, SCu8("constant")) || string_match(capture_name, SCu8("number"))) {
+      color_name = SCu8("defcolor_int_constant");
+    } else if (string_match(capture_name, SCu8("keyword"))) {
+      color_name = SCu8("defcolor_keyword");
+    } else if (string_match(capture_name, SCu8("string"))) {
+      color_name = SCu8("defcolor_str_constant");
+    } else if (string_match(capture_name, SCu8("comment"))) {
+      color_name = SCu8("defcolor_comment");
+    }
+
     Range_i64 highlight_range = tree_sitter_node_to_range(node);
-    Managed_ID color_id = managed_id_get(app, SCu8("colors"), capture_name);
+    Managed_ID color_id = managed_id_get(app, SCu8("colors"), color_name);
     if (color_id != 0)
     {
       paint_text_color_fcolor(app, text_layout_id, highlight_range, fcolor_id(color_id));
     }
+
+    highlighted_any_code = true;
   }
+
+  return highlighted_any_code;
 }
 
 function void
@@ -461,7 +482,11 @@ draw_tree_sitter_node_colors(Application_Links* app, Text_Layout_ID text_layout_
     }
     else
     {
-      tree_sitter_highlight_node(app, query, visible_container, query_cursor, text_layout_id);
+      bool should_ascend = false;
+      do {
+        should_ascend = !tree_sitter_highlight_node(app, query, visible_container, query_cursor, text_layout_id);
+        if (should_ascend) visible_container = ts_node_parent(visible_container);
+      } while (should_ascend && !ts_node_is_null(visible_container));
     }
     ts_query_cursor_delete(query_cursor);
   }
@@ -483,6 +508,10 @@ convert_to_single_line_in_place(String_Const_u8 str)
     {
       while (src < str.size && character_is_whitespace(str.str[src])) { src += 1; }
       if (src >= str.size) break;
+      else {
+        result.str[dst] = ' ';
+        dst += 1;
+      }
     }
     result.str[dst] = str.str[src];
     dst += 1;
@@ -555,15 +584,10 @@ tree_sitter_list_all_query_results(
     {
       Tree_Sitter_Language_Definition* lang = tree_sitter_language_for_buffer(app, buffer);
       if (!lang) continue;
-      TSQuery* ts_query = lang->queries.ptr[query_kind];
+      TSQuery* ts_query = lang->queries.ptr[Tree_Sitter_Language_Query_Tags];
+      if (!ts_query) ts_query = lang->queries.ptr[Tree_Sitter_Language_Query_Highlights];
       Tree_Sitter_Query_Cursor query = tree_sitter_query_init(app, buffer, ts_query);
 
-      i64 last_query_match_id = -1;
-      i64 last_query_match_printed = -1;
-      i64 last_query_line_number = 0;
-      Range_i64 last_query_range = {};
-      String_Const_u8 last_query_prefix = {};
-
       TSQueryMatch query_match;
       u32 capture_index;
       bool reached_end = false;
@@ -571,43 +595,44 @@ tree_sitter_list_all_query_results(
       {
         TSQueryCapture capture = query_match.captures[capture_index];
 
-        if (last_query_match_id != query_match.id)
-        {
-          if (last_query_match_id >= 0)
-          {
-            print_position(
-              app, &out, buffer, buffer_name, last_query_range, last_query_prefix, scratch, true
-            );
-            last_query_match_printed = last_query_match_id;
-          }
-          last_query_range.start = (i64)ts_node_start_byte(capture.node);
-          last_query_range.end = last_query_range.start;
-          last_query_prefix = {};
-        }
-
-        last_query_match_id = query_match.id;
-        last_query_range.end = Max((i64)ts_node_end_byte(capture.node), last_query_range.end);
+        Range_i64 match_range;
+        match_range.start = (i64)ts_node_start_byte(capture.node);
+        match_range.end   = (i64)ts_node_end_byte(capture.node);
 
         String_Const_u8 name;
         u32 name_length;
         name.str = (u8*)ts_query_capture_name_for_id(ts_query, capture.index, &name_length);
         name.size = (u64)name_length;
 
-        String_Const_u8 prefix_identifier = SCu8("prefix");
-        u64 prefix_loc = string_find_first(name, prefix_identifier);
-        if (prefix_loc < name.size)
+        bool matches = false;
+        switch (query_kind)
         {
-          last_query_prefix = name;
-          last_query_prefix.str += prefix_loc + prefix_identifier.size;
-          last_query_prefix.size -= prefix_loc + prefix_identifier.size;
-        }
-      }
+          case Tree_Sitter_Language_Query_Functions:
+          {
+            matches  = string_match(name, SCu8("function"));
+            matches |= string_match(name, SCu8("definition.function"));
 
-      if (last_query_match_printed != last_query_match_id)
-      {
-        print_position(
-          app, &out, buffer, buffer_name, last_query_range, last_query_prefix, scratch, true
-        );
+            // TODO(PS): ideally, these should be provided by the language
+            matches |= string_match(name, SCu8("definition.method"));
+          } break;
+
+          case Tree_Sitter_Language_Query_Types:
+          {
+            matches  = string_match(name, SCu8("type"));
+            matches |= string_match(name, SCu8("definition.type"));
+
+            // TODO(PS): ideally, these should be provided by the language
+            matches |= string_match(name, SCu8("definition.class"));
+            matches |= string_match(name, SCu8("definition.interface"));
+          } break;
+        }
+
+        if (matches)
+        {
+          print_position(
+            app, &out, buffer, buffer_name, match_range, {}, scratch, true
+          );
+        }
       }
       tree_sitter_query_end(&query);
 
diff --git a/code/custom/4coder_tree_sitter.h b/code/custom/4coder_tree_sitter.h
index 3a4bec5e..d2be4721 100644
--- a/code/custom/4coder_tree_sitter.h
+++ b/code/custom/4coder_tree_sitter.h
@@ -8,6 +8,8 @@
 enum Tree_Sitter_Language_Query_Kind
 {
   Tree_Sitter_Language_Query_Highlights,
+  Tree_Sitter_Language_Query_Tags,
+
   Tree_Sitter_Language_Query_Nests,
 
   Tree_Sitter_Language_Query_Functions,
diff --git a/code/custom/languages/tree_sitter_bash.h b/code/custom/languages/tree_sitter_bash.h
index f67eb505..22abd254 100644
--- a/code/custom/languages/tree_sitter_bash.h
+++ b/code/custom/languages/tree_sitter_bash.h
@@ -3,16 +3,65 @@
 #ifndef TREE_SITTER_BASH_H
 #define TREE_SITTER_BASH_H
 
+// Source: https://github.com/tree-sitter/tree-sitter-bash/blob/master/queries/highlights.scm
 String_Const_u8 TS_BASH_HIGHLIGHT_QUERY = string_u8_litexpr(R"DONE(
+[
+  (string)
+  (raw_string)
+  (heredoc_body)
+  (heredoc_start)
+] @string
 
-(command name: (command_name) @defcolor_function)
-(variable_assignment name: (variable_name) @defcolor_macro)
+(command_name) @function.call
 
-; Basic Types
-(string) @defcolor_str_constant
-(number) @defcolor_int_constant
+(variable_name) @property
+
+[
+  "case"
+  "do"
+  "done"
+  "elif"
+  "else"
+  "esac"
+  "export"
+  "fi"
+  "for"
+  "function"
+  "if"
+  "in"
+  "select"
+  "then"
+  "unset"
+  "until"
+  "while"
+] @keyword
+
+(comment) @comment
+
+(function_definition name: (word) @function)
+
+(file_descriptor) @number
+
+[
+  (command_substitution)
+  (process_substitution)
+  (expansion)
+] @embedded
+
+[
+  "$"
+  "&&"
+  ">"
+  ">>"
+  "<"
+  "|"
+] @operator
+
+(
+  (command (_) @constant)
+  (#match? @constant "^-")
+)
 
-(comment) @defcolor_comment
 )DONE");
 
 extern "C" {
diff --git a/code/custom/languages/tree_sitter_cpp.h b/code/custom/languages/tree_sitter_cpp.h
index 4e09a44e..49184546 100644
--- a/code/custom/languages/tree_sitter_cpp.h
+++ b/code/custom/languages/tree_sitter_cpp.h
@@ -5,55 +5,186 @@
 
 String_Const_u8 TS_CPP_NEST_QUERY = string_u8_litexpr("(_ \"{\" @scope_begin \"}\" @scope_end )\n");
 
-String_Const_u8 TS_CPP_FUNCTION_QUERY = string_u8_litexpr(R"DONE(
-  (function_declarator) @function_identifier
+// Source: https://github.com/tree-sitter/tree-sitter-cpp/blob/master/queries/tags.scm
+String_Const_u8 TS_CPP_TAGS_QUERY_SCM = string_u8_litexpr(R"DONE(
+(struct_specifier name: (type_identifier) @name body:(_)) @definition.class
+
+(declaration type: (union_specifier name: (type_identifier) @name)) @definition.class
+
+(function_declarator declarator: (identifier) @name) @definition.function
+
+(function_declarator declarator: (field_identifier) @name) @definition.function
+
+(function_declarator declarator: (qualified_identifier scope: (namespace_identifier) @local.scope name: (identifier) @name)) @definition.method
+
+(type_definition declarator: (type_identifier) @name) @definition.type
+
+(enum_specifier name: (type_identifier) @name) @definition.type
+
+(class_specifier name: (type_identifier) @name) @definition.class
 )DONE");
 
-String_Const_u8 TS_CPP_TYPE_QUERY = string_u8_litexpr(R"DONE(
-  (struct_specifier
-    name: (type_identifier) @prefixStruct
-  )
-  (enum_specifier
-    name: (type_identifier) @prefixEnum
-  )
-  (class_specifier
-    name: (type_identifier) @prefixClass
-  )
-)DONE");
+// Source: https://github.com/tree-sitter/tree-sitter-cpp/blob/master/queries/highlights.scm
+String_Const_u8 TS_CPP_HIGHLIGHT_QUERY_SCM = string_u8_litexpr(R"DONE(
+(identifier) @variable
 
-String_Const_u8 TS_CPP_HIGHLIGHT_QUERY = string_u8_litexpr(R"DONE(
-(call_expression function: [
- (identifier) @defcolor_function
- (field_expression field: (field_identifier) @defcolor_function)])
+((identifier) @constant
+ (#match? @constant "^[A-Z][A-Z\\d_]*$"))
 
-(function_declarator
- declarator: [(identifier) (field_identifier)] @defcolor_function)
+"break" @keyword
+"case" @keyword
+"const" @keyword
+"continue" @keyword
+"default" @keyword
+"do" @keyword
+"else" @keyword
+"enum" @keyword
+"extern" @keyword
+"for" @keyword
+"if" @keyword
+"inline" @keyword
+"return" @keyword
+"sizeof" @keyword
+"static" @keyword
+"struct" @keyword
+"switch" @keyword
+"typedef" @keyword
+"union" @keyword
+"volatile" @keyword
+"while" @keyword
 
-(preproc_def
- name: (identifier) @defcolor_macro)
+"#define" @keyword
+"#elif" @keyword
+"#else" @keyword
+"#endif" @keyword
+"#if" @keyword
+"#ifdef" @keyword
+"#ifndef" @keyword
+"#include" @keyword
+(preproc_directive) @keyword
 
-(preproc_function_def
- name: (identifier) @defcolor_macro)
+"--" @operator
+"-" @operator
+"-=" @operator
+"->" @operator
+"=" @operator
+"!=" @operator
+"*" @operator
+"&" @operator
+"&&" @operator
+"+" @operator
+"++" @operator
+"+=" @operator
+"<" @operator
+"==" @operator
+">" @operator
+"||" @operator
 
-(type_identifier) @defcolor_type
+"." @delimiter
+";" @delimiter
+
+(string_literal) @string
+(system_lib_string) @string
+
+(null) @constant
+(number_literal) @number
+(char_literal) @number
+
+(field_identifier) @property
+(statement_identifier) @label
+(type_identifier) @type
+(primitive_type) @type
+(sized_type_specifier) @type
 
 (call_expression
- function: (parenthesized_expression
-  (identifier) @defcolor_type))
+  function: (identifier) @function)
+(call_expression
+  function: (field_expression
+    field: (field_identifier) @function))
+(function_declarator
+  declarator: (identifier) @function)
+(preproc_function_def
+  name: (identifier) @function.special)
 
-[(primitive_type) (type_qualifier) (storage_class_specifier)
- (break_statement) (continue_statement) "union" "return" "do"
- "while" "for" "if" "class" "struct" "enum" "sizeof"
- "else" "switch" "case"] @defcolor_keyword
+(comment) @comment
 
-[(number_literal) (string_literal) (raw_string_literal)] @defcolor_str_constant
+; Functions
 
-[(preproc_directive) "#define" "#if" "#elif" "#else" "#endif"
- "#include"] @defcolor_preproc
+(call_expression
+  function: (qualified_identifier
+    name: (identifier) @function))
 
-["{" "}" ";" ":" ","] @defcolor_text_default
+(template_function
+  name: (identifier) @function)
 
-(comment) @defcolor_comment
+(template_method
+  name: (field_identifier) @function)
+
+(template_function
+  name: (identifier) @function)
+
+(function_declarator
+  declarator: (qualified_identifier
+    name: (identifier) @function))
+
+(function_declarator
+  declarator: (field_identifier) @function)
+
+; Types
+
+((namespace_identifier) @type
+ (#match? @type "^[A-Z]"))
+
+(auto) @type
+
+; Constants
+
+(this) @variable.builtin
+(null "nullptr" @constant)
+
+; Modules
+(module_name
+  (identifier) @module)
+
+; Keywords
+
+[
+ "catch"
+ "class"
+ "co_await"
+ "co_return"
+ "co_yield"
+ "constexpr"
+ "constinit"
+ "consteval"
+ "delete"
+ "explicit"
+ "final"
+ "friend"
+ "mutable"
+ "namespace"
+ "noexcept"
+ "new"
+ "override"
+ "private"
+ "protected"
+ "public"
+ "template"
+ "throw"
+ "try"
+ "typename"
+ "using"
+ "concept"
+ "requires"
+ "virtual"
+ "import"
+ "export"
+ "module"
+] @keyword
+
+; Strings
+
+(raw_string_literal) @string
 )DONE");
 
 extern "C" {
@@ -64,13 +195,10 @@ void
 tree_sitter_register_cpp(Application_Links* app)
 {
   TSLanguage* language = tree_sitter_cpp();
-  String_Const_u8 highlight_query_str = TS_CPP_HIGHLIGHT_QUERY;
-
   Tree_Sitter_Language_Queries queries = {};
-  queries.ptr[Tree_Sitter_Language_Query_Highlights] = tree_sitter_query_new(app, language, TS_CPP_HIGHLIGHT_QUERY);
+  queries.ptr[Tree_Sitter_Language_Query_Highlights] = tree_sitter_query_new(app, language, TS_CPP_HIGHLIGHT_QUERY_SCM);
+  queries.ptr[Tree_Sitter_Language_Query_Tags]       = tree_sitter_query_new(app, language, TS_CPP_TAGS_QUERY_SCM);
   queries.ptr[Tree_Sitter_Language_Query_Nests]      = tree_sitter_query_new(app, language, TS_CPP_NEST_QUERY);
-  queries.ptr[Tree_Sitter_Language_Query_Functions]  = tree_sitter_query_new(app, language, TS_CPP_FUNCTION_QUERY);
-  queries.ptr[Tree_Sitter_Language_Query_Types]      = tree_sitter_query_new(app, language, TS_CPP_TYPE_QUERY);
   tree_sitter_register_language(SCu8("c"),   language, queries);
   tree_sitter_register_language(SCu8("cpp"), language, queries);
   tree_sitter_register_language(SCu8("h"),   language, queries);
diff --git a/code/custom/languages/tree_sitter_jai.h b/code/custom/languages/tree_sitter_jai.h
index 97fc007b..834b17a6 100644
--- a/code/custom/languages/tree_sitter_jai.h
+++ b/code/custom/languages/tree_sitter_jai.h
@@ -3,143 +3,98 @@
 #ifndef TREE_SITTER_JAI_H
 #define TREE_SITTER_JAI_H
 
-String_Const_u8 TS_JAI_FUNCTION_QUERY = string_u8_litexpr(R"DONE(
+String_Const_u8 TS_JAI_TAGS_QUERY = string_u8_litexpr(R"DONE(
 
 (procedure_declaration
-  name: (identifier) @print1
-  (procedure
-    (named_parameters) @print2
-    (procedure_returns) @print
-  )
+  name: (identifier) @definition.function
 )
 
-)DONE");
-
-String_Const_u8 TS_JAI_TYPE_QUERY = string_u8_litexpr(R"DONE(
-
 (struct_declaration
-  name: (identifier) @prefixStruct
+  name: (identifier) @definition.type
 )
 
 (enum_declaration
-  name: (identifier) @prefixEnum
+  name: (identifier) @definition.type
 )
 
 )DONE");
 
-String_Const_u8 TS_JAI_HIGHLIGHT_QUERY = string_u8_litexpr(R"DONE(
-
-; keywords
-[
-  "if"
-  "else"
-  "break"
-  "continue"
-  "return"
-  "struct"
-  "enum"
-  "for"
-  "defer"
-  "cast"
-  "xx"
-  "ifx"
-  "null"
-] @defcolor_keyword
-
-; # preceeded
-[
-  (compiler_directive)
-  (import)
-  (char_string)
-] @defcolor_macro
-
-(import (identifier) @defcolor_type)
-
-; Identifiers
-(struct_declaration
-  name: (identifier) @defcolor_type
-)
-
-(struct_literal
-  type: (identifier) @defcolor_type
-)
-
-(enum_declaration
-  name: (identifier) @defcolor_type
-)
-
-(enum_declaration "{" (identifier) @defcolor_type)
-
-(variable_declaration
-  type: (types) @defcolor_type
-)
-
-(procedure_declaration
-  name: (identifier) @defcolor_function
-)
-
-(call_expression
-  function: (identifier) @defcolor_function
-)
-
-(procedure
-  result: (procedure_returns) @defcolor_type
-)
-
-; Constants & Literals
-[
-  (string)
-  (string_directive)
-] @defcolor_str_constant
-(escape_sequence) @defcolor_special_character
-(integer) @defcolor_int_constant
-(float) @defcolor_float_constant
-(boolean) @defcolor_bool_constant
-
-(array_literal
-  type: (identifier) @defcolor_type
-)
-
-; Comments
-(note) @defcolor_comment
-(block_comment) @defcolor_comment
-(block_comment_text) @defcolor_comment
-
-)DONE");
-
 // NOTE(PS): source: https://github.com/St0wy/tree-sitter-jai/blob/main/queries/highlights.scm
-String_Const_u8 TS_JAI_HIGHLIGHT_QUERY_ = string_u8_litexpr(R"DONE(
+String_Const_u8 TS_JAI_HIGHLIGHT_QUERY = string_u8_litexpr(R"DONE(
+; Includes
 
 [
-  (compiler_directive)
   (import)
-] @defcolor_macro
+  (load)
+] @include
 
 ; Keywords
-; TODO : complete this list
 [
-  "struct"
-  "enum"
-  "defer"
-  "cast"
+  ; from modules/Jai_Lexer
+  "if"
   "xx"
-  "return"
-] @defcolor_keyword
 
-; Conditionals
+  "ifx"
+  "for"
+
+  "then"
+  "else"
+  "null"
+  "case"
+  "enum"
+  "true"
+  "cast"
+
+  "while"
+  "break"
+  "using"
+  "defer"
+  "false"
+  "union"
+
+  "return"
+  "struct"
+  "inline"
+  "remove"
+
+  ; "size_of"
+  "type_of"
+  ; "code_of"
+  ; "context"
+
+  "continue"
+  "operator"
+
+  ; "type_info"
+  "no_inline"
+  "interface"
+
+  "enum_flags"
+
+  ; "is_constant"
+
+  "push_context"
+
+  ; "initializer_of"
+] @keyword
+
+[
+  "return"
+] @keyword.return
+
 [
   "if"
   "else"
   "case"
   "break"
-] @defcolor_keyword
+] @keyword.conditional
 
 ((if_expression
   [
     "then"
     "ifx"
     "else"
-  ] @defcolor_keyword)
+  ] @keyword.conditional.ternary)
   (#set! "priority" 105))
 
 ; Repeats
@@ -148,132 +103,123 @@ String_Const_u8 TS_JAI_HIGHLIGHT_QUERY_ = string_u8_litexpr(R"DONE(
   "for"
   "while"
   "continue"
-] @defcolor_keyword
+] @keyword.repeat
 
 ; Variables
 
-(identifier) @defcolor_text_default
+; (identifier) @variable
+name: (identifier) @variable
+argument: (identifier) @variable
+named_argument: (identifier) @variable
+(member_expression (identifier) @variable)
+(parenthesized_expression (identifier) @variable)
+
+((identifier) @variable.builtin
+  (#any-of? @variable.builtin "context"))
 
 ; Namespaces
 
-(import (identifier) @defcolor_text_default)
+(import (identifier) @namespace)
 
 ; Parameters
 
-(parameter (identifier) @defcolor_text_default ":" "="? (identifier)? @defcolor_str_constant)
+(parameter (identifier) @parameter ":" "="? (identifier)? @constant)
 
-(default_parameter (identifier) @defcolor_text_default ":=")
-
-(call_expression argument: (identifier) @defcolor_text_default "=")
+; (call_expression argument: (identifier) @parameter "=")
 
 ; Functions
 
-(procedure_declaration (identifier) @defcolor_function)
+; (procedure_declaration (identifier) @function (procedure (block)))
+(procedure_declaration (identifier) @function (block))
 
-(procedure_declaration (identifier) @defcolor_function (procedure (block)))
-
-(call_expression function: (identifier) @defcolor_function)
+(call_expression function: (identifier) @function.call)
 
 ; Types
 
-(type (identifier) @defcolor_type)
+type: (types) @type
+type: (identifier) @type
+((types) @type)
 
-((type (identifier) @defcolor_type)
+modifier: (identifier) @keyword
+keyword: (identifier) @keyword
+
+((types (identifier) @type.builtin)
   (#any-of? @type.builtin
-    "bool"
-    "int" "s8" "s16" "s32" "s64"
+    "bool" "int" "string"
+    "s8" "s16" "s32" "s64"
     "u8" "u16" "u32" "u64"
-    "string"))
+    "Type" "Any"))
 
+(struct_declaration (identifier) @type ":" ":")
 
-(struct_declaration (identifier) @defcolor_type "::")
+(enum_declaration (identifier) @type ":" ":")
 
-(enum_declaration (identifier) @defcolor_type "::")
+; (const_declaration (identifier) @type ":" ":" [(array_type) (pointer_type)])
 
-;(union_declaration (identifier) @defcolor_type "::")
-
-(const_declaration (identifier) @defcolor_type "::" [(array_type) (pointer_type)])
-
-(struct . (identifier) @defcolor_type)
-
-;(field_type . (identifier) @namespace "." (identifier) @defcolor_type)
-
-;(bit_set_type (identifier) @defcolor_type ";")
-
-;(procedure_type (parameters (parameter (identifier) @defcolor_type)))
-
-;(polymorphic_parameters (identifier) @defcolor_type)
-
-((identifier) @defcolor_type
-  (#lua-match? @defcolor_type "^[A-Z][a-zA-Z0-9]*$")
-  (#not-has-parent? @defcolor_type parameter procedure_declaration call_expression))
+; ; I don't like this
+; ((identifier) @type
+;   (#lua-match? @type "^[A-Z][a-zA-Z0-9]*$")
+;   (#not-has-parent? @type parameter procedure_declaration call_expression))
 
 ; Fields
 
-(member_expression "." (identifier) @defcolor_text_default)
+(member_expression "." (identifier) @field)
 
-;(struct_type "{" (identifier) @defcolor_text_default)
-
-(struct_field (identifier) @defcolor_text_default "="?)
-
-(field (identifier) @defcolor_text_default)
+(assignment_statement (identifier) @field "="?)
+(update_statement (identifier) @field)
 
 ; Constants
 
-((identifier) @defcolor_text_default
-  (#lua-match? @defcolor_str_constnat "^_*[A-Z][A-Z0-9_]*$")
-  (#not-has-parent? @text_default type parameter))
+((identifier) @constant
+  (#lua-match? @constant "^_*[A-Z][A-Z0-9_]*$")
+  (#not-has-parent? @constant type parameter))
 
-(member_expression . "." (identifier) @defcolor_text_default)
+(member_expression . "." (identifier) @constant)
 
-(enum_declaration "{" (identifier) @defcolor_text_default)
+(enum_declaration "{" (identifier) @constant)
 
 ; Literals
 
-(number) @defcolor_int_constant
+(integer) @number
+(float) @number
 
-(float) @defcolor_float_constant
+(string) @string
 
-(string) @defcolor_str_constnat
+;(character) @character
 
-;(character) @defcolor_str_constnat
+(string (escape_sequence) @string.escape)
 
-(escape_sequence) @defcolor_str_constant
-
-(boolean) @defcolor_bool_constant
+(boolean) @boolean
 
 [
   (uninitialized)
   (null)
-] @defcolor_text_default
-
-((identifier) @defcolor_text_default
-  (#any-of? @defcolor_text_default "context"))
+] @constant.builtin
 
 ; Operators
 
 [
-  ":="
+  ":"
   "="
   "+"
   "-"
   "*"
   "/"
   "%"
-  "%%"
   ">"
   ">="
   "<"
   "<="
   "=="
   "!="
-  "~="
   "|"
   "~"
   "&"
   "&~"
   "<<"
   ">>"
+  "<<<"
+  ">>>"
   "||"
   "&&"
   "!"
@@ -288,41 +234,51 @@ String_Const_u8 TS_JAI_HIGHLIGHT_QUERY_ = string_u8_litexpr(R"DONE(
   "^="
   "<<="
   ">>="
+  "<<<="
+  ">>>="
   "||="
   "&&="
-  "&~="
-  ;"..="
-  ;"..<"
-  ;"?"
-] @defcolor_operator
+] @operator
 
 ; Punctuation
 
-[ "{" "}" ] @defcolor_text_default
+[ "{" "}" ] @punctuation.bracket
 
-[ "(" ")" ] @defcolor_text_default
+[ "(" ")" ] @punctuation.bracket
 
-[ "[" "]" ] @defcolor_text_default
+[ "[" "]" ] @punctuation.bracket
 
 [
-  "::"
+  "`"
   "->"
   "."
   ","
   ":"
   ";"
-] @defcolor_text_default
+] @punctuation.delimiter
 
 ; Comments
 
 [
-  (comment)
   (block_comment)
-] @defcolor_comment
+  (comment)
+] @comment @spell
 
 ; Errors
 
-(ERROR) @defcolor_comment_pop
+(ERROR) @error
+
+(block_comment) @comment
+
+directive: ("#") @keyword ; #if
+type: ("type_of") @type
+
+(compiler_directive) @keyword
+(heredoc_start) @none
+(heredoc_end) @none
+(heredoc_body) @string
+(note) @string
+
 )DONE");
 
 extern "C" {
@@ -335,8 +291,7 @@ tree_sitter_register_jai (Application_Links* app)
   TSLanguage* language = tree_sitter_jai();
   Tree_Sitter_Language_Queries queries = {};
   queries.ptr[Tree_Sitter_Language_Query_Highlights] = tree_sitter_query_new(app, language, TS_JAI_HIGHLIGHT_QUERY);
-  queries.ptr[Tree_Sitter_Language_Query_Functions]  = tree_sitter_query_new(app, language, TS_JAI_FUNCTION_QUERY);
-  queries.ptr[Tree_Sitter_Language_Query_Types]      = tree_sitter_query_new(app, language, TS_JAI_TYPE_QUERY);
+  queries.ptr[Tree_Sitter_Language_Query_Tags]       = tree_sitter_query_new(app, language, TS_JAI_TAGS_QUERY);
   tree_sitter_register_language(SCu8("jai"), language, queries);
 }