2 files changed, 57 insertions, 6 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index eb7ea04254..9f688da963 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -194,6 +194,32 @@ and predicates. A `capture` allows you to associate names with a specific
 node in a pattern. A `predicate` adds arbitrary metadata and conditional data
 to a match.
 
+Neovim supports to customize the behavior of the queries using a set of
+"modelines", that is comments in the queries starting with `;`. Here are the
+currently supported modeline alternatives:
+
+    `inherits: {lang}...`
+        Specifies that this query should inherit the queries from {lang}.
+        This will recursively descend in the queries of {lang} unless wrapped
+        in parentheses: `({lang})`.
+
+    `extends`
+        Specifies that this query should be used as an extension for the
+        query, i.e. that it should be merged with the others.
+        Note: the order of the extensions, and the query that will be used as
+        a base depends on your 'runtimepath' value.
+
+Note: these modeline comments must be at the top of the query, but can be
+repeated, for example, the following modeline blocks are all valid:
+>
+    ;; inherits: foo,bar
+    ;; extends
+
+    ;; extends
+    ;;
+    ;; inherits: baz
+<
+
 Treesitter Query Predicates                         *lua-treesitter-predicates*
 
 When writing queries for treesitter, one might use `predicates`, that is,
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 78042a74bf..be5b24fd95 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -47,6 +47,9 @@ function M.get_query_files(lang, query_name, is_included)
     return {}
   end
 
+  local base_query = nil
+  local extensions = {}
+
   local base_langs = {}
 
   -- Now get the base languages by looking at the first line of every file
@@ -55,13 +58,26 @@ function M.get_query_files(lang, query_name, is_included)
   --
   -- {language} ::= {lang} | ({lang})
   local MODELINE_FORMAT = '^;+%s*inherits%s*:?%s*([a-z_,()]+)%s*$'
+  local EXTENDS_FORMAT = '^;+%s*extends%s*$'
 
-  for _, file in ipairs(lang_files) do
-    local modeline = safe_read(file, '*l')
+  for _, filename in ipairs(lang_files) do
+    local file, err = io.open(filename, 'r')
+    if not file then
+      error(err)
+    end
 
-    if modeline then
-      local langlist = modeline:match(MODELINE_FORMAT)
+    local extension = false
+
+    for modeline in
+      function()
+        return file:read('*l')
+      end
+    do
+      if not vim.startswith(modeline, ';') then
+        break
+      end
 
+      local langlist = modeline:match(MODELINE_FORMAT)
       if langlist then
         for _, incllang in ipairs(vim.split(langlist, ',', true)) do
           local is_optional = incllang:match('%(.*%)')
@@ -74,16 +90,25 @@ function M.get_query_files(lang, query_name, is_included)
             table.insert(base_langs, incllang)
           end
         end
+      elseif modeline:match(EXTENDS_FORMAT) then
+        extension = true
       end
     end
+
+    if extension then
+      table.insert(extensions, filename)
+    else
+      base_query = filename
+    end
+    io.close(file)
   end
 
-  local query_files = {}
+  local query_files = { base_query }
   for _, base_lang in ipairs(base_langs) do
     local base_files = M.get_query_files(base_lang, query_name, true)
     vim.list_extend(query_files, base_files)
   end
-  vim.list_extend(query_files, lang_files)
+  vim.list_extend(query_files, extensions)
 
   return query_files
 end