From 20bf0c7e51e1e5ebf5f3754332f134a8ca0ce04c Mon Sep 17 00:00:00 2001
From: Dmitriy Pleshevskiy <dmitriy@pleshevski.ru>
Date: Fri, 8 Nov 2024 18:03:08 +0300
Subject: [PATCH] plugins/nvim-treesitter: add extraQueries

---
 modules/plugins/style/nvim-treesitter.nix | 55 ++++++++++++++++++-----
 1 file changed, 45 insertions(+), 10 deletions(-)

diff --git a/modules/plugins/style/nvim-treesitter.nix b/modules/plugins/style/nvim-treesitter.nix
index 719c722..198ad15 100644
--- a/modules/plugins/style/nvim-treesitter.nix
+++ b/modules/plugins/style/nvim-treesitter.nix
@@ -14,8 +14,34 @@ let
     if cfg.grammars == null then cfg.package.withAllGrammars
     else cfg.package.withPlugins cfg.grammars;
 
+
+  postPatchExtraGrammars = lib.concatLines (lib.flip
+    lib.mapAttrsToList
+    cfg.extraGrammars
+    (k: { src, language, ... }:
+      ''
+        if [ -d ${src}/queries ]; then
+          ln -s ${src}/queries queries/${language}
+        fi
+      ''
+    )
+  );
+
+  postPatchExtraQueries = lib.concatLines (lib.flatten
+    (lib.flip lib.mapAttrsToList cfg.extraQueries (
+      lang:
+      lib.mapAttrsToList (queryKind: queries:
+        let file = pkgs.writeText "${lang}-${queryKind}" queries; in
+        ''
+          mkdir -p queries/${lang}
+          cat ${file} >> queries/${lang}/${queryKind}.scm
+        ''
+      )
+    ))
+  );
+
   finalNvimTreeSitter =
-    if cfg.extraGrammars == { } then nvimTreeSitterWithBuiltinGrammars
+    if cfg.extraGrammars == { } && cfg.extraQueries == { } then nvimTreeSitterWithBuiltinGrammars
     else
       nvimTreeSitterWithBuiltinGrammars.overrideAttrs (oldAttrs: {
         passthru.dependencies = oldAttrs.passthru.dependencies
@@ -26,15 +52,9 @@ let
           ''
         ));
 
-        postPatch = oldAttrs.postPatch + (lib.concatLines
-          (lib.flip lib.mapAttrsToList cfg.extraGrammars (k: { src, language, ... }:
-            ''
-              if [ -d ${src}/queries ]; then
-                ln -s ${src}/queries queries/${language}
-              fi
-            ''
-          ))
-        );
+        postPatch = oldAttrs.postPatch
+          + postPatchExtraGrammars
+          + postPatchExtraQueries;
       });
 in
 {
@@ -62,6 +82,21 @@ in
       default = { };
     };
 
+    extraQueries = mkOption {
+      type = with types; attrsOf (submodule {
+        options = {
+          injections = mkOption {
+            type = str;
+            default = "";
+          };
+        };
+      });
+      default = { };
+      example = {
+        javascript.injections = ''(comment) @comment'';
+      };
+    };
+
     settings = mkOption {
       type = types.attrs;
       default = { };