treesitter: runtime queries

Runtime queries just work like ftplugins, that is:

- Queries in the `after` directory are sourced _after_ the "base" query
- Otherwise, the last define query takes precedence.

Queries can be found in the `queries` directory.

Update runtime/lua/vim/treesitter/query.lua

Co-authored-by: Paul Burlumi <paul@burlumi.com>

4 files changed, 261 insertions, 3 deletions

diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 0f497fe434..47114a306f 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -56,7 +56,7 @@ TSHighlighter.hl_map = {
     ["include"] = "Include",
 }
 
-function TSHighlighter.new(query, bufnr, ft)
+function TSHighlighter.new(bufnr, ft, query)
   if bufnr == nil or bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 494fb59fa7..2903c5905c 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -8,6 +8,104 @@ Query.__index = Query
 
 local M = {}
 
+-- Filter the runtime query files, the spec is like regular runtime files but in the new `queries`
+-- directory. They resemble ftplugins, that is that you can override queries by adding things in the
+-- `queries` directory, and extend using the `after/queries` directory.
+local function filter_files(file_list)
+  local main = nil
+  local after = {}
+
+  for _, fname in ipairs(file_list) do
+    -- Only get the name of the directory containing the queries directory
+    if vim.fn.fnamemodify(fname, ":p:h:h:h:t") == "after" then
+      table.insert(after, fname)
+    -- The first one is the one with most priority
+    elseif not main then
+      main = fname
+    end
+  end
+
+  return { main, unpack(after) }
+end
+
+local function runtime_query_path(lang, query_name)
+  return string.format('queries/%s/%s.scm', lang, query_name)
+end
+
+local function filtered_runtime_queries(lang, query_name)
+  return filter_files(a.nvim_get_runtime_file(runtime_query_path(lang, query_name), true) or {})
+end
+
+local function get_query_files(lang, query_name, is_included)
+  local lang_files = filtered_runtime_queries(lang, query_name)
+  local query_files = lang_files
+
+  if #query_files == 0 then return {} end
+
+  local base_langs = {}
+
+  -- Now get the base languages by looking at the first line of every file
+  -- The syntax is the folowing :
+  -- ;+ inherits: ({language},)*{language}
+  --
+  -- {language} ::= {lang} | ({lang})
+  local MODELINE_FORMAT = "^;+%s*inherits%s*:?%s*([a-z_,()]+)%s*$"
+
+  for _, file in ipairs(query_files) do
+    local modeline = vim.fn.readfile(file, "", 1)
+
+    if #modeline == 1 then
+      local langlist = modeline[1]:match(MODELINE_FORMAT)
+
+      if langlist then
+        for _, incllang in ipairs(vim.split(langlist, ',', true)) do
+          local is_optional = incllang:match("%(.*%)")
+
+          if is_optional then
+            if not is_included then
+              table.insert(base_langs, incllang:sub(2, #incllang - 1))
+            end
+          else
+            table.insert(base_langs, incllang)
+          end
+        end
+      end
+    end
+  end
+
+  for _, base_lang in ipairs(base_langs) do
+    local base_files = get_query_files(base_lang, query_name, true)
+    vim.list_extend(query_files, base_files)
+  end
+
+  return query_files
+end
+
+local function read_query_files(filenames)
+  local contents = {}
+
+  for _,filename in ipairs(filenames) do
+    vim.list_extend(contents, vim.fn.readfile(filename))
+  end
+
+  return table.concat(contents, '\n')
+end
+
+--- Returns the runtime query {query_name} for {lang}.
+--
+-- @param lang The language to use for the query
+-- @param query_name The name of the query (i.e. "highlights")
+--
+-- @return The corresponding query, parsed.
+function M.get_query(lang, query_name)
+  local query_files = get_query_files(lang, query_name)
+  local query_string = read_query_files(query_files)
+
+  if #query_string > 0 then
+    return M.parse_query(lang, query_string)
+  end
+end
+
 --- Parses a query.
 --
 -- @param language The language
diff --git a/runtime/queries/c/highlights.scm b/runtime/queries/c/highlights.scm
new file mode 100644
index 0000000000..96b43cf0d0
--- /dev/null
+++ b/runtime/queries/c/highlights.scm
@@ -0,0 +1,151 @@
+(identifier) @variable
+
+[
+  "const"
+  "default"
+  "enum"
+  "extern"
+  "inline"
+  "return"
+  "sizeof"
+  "static"
+  "struct"
+  "typedef"
+  "union"
+  "volatile"
+  "goto"
+] @keyword
+
+[
+  "while"
+  "for"
+  "do"
+  "continue"
+  "break"
+] @repeat
+
+[
+ "if"
+ "else"
+ "case"
+ "switch"
+] @conditional
+
+"#define" @constant.macro
+[
+  "#if"
+  "#ifdef"
+  "#ifndef"
+  "#else"
+  "#elif"
+  "#endif"
+  (preproc_directive)
+] @keyword
+
+"#include" @include
+
+[
+  "="
+
+  "-"
+  "*"
+  "/"
+  "+"
+  "%"
+
+  "~"
+  "|"
+  "&"
+  "^"
+  "<<"
+  ">>"
+
+  "->"
+
+  "<"
+  "<="
+  ">="
+  ">"
+  "=="
+  "!="
+
+  "!"
+  "&&"
+  "||"
+
+  "-="
+  "+="
+  "*="
+  "/="
+  "%="
+  "|="
+  "&="
+  "^="
+  "--"
+  "++"
+] @operator
+
+[
+ (true)
+ (false)
+] @boolean
+
+[ "." ";" ":" "," ] @punctuation.delimiter
+
+(conditional_expression [ "?" ":" ] @conditional)
+
+
+[ "(" ")" "[" "]" "{" "}"] @punctuation.bracket
+
+(string_literal) @string
+(system_lib_string) @string
+
+(null) @constant.builtin
+(number_literal) @number
+(char_literal) @number
+
+(call_expression
+  function: (identifier) @function)
+(call_expression
+  function: (field_expression
+    field: (field_identifier) @function))
+(function_declarator
+  declarator: (identifier) @function)
+(preproc_function_def
+  name: (identifier) @function.macro)
+[
+ (preproc_arg)
+ (preproc_defined)
+]  @function.macro
+; TODO (preproc_arg)  @embedded
+
+(field_identifier) @property
+(statement_identifier) @label
+
+[
+(type_identifier)
+(primitive_type)
+(sized_type_specifier)
+(type_descriptor)
+ ] @type
+
+(declaration type: [(identifier) (type_identifier)] @type)
+(cast_expression type: [(identifier) (type_identifier)] @type)
+(sizeof_expression value: (parenthesized_expression (identifier) @type))
+
+((identifier) @constant
+ (#match? @constant "^[A-Z][A-Z0-9_]+$"))
+
+(comment) @comment
+
+;; Parameters
+(parameter_declaration
+  declarator: (identifier) @parameter)
+
+(parameter_declaration
+  declarator: (pointer_declarator) @parameter)
+
+(preproc_params
+  (identifier)) @parameter
+
+(ERROR) @error
diff --git a/test/functional/lua/treesitter_spec.lua b/test/functional/lua/treesitter_spec.lua
index 74ae6cde2b..35dd6b368a 100644
--- a/test/functional/lua/treesitter_spec.lua
+++ b/test/functional/lua/treesitter_spec.lua
@@ -25,7 +25,6 @@ describe('treesitter API', function()
     eq("Error executing lua: .../language.lua: no parser for 'borklang' language, see :help treesitter-parsers",
        pcall_err(exec_lua, "parser = vim.treesitter.inspect_language('borklang')"))
   end)
-
 end)
 
 describe('treesitter API with C parser', function()
@@ -186,6 +185,16 @@ void ui_refresh(void)
     (field_expression argument: (identifier) @fieldarg)
   ]]
 
+  it("supports runtime queries", function()
+    if not check_parser() then return end
+
+    local ret = exec_lua [[
+      return require"vim.treesitter.query".get_query("c", "highlights").captures[1]
+    ]]
+
+    eq('variable', ret)
+  end)
+
   it('support query and iter by capture', function()
     if not check_parser() then return end
 
@@ -422,7 +431,7 @@ static int nlua_schedule(lua_State *const lstate)
     exec_lua([[
       local highlighter = vim.treesitter.highlighter
       local query = ...
-      test_hl = highlighter.new(query, 0, "c")
+      test_hl = highlighter.new(0, "c", query)
     ]], hl_query)
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |