From c055899dc0739fb48541365dfffda852034438c2 Mon Sep 17 00:00:00 2001
From: Thomas Vigouroux <tomvig38@gmail.com>
Date: Tue, 6 Oct 2020 20:03:39 +0200
Subject: [PATCH] feat(queries): modeline mechanism for base langs

This implements https://github.com/neovim/neovim/pull/13059#issuecomment-704414189

This behaves like modelines and remove the use of the base_language map.
Also, this allows to fine-tune what we actually want to include per
query, which is better IMO.
---
 CONTRIBUTING.md                        | 12 +++++
 lua/nvim-treesitter/query.lua          | 72 ++++++++++++++------------
 queries/cpp/folds.scm                  |  2 +
 queries/cpp/highlights.scm             |  2 +
 queries/cpp/locals.scm                 |  2 +
 queries/javascript/folds.scm           |  1 +
 queries/javascript/highlights.scm      |  1 +
 queries/javascript/locals.scm          |  2 +
 queries/ocaml_interface/folds.scm      |  1 +
 queries/ocaml_interface/highlights.scm |  1 +
 queries/ocaml_interface/locals.scm     |  1 +
 queries/tsx/folds.scm                  |  1 +
 queries/tsx/highlights.scm             |  1 +
 queries/tsx/locals.scm                 |  1 +
 queries/typescript/folds.scm           |  2 +
 queries/typescript/highlights.scm      |  1 +
 queries/typescript/locals.scm          |  2 +
 17 files changed, 73 insertions(+), 32 deletions(-)
 create mode 100644 queries/ocaml_interface/folds.scm
 create mode 100644 queries/ocaml_interface/highlights.scm
 create mode 100644 queries/ocaml_interface/locals.scm
 create mode 100644 queries/tsx/folds.scm
 create mode 100644 queries/tsx/highlights.scm
 create mode 100644 queries/tsx/locals.scm

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 9006308db..09ef589f8 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -67,6 +67,18 @@ Here are some global advices :
 - Examples of queries can be found in [queries/](queries/)
 - Matches in the bottom will override queries that are above of them.
 
+If your language is an extension of a language (TypeScript is an extension of JavaScript for
+example), you can include the queries from your base language by adding the following _as the first
+line of your file_.
+
+```scheme
+; inherits: lang1,(optionallang)
+```
+
+If you want to include a language for a given query, but don't want for the queries including the
+query you qre writing to include it too, you can mark the language as optional (by putting it
+between parenthesis).
+
 ### Highlights
 
 As languages differ quite a lot, here is a set of captures available to you when building a `highlights.scm` query.
diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index 9881cd2dd..a123439fd 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -8,19 +8,9 @@ local M = {}
 
 local query_cache = caching.create_buffer_cache()
 
--- Some treesitter grammars extend others.
--- We can use that to import the queries of the base language
-M.base_language_map = {
-  cpp = {'c'},
-  typescript = {'javascript'},
-  javascript = {'jsx'},
-  tsx = {'typescript', 'javascript', 'jsx'},
-  ocaml_interface = {'ocaml'},
-}
-
 M.built_in_query_groups = {'highlights', 'locals', 'textobjects', 'folds'}
 
--- Creates a function that checks whether a certain query exists
+-- Creates a function that checks whether a given query exists
 -- for a specific language.
 local function get_query_guard(query)
   return function(lang)
@@ -79,19 +69,45 @@ local function filtered_runtime_queries(lang, query_name)
   return filter_files(api.nvim_get_runtime_file(string.format('queries/%s/%s.scm', lang, query_name), true) or {})
 end
 
-local function runtime_query_exists(lang, query_name)
-  local files = api.nvim_get_runtime_file(string.format('queries/%s/%s.scm', lang, query_name), false)
-  return files and #files > 0
-end
-
-local function get_query_files(lang, query_name)
-  local query_files = {}
-
+local function get_query_files(lang, query_name, is_included)
   local lang_files = filtered_runtime_queries(lang, query_name)
-  vim.list_extend(query_files, lang_files)
+  local query_files = lang_files
 
-  for _, base_lang in ipairs(M.base_language_map[lang] or {}) do
-    local base_files = filtered_runtime_queries(base_lang, query_name)
+  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 _, lang in ipairs(vim.split(langlist, ',', true)) do
+          local is_optional = lang:match("%(.*%)")
+
+          if is_optional then
+            if not is_included then
+              table.insert(base_langs, lang:sub(2, #lang - 1))
+            end
+          else
+            table.insert(base_langs, lang)
+          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
 
@@ -99,16 +115,8 @@ local function get_query_files(lang, query_name)
 end
 
 function M.has_query_files(lang, query_name)
-  local langs = {lang}
-  vim.list_extend(langs, M.base_language_map[lang] or {})
-
-  for _, lang in ipairs(langs) do
-    if runtime_query_exists(lang, query_name) then
-      return true
-    end
-  end
-
-  return false
+  local files = get_query_files(lang, query_name)
+  return files and #files > 0
 end
 
 function M.get_query(lang, query_name)
diff --git a/queries/cpp/folds.scm b/queries/cpp/folds.scm
index f43412012..131be208a 100644
--- a/queries/cpp/folds.scm
+++ b/queries/cpp/folds.scm
@@ -1,3 +1,5 @@
+; inherits: c
+
 [
  (for_range_loop)
  (class_specifier)
diff --git a/queries/cpp/highlights.scm b/queries/cpp/highlights.scm
index 30c4b75d2..128ecb9d7 100644
--- a/queries/cpp/highlights.scm
+++ b/queries/cpp/highlights.scm
@@ -1,3 +1,5 @@
+; inherits: c
+
 ((identifier) @field
  (#match? @field "^_"))
 
diff --git a/queries/cpp/locals.scm b/queries/cpp/locals.scm
index aa25a4cd8..a109aabdf 100644
--- a/queries/cpp/locals.scm
+++ b/queries/cpp/locals.scm
@@ -1,3 +1,5 @@
+; inherits: c
+
 ;; Parameters
 (variadic_parameter_declaration
   declarator: (variadic_declarator
diff --git a/queries/javascript/folds.scm b/queries/javascript/folds.scm
index 6bc76af85..0b33ba569 100644
--- a/queries/javascript/folds.scm
+++ b/queries/javascript/folds.scm
@@ -1,3 +1,4 @@
+; inherits: (jsx)
 [
   (for_in_statement)
   (for_statement)
diff --git a/queries/javascript/highlights.scm b/queries/javascript/highlights.scm
index de505366a..0b6504968 100644
--- a/queries/javascript/highlights.scm
+++ b/queries/javascript/highlights.scm
@@ -1,3 +1,4 @@
+; inherits: (jsx)
 ; Types
 
 ; Javascript
diff --git a/queries/javascript/locals.scm b/queries/javascript/locals.scm
index 2f501f3a1..99a40a6be 100644
--- a/queries/javascript/locals.scm
+++ b/queries/javascript/locals.scm
@@ -1,3 +1,5 @@
+; inherits: (jsx)
+
 ; Scopes
 ;-------
 
diff --git a/queries/ocaml_interface/folds.scm b/queries/ocaml_interface/folds.scm
new file mode 100644
index 000000000..6d3dfbcf2
--- /dev/null
+++ b/queries/ocaml_interface/folds.scm
@@ -0,0 +1 @@
+; inherits: ocaml
diff --git a/queries/ocaml_interface/highlights.scm b/queries/ocaml_interface/highlights.scm
new file mode 100644
index 000000000..6d3dfbcf2
--- /dev/null
+++ b/queries/ocaml_interface/highlights.scm
@@ -0,0 +1 @@
+; inherits: ocaml
diff --git a/queries/ocaml_interface/locals.scm b/queries/ocaml_interface/locals.scm
new file mode 100644
index 000000000..6d3dfbcf2
--- /dev/null
+++ b/queries/ocaml_interface/locals.scm
@@ -0,0 +1 @@
+; inherits: ocaml
diff --git a/queries/tsx/folds.scm b/queries/tsx/folds.scm
new file mode 100644
index 000000000..07391231c
--- /dev/null
+++ b/queries/tsx/folds.scm
@@ -0,0 +1 @@
+; inherits: typescript,jsx
diff --git a/queries/tsx/highlights.scm b/queries/tsx/highlights.scm
new file mode 100644
index 000000000..07391231c
--- /dev/null
+++ b/queries/tsx/highlights.scm
@@ -0,0 +1 @@
+; inherits: typescript,jsx
diff --git a/queries/tsx/locals.scm b/queries/tsx/locals.scm
new file mode 100644
index 000000000..07391231c
--- /dev/null
+++ b/queries/tsx/locals.scm
@@ -0,0 +1 @@
+; inherits: typescript,jsx
diff --git a/queries/typescript/folds.scm b/queries/typescript/folds.scm
index f08395f3a..448f4d12a 100644
--- a/queries/typescript/folds.scm
+++ b/queries/typescript/folds.scm
@@ -1,3 +1,5 @@
+; inherits: javascript
+
 [
   (interface_declaration)
   (internal_module)
diff --git a/queries/typescript/highlights.scm b/queries/typescript/highlights.scm
index fa7db9d92..87272f98f 100644
--- a/queries/typescript/highlights.scm
+++ b/queries/typescript/highlights.scm
@@ -1,3 +1,4 @@
+; inherits: javascript
 [
 "abstract"
 "declare"
diff --git a/queries/typescript/locals.scm b/queries/typescript/locals.scm
index cb064b823..3e05f536e 100644
--- a/queries/typescript/locals.scm
+++ b/queries/typescript/locals.scm
@@ -1,2 +1,4 @@
+; inherits: javascript
+
 (required_parameter (identifier) @definition)
 (optional_parameter (identifier) @definition)