2 files changed, 45 insertions, 7 deletions
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 6b833c6d35..afb72dc069 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -12,6 +12,16 @@ TSHighlighterQuery.__index = TSHighlighterQuery
 
 local ns = a.nvim_create_namespace("treesitter/highlighter")
 
+local _default_highlights = {}
+local _link_default_highlight_once = function(from, to)
+  if not _default_highlights[from] then
+    _default_highlights[from] = true
+    vim.cmd(string.format("highlight default link %s %s", from, to))
+  end
+
+  return from
+end
+
 -- These are conventions defined by nvim-treesitter, though it
 -- needs to be user extensible also.
 TSHighlighter.hl_map = {
@@ -70,9 +80,12 @@ function TSHighlighterQuery.new(lang, query_string)
 
   self.hl_cache = setmetatable({}, {
     __index = function(table, capture)
-      local hl = self:get_hl_from_capture(capture)
-      rawset(table, capture, hl)
+      local hl, is_vim_highlight = self:_get_hl_from_capture(capture)
+      if not is_vim_highlight then
+        hl = _link_default_highlight_once(lang .. hl, hl)
+      end
 
+      rawset(table, capture, hl)
       return hl
     end
   })
@@ -90,16 +103,16 @@ function TSHighlighterQuery:query()
   return self._query
 end
 
-function TSHighlighterQuery:get_hl_from_capture(capture)
+--- Get the hl from capture.
+--- Returns a tuple { highlight_name: string, is_builtin: bool }
+function TSHighlighterQuery:_get_hl_from_capture(capture)
   local name = self._query.captures[capture]
 
   if is_highlight_name(name) then
     -- From "Normal.left" only keep "Normal"
-    return vim.split(name, '.', true)[1]
+    return vim.split(name, '.', true)[1], true
   else
-    -- Default to false to avoid recomputing
-    local hl = TSHighlighter.hl_map[name]
-    return hl and a.nvim_get_hl_id_by_name(hl) or 0
+    return TSHighlighter.hl_map[name] or name, false
   end
 end
 
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 1b29618997..188ec94a6a 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -90,6 +90,27 @@ local function new_match_metadata()
   return setmetatable({}, match_metatable)
 end
 
+--- The explicitly set queries from |vim.treesitter.query.set_query()|
+local explicit_queries = setmetatable({}, {
+  __index = function(t, k)
+    local lang_queries = {}
+    rawset(t, k, lang_queries)
+
+    return lang_queries
+  end,
+})
+
+--- Sets the runtime query {query_name} for {lang}
+---
+--- This allows users to override any runtime files and/or configuration
+--- set by plugins.
+---@param lang string: The language to use for the query
+---@param query_name string: The name of the query (i.e. "highlights")
+---@param text string: The query text (unparsed).
+function M.set_query(lang, query_name, text)
+  explicit_queries[lang][query_name] = M.parse_query(lang, text)
+end
+
 --- Returns the runtime query {query_name} for {lang}.
 --
 -- @param lang The language to use for the query
@@ -97,6 +118,10 @@ end
 --
 -- @return The corresponding query, parsed.
 function M.get_query(lang, query_name)
+  if explicit_queries[lang][query_name] then
+    return explicit_queries[lang][query_name]
+  end
+
   local query_files = M.get_query_files(lang, query_name)
   local query_string = read_query_files(query_files)