From 8f046051105196742f94f52c3ca450fb02d8975d Mon Sep 17 00:00:00 2001
From: Pedro Alves <pta2002@pta2002.com>
Date: Mon, 1 Feb 2021 16:21:42 +0000
Subject: [PATCH] tresitter: enable

---
 plugins/default.nix              |  2 +
 plugins/languages/treesitter.nix | 80 ++++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/plugins/default.nix b/plugins/default.nix
index 6e188d09..72fe0349 100644
--- a/plugins/default.nix
+++ b/plugins/default.nix
@@ -7,6 +7,8 @@
 
     ./git/gitgutter.nix
 
+    ./languages/treesitter.nix
+
     ./nvim-lsp
   ];
 }
diff --git a/plugins/languages/treesitter.nix b/plugins/languages/treesitter.nix
index 80365ceb..b6d7b83a 100644
--- a/plugins/languages/treesitter.nix
+++ b/plugins/languages/treesitter.nix
@@ -2,18 +2,92 @@
 with lib;
 let
   cfg = programs.nixvim.plugins.treesitter;
+  helpers = import ../helpers.nix { inherit lib; };
 in
 {
   options = {
-    # TODO we need some treesitter configuration, all done in lua!
     programs.nixvim.plugins.treesitter = {
-      enable = mkEnableOption "Enable treesitter highlighting";
+      enable = mkEnableOption "Enable tree-sitter syntax highlighting";
+
+      ensureInstalled = mkOption {
+        type = with types; oneOf [ (enum [ "all" "maintained" ]) (listOf str) ];
+        default = "maintained";
+        description = "Either \"all\", \"maintained\" or a list of languages";
+      };
+
+      disabledLanguages = mkOption {
+        type = types.listOf types.str;
+        default = [];
+        description = "A list of languages to disable";
+      };
+
+      customCaptures = mkOption {
+        type = types.attrsOf types.str;
+        default = {};
+        description = "Custom capture group highlighting";
+      };
+
+      incrementalSelection = let
+        keymap = default: mkOption {
+          type = types.str;
+          default = default;
+        };
+      in mkOption {
+        type = types.nullOr (types.submodule {
+          enable = mkEnableOption = "Incremental selection based on the named nodes from the grammar";
+          keymaps = {
+            initSelection = keymap "gnn";
+            nodeIncremental = keymap "grn";
+            scopeIncremental = keymap "grc";
+            nodeDecremental = keymap "grm";
+          };
+        });
+        default = null;
+      };
+
+      indent = mkEnableOption "Enable tree-sitter based indentation";
+
+      folding = mkEnableOption "Enable tree-sitter based folding";
     };
   };
 
-  config = mkIf cfg.enable {
+  config = let
+    tsOptions = {
+      highlight = {
+        enable = cfg.enable;
+        disable = mkIf (cfg.disabledLanguages != []) cfg.disabledLanguages;
+
+        custom_captures = mkIf (cfg.customCaptures != {}) cfg.customCaptures;
+      };
+
+      incremental_selection = mkIf cfg.incrementalSelection.enable {
+        enable = true;
+        keymaps = {
+          init_selection = cfg.incrementalSelection.keymaps.initSelection;
+          node_incremental = cfg.incrementalSelection.keymaps.nodeIncremental;
+          scope_incremental = cfg.incrementalSelection.keymaps.scopeIncremental;
+          node_decremental = cfg.incrementalSelection.keymaps.nodeDecremental;
+        };
+      };
+
+      indent = mkIf cfg.indent.enable {
+        enable = true;
+      };
+
+      ensure_installed = cfg.ensureInstalled;
+    };
+  in mkIf cfg.enable {
     programs.nixvim = {
       extraPlugins = [ pkgs.vimPlugins.nvim-treesitter ];
+
+      extraConfigLua = ''
+        require('nvim-treesitter.configs').setup(${helpers.toLuaObject tsOptions})
+      '';
+
+      options = mkIf cfg.folding {
+        foldmethod = "expr";
+        foldexpr = "nvim_treesitter#foldexpr()";
+      };
     };
   }
 }