feat(treesitter): add filetype -> lang API

Problem:

  vim.treesitter does not know how to map a specific filetype to a parser.

  This creates problems since in a few places (including in vim.treesitter itself), the filetype is incorrectly used in place of lang.

Solution:

  Add an API to enable this:

  - Add vim.treesitter.language.add() as a replacement for vim.treesitter.language.require_language().
    - Optional arguments are now passed via an opts table.
    - Also takes a filetype (or list of filetypes) so we can keep track of what filetypes are associated with which langs.
    - Deprecated vim.treesitter.language.require_language().
  - Add vim.treesitter.language.get_lang() which returns the associated lang for a given filetype.
  - Add vim.treesitter.language.register() to associate filetypes to a lang without loading the parser.

9 files changed, 141 insertions, 44 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index ae21bc47ca..8b41f2d104 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -164,6 +164,9 @@ The following new APIs or features were added.
 • |vim.treesitter.query.get_node_text()| now accepts a `metadata` option for
   writing custom directives using |vim.treesitter.query.add_directive()|.
 
+• |vim.treesitter.language.add()| as a new replacement for
+  `vim.treesitter.language.require_language`.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changes*
 
@@ -202,6 +205,8 @@ DEPRECATIONS                                                *news-deprecations*
 The following functions are now deprecated and will be removed in the next
 release.
 
+• `vim.treesitter.language.require_language()` has been deprecated in favour
+  of |vim.treesitter.language.add()|.
 
 
  vim:tw=78:ts=8:sw=2:et:ft=help:norl:
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index d4799053f7..e35e145301 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -666,6 +666,36 @@ stop({bufnr})                                          *vim.treesitter.stop()*
 ==============================================================================
 Lua module: vim.treesitter.language                  *lua-treesitter-language*
 
+add({lang}, {opts})                            *vim.treesitter.language.add()*
+    Asserts that a parser for the language {lang} is installed.
+
+    Parsers are searched in the `parser` runtime directory, or the provided
+    {path}
+
+    Parameters: ~
+      • {lang}  (string) Language the parser should parse (alphanumerical and
+                `_` only)
+      • {opts}  (table|nil) Options:
+                • filetype (string|string[]) Filetype(s) that lang can be
+                  parsed with. Note this is not strictly the same as lang
+                  since a single lang can parse multiple filetypes. Defaults
+                  to lang.
+                • path (string|nil) Optional path the parser is located at
+                • symbol_name (string|nil) Internal symbol name for the
+                  language to load
+                • silent (boolean|nil) Don't throw an error if language not
+                  found
+
+    Return: ~
+        (boolean) If the specified language is installed
+
+get_lang({filetype})                      *vim.treesitter.language.get_lang()*
+    Parameters: ~
+      • {filetype}  (string)
+
+    Return: ~
+        (string|nil)
+
 inspect_language({lang})          *vim.treesitter.language.inspect_language()*
     Inspects the provided language.
 
@@ -678,24 +708,12 @@ inspect_language({lang})          *vim.treesitter.language.inspect_language()*
     Return: ~
         (table)
 
-                                  *vim.treesitter.language.require_language()*
-require_language({lang}, {path}, {silent}, {symbol_name})
-    Asserts that a parser for the language {lang} is installed.
-
-    Parsers are searched in the `parser` runtime directory, or the provided
-    {path}
+register({lang}, {filetype})              *vim.treesitter.language.register()*
+    Register a lang to be used for a filetype (or list of filetypes).
 
     Parameters: ~
-      • {lang}         (string) Language the parser should parse
-                       (alphanumerical and `_` only)
-      • {path}         (string|nil) Optional path the parser is located at
-      • {silent}       (boolean|nil) Don't throw an error if language not
-                       found
-      • {symbol_name}  (string|nil) Internal symbol name for the language to
-                       load
-
-    Return: ~
-        (boolean) If the specified language is installed
+      • {lang}      (string) Language to register
+      • {filetype}  string|string[] Filetype(s) to associate with lang
 
 
 ==============================================================================
diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 4127198576..c524b20bc2 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -29,6 +29,8 @@ setmetatable(M, {
   end,
 })
 
+---@diagnostic disable:invisible
+
 --- Creates a new parser
 ---
 --- It is not recommended to use this; use |get_parser()| instead.
@@ -39,16 +41,15 @@ setmetatable(M, {
 ---
 ---@return LanguageTree object to use for parsing
 function M._create_parser(bufnr, lang, opts)
-  language.require_language(lang)
   if bufnr == 0 then
     bufnr = vim.api.nvim_get_current_buf()
   end
 
   vim.fn.bufload(bufnr)
 
-  local self = LanguageTree.new(bufnr, lang, opts)
+  language.add(lang, { filetype = vim.bo[bufnr].filetype })
 
-  ---@diagnostic disable:invisible
+  local self = LanguageTree.new(bufnr, lang, opts)
 
   ---@private
   local function bytes_cb(_, ...)
@@ -96,13 +97,16 @@ function M.get_parser(bufnr, lang, opts)
   if bufnr == nil or bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
+  if lang == nil then
+    local ft = vim.bo[bufnr].filetype
+    lang = language.get_lang(ft) or ft
+    -- TODO(lewis6991): we should error here and not default to ft
+    -- if not lang then
+    --   error(string.format('filetype %s of buffer %d is not associated with any lang', ft, bufnr))
+    -- end
+  end
 
-  if parsers[bufnr] == nil then
-    lang = lang or a.nvim_buf_get_option(bufnr, 'filetype')
-    parsers[bufnr] = M._create_parser(bufnr, lang, opts)
-  elseif lang and parsers[bufnr]:lang() ~= lang then
-    -- Only try to create a new parser if lang is provided
-    -- and it doesn't match the stored parser
+  if parsers[bufnr] == nil or parsers[bufnr]:lang() ~= lang then
     parsers[bufnr] = M._create_parser(bufnr, lang, opts)
   end
 
@@ -123,7 +127,7 @@ function M.get_string_parser(str, lang, opts)
     str = { str, 'string' },
     lang = { lang, 'string' },
   })
-  language.require_language(lang)
+  language.add(lang)
 
   return LanguageTree.new(str, lang, opts)
 end
diff --git a/runtime/lua/vim/treesitter/health.lua b/runtime/lua/vim/treesitter/health.lua
index c0a1eca0ce..1abcdd0b31 100644
--- a/runtime/lua/vim/treesitter/health.lua
+++ b/runtime/lua/vim/treesitter/health.lua
@@ -17,7 +17,7 @@ function M.check()
 
   for _, parser in pairs(parsers) do
     local parsername = vim.fn.fnamemodify(parser, ':t:r')
-    local is_loadable, ret = pcall(ts.language.require_language, parsername)
+    local is_loadable, ret = pcall(ts.language.add, parsername)
 
     if not is_loadable or not ret then
       health.report_error(
diff --git a/runtime/lua/vim/treesitter/language.lua b/runtime/lua/vim/treesitter/language.lua
index 8634e53b7b..8637d7d544 100644
--- a/runtime/lua/vim/treesitter/language.lua
+++ b/runtime/lua/vim/treesitter/language.lua
@@ -2,19 +2,66 @@ local a = vim.api
 
 local M = {}
 
+---@type table<string,string>
+local ft_to_lang = {}
+
+---@param filetype string
+---@return string|nil
+function M.get_lang(filetype)
+  return ft_to_lang[filetype]
+end
+
+---@deprecated
+function M.require_language(lang, path, silent, symbol_name)
+  return M.add(lang, {
+    silent = silent,
+    path = path,
+    symbol_name = symbol_name,
+  })
+end
+
+---@class treesitter.RequireLangOpts
+---@field path? string
+---@field silent? boolean
+---@field filetype? string|string[]
+---@field symbol_name? string
+
 --- Asserts that a parser for the language {lang} is installed.
 ---
 --- Parsers are searched in the `parser` runtime directory, or the provided {path}
 ---
 ---@param lang string Language the parser should parse (alphanumerical and `_` only)
----@param path (string|nil) Optional path the parser is located at
----@param silent (boolean|nil) Don't throw an error if language not found
----@param symbol_name (string|nil) Internal symbol name for the language to load
+---@param opts (table|nil) Options:
+---                        - filetype (string|string[]) Filetype(s) that lang can be parsed with.
+---                          Note this is not strictly the same as lang since a single lang can
+---                          parse multiple filetypes.
+---                          Defaults to lang.
+---                        - path (string|nil) Optional path the parser is located at
+---                        - symbol_name (string|nil) Internal symbol name for the language to load
+---                        - silent (boolean|nil) Don't throw an error if language not found
 ---@return boolean If the specified language is installed
-function M.require_language(lang, path, silent, symbol_name)
+function M.add(lang, opts)
+  ---@cast opts treesitter.RequireLangOpts
+  opts = opts or {}
+  local path = opts.path
+  local silent = opts.silent
+  local filetype = opts.filetype or lang
+  local symbol_name = opts.symbol_name
+
+  vim.validate({
+    lang = { lang, 'string' },
+    path = { path, 'string', true },
+    silent = { silent, 'boolean', true },
+    symbol_name = { symbol_name, 'string', true },
+    filetype = { filetype, { 'string', 'table' }, true },
+  })
+
+  M.register(lang, filetype or lang)
+
   if vim._ts_has_language(lang) then
     return true
   end
+
   if path == nil then
     if not (lang and lang:match('[%w_]+') == lang) then
       if silent then
@@ -35,9 +82,9 @@ function M.require_language(lang, path, silent, symbol_name)
   end
 
   if silent then
-    return pcall(function()
-      vim._ts_add_language(path, lang, symbol_name)
-    end)
+    if not pcall(vim._ts_add_language, path, lang, symbol_name) then
+      return false
+    end
   else
     vim._ts_add_language(path, lang, symbol_name)
   end
@@ -45,6 +92,27 @@ function M.require_language(lang, path, silent, symbol_name)
   return true
 end
 
+--- Register a lang to be used for a filetype (or list of filetypes).
+---@param lang string Language to register
+---@param filetype string|string[] Filetype(s) to associate with lang
+function M.register(lang, filetype)
+  vim.validate({
+    lang = { lang, 'string' },
+    filetype = { filetype, { 'string', 'table' } },
+  })
+
+  local filetypes ---@type string[]
+  if type(filetype) == 'string' then
+    filetypes = { filetype }
+  else
+    filetypes = filetype
+  end
+
+  for _, f in ipairs(filetypes) do
+    ft_to_lang[f] = lang
+  end
+end
+
 --- Inspects the provided language.
 ---
 --- Inspecting provides some useful information on the language like node names, ...
@@ -52,7 +120,7 @@ end
 ---@param lang string Language
 ---@return table
 function M.inspect_language(lang)
-  M.require_language(lang)
+  M.add(lang)
   return vim._ts_inspect_language(lang)
 end
 
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 8255c6f4fe..81ad83db2c 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -51,7 +51,7 @@ LanguageTree.__index = LanguageTree
 ---                                query per language.
 ---@return LanguageTree parser object
 function LanguageTree.new(source, lang, opts)
-  language.require_language(lang)
+  language.add(lang)
   ---@type LanguageTreeOpts
   opts = opts or {}
 
@@ -170,7 +170,7 @@ function LanguageTree:parse()
   local seen_langs = {} ---@type table<string,boolean>
 
   for lang, injection_ranges in pairs(injections_by_lang) do
-    local has_lang = language.require_language(lang, nil, true)
+    local has_lang = language.add(lang, { silent = true })
 
     -- Child language trees should just be ignored if not found, since
     -- they can depend on the text of a node. Intermediate strings
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 008e5a54d7..83910316a6 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -252,7 +252,7 @@ end)
 ---
 ---@return Query Parsed query
 function M.parse_query(lang, query)
-  language.require_language(lang)
+  language.add(lang)
   local cached = query_cache[lang][query]
   if cached then
     return cached
diff --git a/test/functional/treesitter/language_spec.lua b/test/functional/treesitter/language_spec.lua
index 51d8eb62f3..a8831b9a9b 100644
--- a/test/functional/treesitter/language_spec.lua
+++ b/test/functional/treesitter/language_spec.lua
@@ -18,22 +18,23 @@ describe('treesitter language API', function()
 
     -- actual message depends on platform
     matches("Failed to load parser for language 'borklang': uv_dlopen: .+",
-       pcall_err(exec_lua, "parser = vim.treesitter.require_language('borklang', 'borkbork.so')"))
+       pcall_err(exec_lua, "parser = vim.treesitter.add('borklang', { path = 'borkbork.so' })"))
 
     -- Should not throw an error when silent
-    eq(false, exec_lua("return vim.treesitter.require_language('borklang', nil, true)"))
-    eq(false, exec_lua("return vim.treesitter.require_language('borklang', 'borkbork.so', true)"))
+    eq(false, exec_lua("return vim.treesitter.add('borklang', { silent = true })"))
+
+    eq(false, exec_lua("return vim.treesitter.add('borklang', { path = 'borkbork.so', silent = true })"))
 
     eq(".../language.lua:0: no parser for 'borklang' language, see :help treesitter-parsers",
        pcall_err(exec_lua, "parser = vim.treesitter.inspect_language('borklang')"))
 
     matches("Failed to load parser: uv_dlsym: .+",
-       pcall_err(exec_lua, 'vim.treesitter.require_language("c", nil, false, "borklang")'))
+       pcall_err(exec_lua, 'vim.treesitter.add("c", { symbol_name = "borklang" })'))
   end)
 
   it('shows error for invalid language name', function()
     eq(".../language.lua:0: '/foo/' is not a valid language name",
-      pcall_err(exec_lua, 'vim.treesitter.require_language("/foo/", nil, false)'))
+      pcall_err(exec_lua, 'vim.treesitter.add("/foo/", nil, false)'))
   end)
 
   it('inspects language', function()
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 2d0d57fbd0..2430021617 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -182,8 +182,9 @@ void ui_refresh(void)
     local firstrun = q(1)
     local manyruns = q(100)
 
+    local factor = is_os('win') and 3 or 4
     -- First run should be at least 4x slower.
-    assert(400 * manyruns < firstrun, ('firstrun: %d ms, manyruns: %d ms'):format(firstrun / 1000, manyruns / 1000))
+    assert(factor * 100 * manyruns < firstrun, ('firstrun: %d ms, manyruns: %d ms'):format(firstrun / 1000, manyruns / 1000))
   end)
 
   it('support query and iter by capture', function()