From baf159c9e5a054a659d94be81ea64d15119d8e2b Mon Sep 17 00:00:00 2001
From: Iron-E <code.iron.e@gmail.com>
Date: Thu, 4 May 2023 12:09:42 -0400
Subject: [PATCH] ref(rust): improve TS highlighting

---
 after/queries/rust/highlights.scm | 54 +++++++++++++++++++++++++------
 ftplugin/rust/highlite.lua        |  4 ++-
 lua/highlite/groups/default.lua   |  4 +--
 3 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/after/queries/rust/highlights.scm b/after/queries/rust/highlights.scm
index 6a9b142..0143c85 100644
--- a/after/queries/rust/highlights.scm
+++ b/after/queries/rust/highlights.scm
@@ -41,10 +41,15 @@
 )
 
 ; enum variants
-(enum_variant name: (identifier) @constant !body)
-(enum_variant
-	name: (identifier) @type
-	body: [(field_declaration_list) (ordered_field_declaration_list)]
+(call_expression
+	function: (scoped_identifier
+		name: (identifier) @type (#lua-match? @type "^%u")
+	)
+)
+
+(enum_variant name: (identifier) @type)
+((identifier) @type.builtin
+	(#any-of? @type.builtin "Some" "None" "Ok" "Err")
 )
 
 ; HRTB closures
@@ -52,13 +57,13 @@
 
 ; macros
 (macro_invocation
-	macro: (identifier) @_id @debug (#contains? @_id "dbg" "debug")
+	macro: (identifier) @_id @debug (#contains? @_id "assert" "dbg" "debug")
 	"!" @debug
 	(#set! "priority" 101)
 )
 
 (macro_invocation
-	macro: (identifier) @_id @include (#lua-match? @_id "^include")
+	macro: (identifier) @_id @include (#contains? @_id "include")
 	"!" @include
 	(#set! "priority" 101)
 )
@@ -67,18 +72,47 @@
 (match_pattern "_" @variable.builtin)
 
 ; raw identifiers
-((field_identifier) @punctuation.special
+(
+	[
+		(field_identifier) @punctuation.special
+		(identifier) @punctuation.special
+	]
 	(#lua-match? @punctuation.special "^r#")
 	(#offset-from-start! @punctuation.special 0 0 0 2)
 )
 
-((identifier) @punctuation.special
+((raw_string_literal) @punctuation.special
 	(#lua-match? @punctuation.special "^r#")
 	(#offset-from-start! @punctuation.special 0 0 0 2)
 )
 
-; trait bounds
-(removed_trait_bound "?" @punctuation.special)
+((raw_string_literal) @punctuation.special
+	(#lua-match? @punctuation.special "^r#")
+	(#offset-from-end! @punctuation.special 0 -1 0 0)
+)
+
+((raw_string_literal) @punctuation.special
+	(#lua-match? @punctuation.special "^r##")
+	(#offset-from-start! @punctuation.special 0 0 0 3)
+)
+
+((raw_string_literal) @punctuation.special
+	(#lua-match? @punctuation.special "^r##")
+	(#offset-from-end! @punctuation.special 0 -2 0 0)
+)
+
+((raw_string_literal) @punctuation.special
+	(#lua-match? @punctuation.special "^r###")
+	(#offset-from-start! @punctuation.special 0 0 0 4)
+)
+
+((raw_string_literal) @punctuation.special
+	(#lua-match? @punctuation.special "^r###")
+	(#offset-from-end! @punctuation.special 0 -3 0 0)
+)
 
 ; `try!`
 (try_expression ("?" @exception))
+
+; unsafe
+"unsafe" @exception
diff --git a/ftplugin/rust/highlite.lua b/ftplugin/rust/highlite.lua
index 3c019fc..4f62be3 100644
--- a/ftplugin/rust/highlite.lua
+++ b/ftplugin/rust/highlite.lua
@@ -1,3 +1,5 @@
 if vim.api.nvim_get_var('colors_name'):find '^highlite' then
-	require('highlite.treesitter').add_directive.offset_from_start()
+	local Treesitter = require 'highlite.treesitter' --- @type highlite.Treesitter
+	Treesitter.add_directive.offset_from_end()
+	Treesitter.add_directive.offset_from_start()
 end
diff --git a/lua/highlite/groups/default.lua b/lua/highlite/groups/default.lua
index f8c4b5f..39d5656 100644
--- a/lua/highlite/groups/default.lua
+++ b/lua/highlite/groups/default.lua
@@ -377,13 +377,13 @@ local function from_palette(palette, opts)
 		['@lsp.type.builtinType.rust'] = '@type.builtin.rust',
 		['@lsp.type.decorator.rust'] = '@preproc.rust',
 		['@lsp.type.derive.rust'] = '@macro.rust',
-		['@lsp.type.enumMember.rust'] = NONE,
+		['@lsp.type.enumMember.rust'] = '@type',
 		['@lsp.type.keyword.rust'] = NONE,
 		['@lsp.type.macro.rust'] = NONE,
 		['@lsp.type.operator.rust'] = NONE,
 		['@lsp.type.selfKeyword.rust'] = '@variable.builtin.rust',
 		['@lsp.type.selfTypeKeyword.rust'] = '@lsp.type.typeAlias.rust',
-		['@lsp.typemod.enumMember.injected.rust'] = '@lsp.type.enumMember',
+		['@lsp.typemod.enumMember.defaultLibrary.rust'] = '@type.builtin',
 		['@lsp.typemod.keyword.injected.rust'] = '@keyword.rust',
 		['@lsp.typemod.macro.injected.rust'] = '@macro.rust',
 		['@lsp.typemod.operator.injected.rust'] = '@operator.rust',