From 5e08ed42e74e1aeb5ccec71d0e50de5fb50337f3 Mon Sep 17 00:00:00 2001
From: NotAShelf <raf@notashelf.dev>
Date: Tue, 23 Apr 2024 16:48:39 +0300
Subject: [PATCH] plugins/treesitter: allow `highlight` to be fine-grained

---
 modules/plugins/treesitter/config.nix     |  6 ++--
 modules/plugins/treesitter/treesitter.nix | 40 +++++++++++++++++++++--
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/modules/plugins/treesitter/config.nix b/modules/plugins/treesitter/config.nix
index e91e5cd..8fee4dc 100644
--- a/modules/plugins/treesitter/config.nix
+++ b/modules/plugins/treesitter/config.nix
@@ -6,7 +6,9 @@
 }: let
   inherit (lib.modules) mkIf mkMerge;
   inherit (lib.lists) optional;
+  inherit (lib.trivial) boolToString;
   inherit (lib.nvim.binds) mkSetBinding addDescriptionsToMappings;
+  inherit (lib.nvim.lua) toLuaObject;
   inherit (lib.nvim.dag) entryBefore entryAnywhere;
 
   cfg = config.vim.treesitter;
@@ -50,8 +52,8 @@ in {
           ensure_installed = {},
 
           highlight = {
-            enable = true,
-            disable = {},
+            enable = ${boolToString cfg.highlight.enable},
+            disable = ${toLuaObject cfg.highlight.disable},
           },
 
           incremental_selection = {
diff --git a/modules/plugins/treesitter/treesitter.nix b/modules/plugins/treesitter/treesitter.nix
index 373f0ee..74c1b62 100644
--- a/modules/plugins/treesitter/treesitter.nix
+++ b/modules/plugins/treesitter/treesitter.nix
@@ -3,9 +3,10 @@
   lib,
   ...
 }: let
-  inherit (lib.options) mkOption mkEnableOption;
+  inherit (lib.options) mkOption mkEnableOption literalMD;
   inherit (lib.nvim.binds) mkMappingOption;
-  inherit (lib.types) listOf package;
+  inherit (lib.types) listOf package str either;
+  inherit (lib.nvim.types) luaInline;
 
   inherit (pkgs.vimPlugins.nvim-treesitter) builtGrammars;
 in {
@@ -33,6 +34,41 @@ in {
       '';
     };
 
+    highlight = {
+      enable = mkEnableOption "highlighting with treesitter";
+      disable = mkOption {
+        type = either (luaInline (listOf str));
+        default = [];
+        example = literalMD ''
+          ```lua
+          -- Disable slow treesitter highlight for large files
+          disable = function(lang, buf)
+            local max_filesize = 1000 * 1024 -- 1MB
+            local ok, stats = pcall(vim.loop.fs_stat, vim.api.nvim_buf_get_name(buf))
+            if ok and stats and stats.size > max_filesize then
+                return true
+            end
+          end
+          ```
+        '';
+
+        description = ''
+          List of treesitter grammars to disable highlighting for.
+
+          This option can be either a list, in which case it will be
+          converted to a Lua table containing grammars to disable
+          highlighting for, or a string containing a **lua function**
+          that will be read as is.
+
+          ::: {.warning}
+          A comma will be added at the end of your function, so you
+          do not need to add it yourself. Doing so will cause in
+          syntax errors within your Neovim configuration.
+          :::
+        '';
+      };
+    };
+
     defaultGrammars = mkOption {
       internal = true;
       readOnly = true;