diff --git a/modules/treesitter/config.nix b/modules/treesitter/config.nix
index d197129..92802a8 100644
--- a/modules/treesitter/config.nix
+++ b/modules/treesitter/config.nix
@@ -3,60 +3,64 @@
   lib,
   ...
 }: let
-  inherit (lib) addDescriptionsToMappings mkIf optional mkSetBinding mkMerge nvim;
+  inherit (lib.modules) mkIf mkMerge;
+  inherit (lib.lists) optional;
+  inherit (lib.nvim.binds) mkSetBinding addDescriptionsToMappings;
+  inherit (lib.nvim.dag) entryBefore entryAnywhere;
 
   cfg = config.vim.treesitter;
   usingNvimCmp = config.vim.autocomplete.enable && config.vim.autocomplete.type == "nvim-cmp";
 
   self = import ./treesitter.nix {inherit lib;};
-
   mappingDefinitions = self.options.vim.treesitter.mappings;
   mappings = addDescriptionsToMappings cfg.mappings mappingDefinitions;
 in {
   config = mkIf cfg.enable {
-    vim.startPlugins =
-      ["nvim-treesitter"]
-      ++ optional usingNvimCmp "cmp-treesitter";
+    vim = {
+      startPlugins = ["nvim-treesitter"] ++ optional usingNvimCmp "cmp-treesitter";
 
-    vim.autocomplete.sources = {"treesitter" = "[Treesitter]";};
+      autocomplete.sources = {"treesitter" = "[Treesitter]";};
 
-    # For some reason, using mkSetLuaBinding and putting the lua code does not work. It just selects the whole file.
-    # This works though, and if it ain't broke, don't fix it.
-    vim.maps.normal = mkSetBinding mappings.incrementalSelection.init ":lua require('nvim-treesitter.incremental_selection').init_selection()<CR>";
+      maps = {
+        # HACK: Using mkSetLuaBinding and putting the lua code does not work for some reason: It just selects the whole file.
+        # This works though, and if it ain't broke, don't fix it.
+        normal = mkSetBinding mappings.incrementalSelection.init ":lua require('nvim-treesitter.incremental_selection').init_selection()<CR>";
 
-    vim.maps.visualOnly = mkMerge [
-      (mkSetBinding mappings.incrementalSelection.incrementByNode ":lua require('nvim-treesitter.incremental_selection').node_incremental()<CR>")
-      (mkSetBinding mappings.incrementalSelection.incrementByScope ":lua require('nvim-treesitter.incremental_selection').scope_incremental()<CR>")
-      (mkSetBinding mappings.incrementalSelection.decrementByNode ":lua require('nvim-treesitter.incremental_selection').node_decremental()<CR>")
-    ];
+        visualOnly = mkMerge [
+          (mkSetBinding mappings.incrementalSelection.incrementByNode ":lua require('nvim-treesitter.incremental_selection').node_incremental()<CR>")
+          (mkSetBinding mappings.incrementalSelection.incrementByScope ":lua require('nvim-treesitter.incremental_selection').scope_incremental()<CR>")
+          (mkSetBinding mappings.incrementalSelection.decrementByNode ":lua require('nvim-treesitter.incremental_selection').node_decremental()<CR>")
+        ];
+      };
 
-    # For some reason treesitter highlighting does not work on start if this is set before syntax on
-    vim.configRC.treesitter-fold = mkIf cfg.fold (nvim.dag.entryBefore ["basic"] ''
-      set foldmethod=expr
-      set foldexpr=nvim_treesitter#foldexpr()
-      set nofoldenable
-    '');
+      # For some reason treesitter highlighting does not work on start if this is set before syntax on
+      configRC.treesitter-fold = mkIf cfg.fold (entryBefore ["basic"] ''
+        set foldmethod=expr
+        set foldexpr=nvim_treesitter#foldexpr()
+        set nofoldenable
+      '');
 
-    vim.luaConfigRC.treesitter = nvim.dag.entryAnywhere ''
-      require'nvim-treesitter.configs'.setup {
-        highlight = {
-          enable = true,
-          disable = {},
-        },
-
-        auto_install = false,
-        ensure_installed = {},
-
-        incremental_selection = {
-          enable = true,
-          keymaps = {
-            init_selection = false,
-            node_incremental = false,
-            scope_incremental = false,
-            node_decremental = false,
+      luaConfigRC.treesitter = entryAnywhere ''
+        require'nvim-treesitter.configs'.setup {
+          highlight = {
+            enable = true,
+            disable = {},
           },
-        },
-      }
-    '';
+
+          auto_install = false,
+          ensure_installed = {},
+
+          incremental_selection = {
+            enable = true,
+            keymaps = {
+              init_selection = false,
+              node_incremental = false,
+              scope_incremental = false,
+              node_decremental = false,
+            },
+          },
+        }
+      '';
+    };
   };
 }
diff --git a/modules/treesitter/context.nix b/modules/treesitter/context.nix
index d6859d4..8bb65b5 100644
--- a/modules/treesitter/context.nix
+++ b/modules/treesitter/context.nix
@@ -1,12 +1,16 @@
 {
-  pkgs,
   config,
   lib,
   ...
 }: let
-  inherit (lib) mkEnableOption mkOption types mkIf nvim boolToString;
+  inherit (lib.options) mkOption mkEnableOption;
+  inherit (lib.types) int bool str nullOr enum;
+  inherit (lib.modules) mkIf;
+  inherit (lib.trivial) boolToString;
+  inherit (lib.nvim.lua) nullString;
+  inherit (lib.nvim.dag) entryAnywhere;
 
-  treesitter = config.vim.treesitter;
+  inherit (config.vim) treesitter;
   cfg = treesitter.context;
 in {
   options.vim.treesitter.context = {
@@ -14,37 +18,37 @@ in {
 
     maxLines = mkOption {
       description = "How many lines the window should span. Values &lt;=0 mean no limit.";
-      type = types.int;
+      type = int;
       default = 0;
     };
 
     minWindowHeight = mkOption {
       description = "Minimum editor window height to enable context. Values &lt;= 0 mean no limit.";
-      type = types.int;
+      type = int;
       default = 0;
     };
 
     lineNumbers = mkOption {
       description = "";
-      type = types.bool;
+      type = bool;
       default = true;
     };
 
     multilineThreshold = mkOption {
       description = "Maximum number of lines to collapse for a single context line.";
-      type = types.int;
+      type = int;
       default = 20;
     };
 
     trimScope = mkOption {
       description = "Which context lines to discard if [](#opt-vim.treesitter.context.maxLines) is exceeded.";
-      type = types.enum ["inner" "outer"];
+      type = enum ["inner" "outer"];
       default = "outer";
     };
 
     mode = mkOption {
       description = "Line used to calculate context.";
-      type = types.enum ["cursor" "topline"];
+      type = enum ["cursor" "topline"];
       default = "cursor";
     };
 
@@ -54,13 +58,13 @@ in {
 
         When separator is set, the context will only show up when there are at least 2 lines above cursorline.
       '';
-      type = with types; nullOr str;
+      type = nullOr str;
       default = null;
     };
 
     zindex = mkOption {
       description = "The Z-index of the context window.";
-      type = types.int;
+      type = int;
       default = 20;
     };
   };
@@ -68,7 +72,7 @@ in {
   config = mkIf (treesitter.enable && cfg.enable) {
     vim.startPlugins = ["nvim-treesitter-context"];
 
-    vim.luaConfigRC.treesitter-context = nvim.dag.entryAnywhere ''
+    vim.luaConfigRC.treesitter-context = entryAnywhere ''
       require'treesitter-context'.setup {
         enable = true,
         max_lines = ${toString cfg.maxLines},
@@ -77,7 +81,7 @@ in {
         multiline_threshold = ${toString cfg.multilineThreshold},
         trim_scope = '${cfg.trimScope}',
         mode = '${cfg.mode}',
-        separator = ${nvim.lua.nullString cfg.separator},
+        separator = ${nullString cfg.separator},
         max_lines = ${toString cfg.zindex},
       }
     '';
diff --git a/modules/treesitter/default.nix b/modules/treesitter/default.nix
index 30d51d5..5520cfe 100644
--- a/modules/treesitter/default.nix
+++ b/modules/treesitter/default.nix
@@ -1,4 +1,4 @@
-{...}: {
+{
   imports = [
     ./treesitter.nix
     ./context.nix
diff --git a/modules/treesitter/treesitter.nix b/modules/treesitter/treesitter.nix
index 413f303..feae8d7 100644
--- a/modules/treesitter/treesitter.nix
+++ b/modules/treesitter/treesitter.nix
@@ -1,29 +1,26 @@
 {lib, ...}: let
-  inherit (lib) mkEnableOption mkMappingOption mkOption types;
+  inherit (lib.options) mkOption mkEnableOption;
+  inherit (lib.nvim.binds) mkMappingOption;
+  inherit (lib.types) listOf package;
 in {
   options.vim.treesitter = {
     enable = mkEnableOption "treesitter, also enabled automatically through language options";
-
     fold = mkEnableOption "fold with treesitter";
-
     autotagHtml = mkEnableOption "autoclose and rename html tag";
-
-    mappings = {
-      incrementalSelection = {
-        init = mkMappingOption "Init selection [treesitter]" "gnn";
-        incrementByNode = mkMappingOption "Increment selection by node [treesitter]" "grn";
-        incrementByScope = mkMappingOption "Increment selection by scope [treesitter]" "grc";
-        decrementByNode = mkMappingOption "Decrement selection by node [treesitter]" "grm";
-      };
-    };
-
     grammars = mkOption {
-      type = with types; listOf package;
+      type = listOf package;
       default = [];
       description = ''
         List of treesitter grammars to install. For supported languages
         use the `vim.language.<lang>.treesitter` option
       '';
     };
+
+    mappings.incrementalSelection = {
+      init = mkMappingOption "Init selection [treesitter]" "gnn";
+      incrementByNode = mkMappingOption "Increment selection by node [treesitter]" "grn";
+      incrementByScope = mkMappingOption "Increment selection by scope [treesitter]" "grc";
+      decrementByNode = mkMappingOption "Decrement selection by node [treesitter]" "grm";
+    };
   };
 }