feat(install): support custom queries

3 files changed, 104 insertions, 38 deletions

diff --git a/lua/nvim-treesitter/_meta/parsers.lua b/lua/nvim-treesitter/_meta/parsers.lua
index 9ec2f0b9a..5af809bb0 100644
--- a/lua/nvim-treesitter/_meta/parsers.lua
+++ b/lua/nvim-treesitter/_meta/parsers.lua
@@ -23,6 +23,9 @@ error('Cannot require a meta file')
 ---
 ---Parser repo is a local directory; overrides `url`, `revision`, and `branch`
 ---@field path? string
+---
+---Directory with queries to be installed
+---@field queries? string
 
 ---@class ParserInfo
 ---
diff --git a/lua/nvim-treesitter/install.lua b/lua/nvim-treesitter/install.lua
index c56069621..336f83717 100644
--- a/lua/nvim-treesitter/install.lua
+++ b/lua/nvim-treesitter/install.lua
@@ -14,6 +14,9 @@ local uv_copyfile = a.awrap(4, uv.fs_copyfile)
 ---@type fun(path: string, mode: integer): string?
 local uv_mkdir = a.awrap(3, uv.fs_mkdir)
 
+---@type fun(path: string): string?
+local uv_rmdir = a.awrap(2, uv.fs_rmdir)
+
 ---@type fun(path: string, new_path: string): string?
 local uv_rename = a.awrap(3, uv.fs_rename)
 
@@ -23,6 +26,36 @@ local uv_symlink = a.awrap(4, uv.fs_symlink)
 ---@type fun(path: string): string?
 local uv_unlink = a.awrap(2, uv.fs_unlink)
 
+---@async
+---@param path string
+---@return string? err
+local function mkpath(path)
+  local parent = fs.dirname(path)
+  if not parent:match('^[./]$') and not uv.fs_stat(parent) then
+    mkpath(parent)
+  end
+
+  return uv_mkdir(path, 493) -- tonumber('755', 8)
+end
+
+---@async
+---@param path string
+local function rmdir(path)
+  local stat = uv.fs_lstat(path)
+  if not stat then
+    return
+  end
+
+  if stat.type == 'directory' then
+    for file in fs.dir(path) do
+      rmdir(fs.joinpath(path, file))
+    end
+    return uv_rmdir(path)
+  else
+    return uv_unlink(path)
+  end
+end
+
 local MAX_JOBS = 100
 local INSTALL_TIMEOUT = 60000
 
@@ -95,18 +128,6 @@ local function download_file(url, output)
   end
 end
 
----@async
----@param path string
----@return string? err
-local function mkpath(path)
-  local parent = fs.dirname(path)
-  if not parent:match('^[./]$') and not uv.fs_stat(parent) then
-    mkpath(parent)
-  end
-
-  return uv_mkdir(path, 493) -- tonumber('755', 8)
-end
-
 local M = {}
 
 ---
@@ -199,7 +220,8 @@ local function do_download(logger, url, project_name, cache_dir, revision, outpu
 
   local tmp = output_dir .. '-tmp'
 
-  util.delete(tmp)
+  rmdir(tmp)
+  a.schedule()
 
   url = url:gsub('.git$', '')
   local target = is_gitlab
@@ -258,7 +280,8 @@ local function do_download(logger, url, project_name, cache_dir, revision, outpu
     end
   end
 
-  util.delete(tmp)
+  rmdir(tmp)
+  a.schedule()
 end
 
 ---@async
@@ -301,6 +324,38 @@ local function do_install(logger, compile_location, target_location)
 end
 
 ---@async
+---@param logger Logger
+---@param query_src string
+---@param query_dir string
+---@return string? err
+local function do_link_queries(logger, query_src, query_dir)
+  uv_unlink(query_dir)
+  local err = uv_symlink(query_src, query_dir, { dir = true, junction = true })
+  a.schedule()
+  if err then
+    return logger:error(err)
+  end
+end
+
+---@async
+---@param logger Logger
+---@param query_src string
+---@param query_dir string
+---@return string? err
+local function do_copy_queries(logger, query_src, query_dir)
+  rmdir(query_dir)
+  local err = uv_mkdir(query_dir, 493) -- tonumber('755', 8)
+
+  for f in fs.dir(query_src) do
+    err = uv_copyfile(fs.joinpath(query_src, f), fs.joinpath(query_dir, f))
+  end
+  a.schedule()
+  if err then
+    return logger:error(err)
+  end
+end
+
+---@async
 ---@param lang string
 ---@param cache_dir string
 ---@param install_dir string
@@ -310,9 +365,8 @@ local function try_install_lang(lang, cache_dir, install_dir, generate)
   local logger = log.new('install/' .. lang)
 
   local repo = get_parser_install_info(lang)
+  local project_name = 'tree-sitter-' .. lang
   if repo then
-    local project_name = 'tree-sitter-' .. lang
-
     local revision = repo.revision
 
     local compile_location ---@type string
@@ -320,7 +374,7 @@ local function try_install_lang(lang, cache_dir, install_dir, generate)
       compile_location = fs.normalize(repo.path)
     else
       local project_dir = fs.joinpath(cache_dir, project_name)
-      util.delete(project_dir)
+      rmdir(project_dir)
 
       revision = revision or repo.branch or 'main'
 
@@ -362,26 +416,37 @@ local function try_install_lang(lang, cache_dir, install_dir, generate)
       local revfile = fs.joinpath(config.get_install_dir('parser-info') or '', lang .. '.revision')
       util.write_file(revfile, revision or '')
     end
-
-    if not repo.path then
-      util.delete(fs.joinpath(cache_dir, project_name))
-    end
   end
 
   do -- install queries
-    local queries_src = M.get_package_path('runtime', 'queries', lang)
-    if uv.fs_stat(queries_src) then
-      local queries = fs.joinpath(config.get_install_dir('queries'), lang)
+    local query_src = M.get_package_path('runtime', 'queries', lang)
+    local query_dir = fs.joinpath(config.get_install_dir('queries'), lang)
+    local task ---@type function
 
-      uv_unlink(queries)
-      local err = uv_symlink(queries_src, queries, { dir = true, junction = true })
-      a.schedule()
+    if repo and repo.queries and repo.path then -- link queries from local repo
+      query_src = fs.joinpath(fs.normalize(repo.path), repo.queries)
+      task = do_link_queries
+    elseif repo and repo.queries then -- copy queries from tarball
+      query_src = fs.joinpath(cache_dir, project_name, repo.queries)
+      task = do_copy_queries
+    elseif uv.fs_stat(query_src) then -- link queries from runtime
+      task = do_link_queries
+    end
+
+    if task then
+      local err = task(logger, query_src, query_dir)
       if err then
-        return logger:error(err)
+        return err
       end
     end
   end
 
+  -- clean up
+  if repo and not repo.path then
+    rmdir(fs.joinpath(cache_dir, project_name))
+    a.schedule()
+  end
+
   logger:info('Language installed')
 end
 
@@ -520,17 +585,21 @@ local function uninstall_lang(logger, lang, parser, queries)
     logger:debug('Unlinking ' .. parser)
     local perr = uv_unlink(parser)
     a.schedule()
-
     if perr then
       return logger:error(perr)
     end
   end
 
-  if fn.isdirectory(queries) == 1 then
+  local stat = uv.fs_lstat(queries)
+  if stat then
     logger:debug('Unlinking ' .. queries)
-    local qerr = uv_unlink(queries)
+    local qerr ---@type string?
+    if stat.type == 'link' then
+      qerr = uv_unlink(queries)
+    else
+      qerr = rmdir(queries)
+    end
     a.schedule()
-
     if qerr then
       return logger:error(qerr)
     end
diff --git a/lua/nvim-treesitter/util.lua b/lua/nvim-treesitter/util.lua
index 8a9779617..f2ff2f77f 100644
--- a/lua/nvim-treesitter/util.lua
+++ b/lua/nvim-treesitter/util.lua
@@ -17,10 +17,4 @@ function M.write_file(filename, content)
   file:close()
 end
 
---- Recursively delete a directory
---- @param name string
-function M.delete(name)
-  vim.fs.rm(name, { recursive = true, force = true })
-end
-
 return M