From 9c8bee9da6070ea8b1a95fc5133ed1d88b9b959a Mon Sep 17 00:00:00 2001
From: traxys <quentin+dev@familleboyer.net>
Date: Fri, 21 Apr 2023 20:29:33 +0200
Subject: [PATCH] plugins/treesitter: Add injections to highlight lua in nixvim
 options (#344)

---
 plugins/languages/treesitter/treesitter.nix    | 18 ++++++++++++++++++
 .../languages/treesitter/treesitter.nix        |  7 +++++++
 2 files changed, 25 insertions(+)

diff --git a/plugins/languages/treesitter/treesitter.nix b/plugins/languages/treesitter/treesitter.nix
index 1b7988dd..37e5559a 100644
--- a/plugins/languages/treesitter/treesitter.nix
+++ b/plugins/languages/treesitter/treesitter.nix
@@ -91,6 +91,10 @@ in {
         default = {};
         description = "This is the configuration for extra modules. It should not be used directly";
       };
+
+      nixvimInjections =
+        mkEnableOption
+        "nixvim specific injections, like lua highlighting in extraConfigLua";
     };
   };
 
@@ -148,6 +152,20 @@ in {
           require('nvim-treesitter.configs').setup(${helpers.toLuaObject tsOptions})
         '';
 
+      extraFiles = mkIf cfg.nixvimInjections {
+        "queries/nix/injections.scm" = ''
+          ;; extends
+
+          (binding
+            attrpath: (attrpath (identifier) @_path)
+            expression: [
+              (string_expression (string_fragment) @lua)
+              (indented_string_expression (string_fragment) @lua)
+            ]
+            (#match? @_path "^extraConfigLua(Pre|Post)?$"))
+        '';
+      };
+
       extraPlugins = with pkgs;
         if cfg.nixGrammars
         then [(cfg.package.withPlugins (_: cfg.grammarPackages))]
diff --git a/tests/test-sources/plugins/languages/treesitter/treesitter.nix b/tests/test-sources/plugins/languages/treesitter/treesitter.nix
index 2b816e9e..24d10a51 100644
--- a/tests/test-sources/plugins/languages/treesitter/treesitter.nix
+++ b/tests/test-sources/plugins/languages/treesitter/treesitter.nix
@@ -12,6 +12,13 @@
     };
   };
 
+  nixvimInjections = {
+    plugins.treesitter = {
+      enable = true;
+      nixvimInjections = true;
+    };
+  };
+
   # This needs a custom input
   # custom = {
   #   plugins.treesitter = {