From e48d3993e3aba661bd48a36611d490c8eb863647 Mon Sep 17 00:00:00 2001
From: Marc Jakobi <marc.jakobi@tiko.energy>
Date: Wed, 28 Aug 2024 10:18:28 +0200
Subject: [PATCH] feat: more fine-grained control for disabling highlighting

---
 README.md                          |  6 ++++
 lua/rocks_treesitter/config.lua    | 50 +++++++++++++++++++++++++-----
 lua/rocks_treesitter/highlight.lua |  5 ++-
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index a7ac3e2..24d046e 100644
--- a/README.md
+++ b/README.md
@@ -90,6 +90,10 @@ auto_highlight = [
   "toml"
 ]
 auto_install = "prompt" # true | false
+# NOTE: These are filetypes
+disable = [
+  "lhaskell"
+]
 
 [treesitter.parser_map]
 # You can add custom filetype to parser mappings.
@@ -112,6 +116,8 @@ vim.g.rocks_nvim = {
         auto_highlight = { },
         auto_install = "prompt",
         parser_map = { },
+        ---@type string[] | fun(lang: string, bufnr: integer):boolean
+        disable = { }, -- filetypes or a function
     },
 }
 ```
diff --git a/lua/rocks_treesitter/config.lua b/lua/rocks_treesitter/config.lua
index 99ea921..061ace9 100644
--- a/lua/rocks_treesitter/config.lua
+++ b/lua/rocks_treesitter/config.lua
@@ -7,6 +7,8 @@ local config = {
     auto_install = "prompt",
     ---@type "all" | table<string, boolean>
     auto_highlight = {},
+    ---@type table<string, boolean> | fun(lang: string, bufnr: integer):boolean
+    disable = {},
     ---@type table<string, string>
     parser_map = require("rocks_treesitter.ft_parser_map"),
 }
@@ -14,9 +16,22 @@ local config = {
 local api = require("rocks.api")
 
 ---@class RocksTreesitterOpts
----@field auto_highlight? string[] | "all" Parsers to automatically enable syntax highlighting for. Note that these are parser languages, not file types. Defaults to an empty list.
----@field auto_install? boolean | "prompt" Auto-install parsers in `auto_highlight` as needed? Default: 'prompt'
----@field parser_map? table<string, string> Map from filetypes to parser languages
+---
+--- Parsers to automatically enable syntax highlighting for.
+--- Note that these are parser languages, not file types.
+--- Defaults to an empty list.
+---@field auto_highlight? string[] | "all"
+---
+--- Auto-install parsers in `auto_highlight` as needed?
+--- Default: 'prompt'
+---@field auto_install? boolean | "prompt"
+---
+--- Map from filetypes to parser languages
+---@field parser_map? table<string, string>
+---
+--- File types to disable highlighting for (lua + toml config)
+--- or a function to conditionally disable highlighting (lua only)
+---@field disable? string[] | fun(lang: string, bufnr: integer):boolean
 
 local toml_config = api.get_rocks_toml()["treesitter"]
 ---@type RocksTreesitterOpts
@@ -32,13 +47,32 @@ if type(lua_config) == "table" then
     opts = vim.tbl_deep_extend("force", opts, lua_config)
 end
 
+---@generic T
+---@param lst? T[]
+---@return table<T, boolean>
+local function list_to_bool_map(lst)
+    return vim
+        .iter(lst or {})
+        ---@generic T
+        ---@param acc table<T, boolean>
+        ---@param val T
+        :fold(
+            {},
+            function(acc, val)
+                acc[val] = true
+                return acc
+            end
+        )
+end
+
 --- Map opts to configs, preserving defaults if not set
 config.auto_highlight = opts.auto_highlight == "all" and "all"
-    or vim.iter(opts.auto_highlight or {}):fold({}, function(acc, lang)
-        ---@cast lang string
-        acc[lang] = true
-        return acc
-    end)
+    ---@diagnostic disable-next-line: param-type-mismatch
+    or list_to_bool_map(opts.auto_highlight)
+config.disable = type(opts.disable) == "function" and opts.disable
+    ---@diagnostic disable-next-line: param-type-mismatch
+    or list_to_bool_map(opts.disable)
+---@diagnostic disable-next-line: assign-type-mismatch
 config.auto_install = vim.F.if_nil(opts.auto_install, config.auto_install)
 config.parser_map = vim.tbl_extend("force", config.parser_map, opts.parser_map or {})
 
diff --git a/lua/rocks_treesitter/highlight.lua b/lua/rocks_treesitter/highlight.lua
index a814024..102b4d0 100644
--- a/lua/rocks_treesitter/highlight.lua
+++ b/lua/rocks_treesitter/highlight.lua
@@ -137,7 +137,10 @@ function highlight.create_autocmd()
             local bufnr = ctx.buf
             local filetype = vim.bo[bufnr].filetype
             local lang = get_lang(filetype)
-            if config.auto_highlight == "all" or config.auto_highlight[lang] then
+            ---@type boolean | nil
+            local disable = type(config.disable) == "function" and config.disable(lang, bufnr)
+                or config.disable[filetype]
+            if not disable and config.auto_highlight == "all" or config.auto_highlight[lang] then
                 do_highlight(lang)
             end
         end,