feat(highlight): re-apply default highlights on colorscheme changes

If the plugin is loaded after a colorscheme is set that defines any of
these highlight groups, the default won't be applied. Subsequent
"highlight clear" commands (common when switching colorschemes) will
then clear any of those highlights, but these defaults never have
another opportunity to be initialized.

Effectively, if you load neovim with a colorscheme that has definitions
for some of these highlight groups, then load treesitter, then switch
colorschemes, many of these default links will be absent resulting in
colors that do not appear the same as if that colorscheme had been
used at startup.

Hooking the ColorScheme event with an autocmd that just reapplies these
defaults gives every colorscheme switch the opportunity to get the
defaults for non-explicitly-defined groups.

2 files changed, 96 insertions, 85 deletions

diff --git a/lua/nvim-treesitter/highlight.lua b/lua/nvim-treesitter/highlight.lua
index caca7ee00..b35467126 100644
--- a/lua/nvim-treesitter/highlight.lua
+++ b/lua/nvim-treesitter/highlight.lua
@@ -156,4 +156,92 @@ function M.set_custom_captures(captures)
   end
 end
 
+function M.set_default_hlgroups()
+  local highlights = {
+    TSNone = { default = true },
+    TSPunctDelimiter = { link = "Delimiter", default = true },
+    TSPunctBracket = { link = "Delimiter", default = true },
+    TSPunctSpecial = { link = "Delimiter", default = true },
+
+    TSConstant = { link = "Constant", default = true },
+    TSConstBuiltin = { link = "Special", default = true },
+    TSConstMacro = { link = "Define", default = true },
+    TSString = { link = "String", default = true },
+    TSStringRegex = { link = "String", default = true },
+    TSStringEscape = { link = "SpecialChar", default = true },
+    TSStringSpecial = { link = "SpecialChar", default = true },
+    TSCharacter = { link = "Character", default = true },
+    TSCharacterSpecial = { link = "SpecialChar", default = true },
+    TSNumber = { link = "Number", default = true },
+    TSBoolean = { link = "Boolean", default = true },
+    TSFloat = { link = "Float", default = true },
+
+    TSFunction = { link = "Function", default = true },
+    TSFuncBuiltin = { link = "Special", default = true },
+    TSFuncMacro = { link = "Macro", default = true },
+    TSParameter = { link = "Identifier", default = true },
+    TSParameterReference = { link = "TSParameter", default = true },
+    TSMethod = { link = "Function", default = true },
+    TSField = { link = "Identifier", default = true },
+    TSProperty = { link = "Identifier", default = true },
+    TSConstructor = { link = "Special", default = true },
+    TSAnnotation = { link = "PreProc", default = true },
+    TSAttribute = { link = "PreProc", default = true },
+    TSNamespace = { link = "Include", default = true },
+    TSSymbol = { link = "Identifier", default = true },
+
+    TSConditional = { link = "Conditional", default = true },
+    TSRepeat = { link = "Repeat", default = true },
+    TSLabel = { link = "Label", default = true },
+    TSOperator = { link = "Operator", default = true },
+    TSKeyword = { link = "Keyword", default = true },
+    TSKeywordFunction = { link = "Keyword", default = true },
+    TSKeywordOperator = { link = "TSOperator", default = true },
+    TSKeywordReturn = { link = "TSKeyword", default = true },
+    TSException = { link = "Exception", default = true },
+    TSDebug = { link = "Debug", default = true },
+    TSDefine = { link = "Define", default = true },
+    TSPreProc = { link = "PreProc", default = true },
+    TSStorageClass = { link = "StorageClass", default = true },
+
+    TSTodo = { link = "Todo", default = true },
+
+    TSType = { link = "Type", default = true },
+    TSTypeBuiltin = { link = "Type", default = true },
+    TSTypeQualifier = { link = "Type", default = true },
+    TSTypeDefinition = { link = "Typedef", default = true },
+
+    TSInclude = { link = "Include", default = true },
+
+    TSVariableBuiltin = { link = "Special", default = true },
+
+    TSText = { link = "TSNone", default = true },
+    TSStrong = { bold = true, default = true },
+    TSEmphasis = { italic = true, default = true },
+    TSUnderline = { underline = true },
+    TSStrike = { strikethrough = true },
+
+    TSMath = { link = "Special", default = true },
+    TSTextReference = { link = "Constant", default = true },
+    TSEnvironment = { link = "Macro", default = true },
+    TSEnvironmentName = { link = "Type", default = true },
+    TSTitle = { link = "Title", default = true },
+    TSLiteral = { link = "String", default = true },
+    TSURI = { link = "Underlined", default = true },
+
+    TSComment = { link = "Comment", default = true },
+    TSNote = { link = "SpecialComment", default = true },
+    TSWarning = { link = "Todo", default = true },
+    TSDanger = { link = "WarningMsg", default = true },
+
+    TSTag = { link = "Label", default = true },
+    TSTagDelimiter = { link = "Delimiter", default = true },
+    TSTagAttribute = { link = "TSProperty", default = true },
+  }
+
+  for k, v in pairs(highlights) do
+    api.nvim_set_hl(0, k, v)
+  end
+end
+
 return M
diff --git a/plugin/nvim-treesitter.lua b/plugin/nvim-treesitter.lua
index 5109c80f2..2eedd3372 100644
--- a/plugin/nvim-treesitter.lua
+++ b/plugin/nvim-treesitter.lua
@@ -9,6 +9,7 @@ vim.g.loaded_nvim_treesitter = true
 require("nvim-treesitter").setup()
 
 local api = vim.api
+local highlight = require "nvim-treesitter.highlight"
 
 -- define autocommands
 local augroup = api.nvim_create_augroup("NvimTreesitter", {})
@@ -33,89 +34,11 @@ api.nvim_create_autocmd("Filetype", {
   desc = "Reload query",
 })
 
--- define highlights
-local highlights = {
-  TSNone = { default = true },
-  TSPunctDelimiter = { link = "Delimiter", default = true },
-  TSPunctBracket = { link = "Delimiter", default = true },
-  TSPunctSpecial = { link = "Delimiter", default = true },
-
-  TSConstant = { link = "Constant", default = true },
-  TSConstBuiltin = { link = "Special", default = true },
-  TSConstMacro = { link = "Define", default = true },
-  TSString = { link = "String", default = true },
-  TSStringRegex = { link = "String", default = true },
-  TSStringEscape = { link = "SpecialChar", default = true },
-  TSStringSpecial = { link = "SpecialChar", default = true },
-  TSCharacter = { link = "Character", default = true },
-  TSCharacterSpecial = { link = "SpecialChar", default = true },
-  TSNumber = { link = "Number", default = true },
-  TSBoolean = { link = "Boolean", default = true },
-  TSFloat = { link = "Float", default = true },
-
-  TSFunction = { link = "Function", default = true },
-  TSFuncBuiltin = { link = "Special", default = true },
-  TSFuncMacro = { link = "Macro", default = true },
-  TSParameter = { link = "Identifier", default = true },
-  TSParameterReference = { link = "TSParameter", default = true },
-  TSMethod = { link = "Function", default = true },
-  TSField = { link = "Identifier", default = true },
-  TSProperty = { link = "Identifier", default = true },
-  TSConstructor = { link = "Special", default = true },
-  TSAnnotation = { link = "PreProc", default = true },
-  TSAttribute = { link = "PreProc", default = true },
-  TSNamespace = { link = "Include", default = true },
-  TSSymbol = { link = "Identifier", default = true },
-
-  TSConditional = { link = "Conditional", default = true },
-  TSRepeat = { link = "Repeat", default = true },
-  TSLabel = { link = "Label", default = true },
-  TSOperator = { link = "Operator", default = true },
-  TSKeyword = { link = "Keyword", default = true },
-  TSKeywordFunction = { link = "Keyword", default = true },
-  TSKeywordOperator = { link = "TSOperator", default = true },
-  TSKeywordReturn = { link = "TSKeyword", default = true },
-  TSException = { link = "Exception", default = true },
-  TSDebug = { link = "Debug", default = true },
-  TSDefine = { link = "Define", default = true },
-  TSPreProc = { link = "PreProc", default = true },
-  TSStorageClass = { link = "StorageClass", default = true },
-
-  TSTodo = { link = "Todo", default = true },
-
-  TSType = { link = "Type", default = true },
-  TSTypeBuiltin = { link = "Type", default = true },
-  TSTypeQualifier = { link = "Type", default = true },
-  TSTypeDefinition = { link = "Typedef", default = true },
-
-  TSInclude = { link = "Include", default = true },
-
-  TSVariableBuiltin = { link = "Special", default = true },
-
-  TSText = { link = "TSNone", default = true },
-  TSStrong = { bold = true, default = true },
-  TSEmphasis = { italic = true, default = true },
-  TSUnderline = { underline = true },
-  TSStrike = { strikethrough = true },
-
-  TSMath = { link = "Special", default = true },
-  TSTextReference = { link = "Constant", default = true },
-  TSEnvironment = { link = "Macro", default = true },
-  TSEnvironmentName = { link = "Type", default = true },
-  TSTitle = { link = "Title", default = true },
-  TSLiteral = { link = "String", default = true },
-  TSURI = { link = "Underlined", default = true },
-
-  TSComment = { link = "Comment", default = true },
-  TSNote = { link = "SpecialComment", default = true },
-  TSWarning = { link = "Todo", default = true },
-  TSDanger = { link = "WarningMsg", default = true },
-
-  TSTag = { link = "Label", default = true },
-  TSTagDelimiter = { link = "Delimiter", default = true },
-  TSTagAttribute = { link = "TSProperty", default = true },
-}
+api.nvim_create_autocmd("ColorScheme", {
+  group = augroup,
+  callback = highlight.set_default_hlgroups,
+  desc = "Set default highlights",
+})
 
-for k, v in pairs(highlights) do
-  api.nvim_set_hl(0, k, v)
-end
+-- define highlights
+highlight.set_default_hlgroups()