From d2f9e011d9c6e3d765422b78039e61ddf9e69749 Mon Sep 17 00:00:00 2001
From: Quentin Boyer <quentin+dev@familleboyer.net>
Date: Mon, 19 Aug 2024 23:31:57 +0200
Subject: [PATCH] lib/neovim-plugin: Add lua configuration scoped to the plugin

This commit adds a `plugins.<name>.luaConfig` section controlling the
plugin specific configuration.

The section contains the internal `init` option, containing the plugin's
initialization code.

It also contains the public `pre` and `post` options, that allow to add
code before & after the `init` section

Finally, it contains the `final` option, being the concatenation of the
three previous options.
---
 lib/neovim-plugin.nix                     | 63 ++++++++++++++++-------
 lib/types.nix                             | 34 ++++++++++++
 lib/utils.nix                             |  6 +++
 plugins/by-name/ccc/default.nix           |  2 +-
 plugins/by-name/coq-nvim/default.nix      |  2 +-
 plugins/by-name/hydra/default.nix         |  2 +-
 plugins/by-name/mini/default.nix          |  2 +-
 plugins/by-name/rustaceanvim/default.nix  |  2 +
 plugins/by-name/telescope/default.nix     |  2 +-
 plugins/by-name/treesitter/default.nix    |  2 +-
 plugins/by-name/which-key/default.nix     |  2 +-
 plugins/by-name/yanky/default.nix         |  2 +-
 plugins/cmp/default.nix                   |  2 +-
 plugins/colorschemes/ayu.nix              |  2 +-
 plugins/colorschemes/base16/default.nix   |  2 +-
 plugins/colorschemes/onedark.nix          |  2 +-
 plugins/colorschemes/vscode.nix           |  2 +-
 plugins/none-ls/default.nix               |  2 +-
 plugins/pluginmanagers/lz-n.nix           |  2 +-
 plugins/utils/rest.nix                    |  3 +-
 tests/test-sources/plugins/lua-config.nix | 26 ++++++++++
 21 files changed, 129 insertions(+), 35 deletions(-)
 create mode 100644 tests/test-sources/plugins/lua-config.nix

diff --git a/lib/neovim-plugin.nix b/lib/neovim-plugin.nix
index fce2c0a7fe..a4015424f6 100644
--- a/lib/neovim-plugin.nix
+++ b/lib/neovim-plugin.nix
@@ -25,6 +25,11 @@
       # colorscheme
       isColorscheme ? false,
       colorscheme ? name,
+      # luaConfig
+      configLocation ? if isColorscheme then "extraConfigLuaPre" else "extraConfigLua",
+      # For some plugins it may not make sense to have a configuration attribute, as they are
+      # configured through some other mean, like global variables
+      hasConfigAttrs ? true,
       # options
       originalName ? name,
       # Can be a string, a list of strings, or a module option:
@@ -60,7 +65,14 @@
         let
           cfg = config.${namespace}.${name};
           opt = options.${namespace}.${name};
-          extraConfigNamespace = if isColorscheme then "extraConfigLuaPre" else "extraConfigLua";
+
+          setupCode = ''
+            require('${luaName}')${setup}(${
+              lib.optionalString (cfg ? settings) (helpers.toLuaObject cfg.settings)
+            })
+          '';
+
+          setLuaConfig = lib.setAttrByPath (lib.toList configLocation);
         in
         {
           meta = {
@@ -100,27 +112,40 @@
                 example = settingsExample;
               };
             }
+            // lib.optionalAttrs hasConfigAttrs {
+              luaConfig = lib.mkOption {
+                type = lib.nixvim.nixvimTypes.pluginLuaConfig;
+                default = { };
+                description = "The plugin's lua configuration";
+              };
+            }
             // extraOptions;
 
-          config = lib.mkIf cfg.enable (
-            lib.mkMerge [
-              {
-                extraPlugins = (lib.optional installPackage cfg.package) ++ extraPlugins;
-                inherit extraPackages;
-              }
-              (lib.optionalAttrs callSetup {
-                ${extraConfigNamespace} = ''
-                  require('${luaName}')${setup}(${
-                    lib.optionalString (cfg ? settings) (helpers.toLuaObject cfg.settings)
+          config =
+            assert lib.assertMsg (
+              callSetup -> configLocation != null
+            ) "When a plugin has no config attrs and has a setup function it must have a config location";
+            lib.mkIf cfg.enable (
+              lib.mkMerge (
+                [
+                  {
+                    extraPlugins = (lib.optional installPackage cfg.package) ++ extraPlugins;
+                    inherit extraPackages;
+                  }
+                  (lib.optionalAttrs (isColorscheme && (colorscheme != null)) {
+                    colorscheme = lib.mkDefault colorscheme;
                   })
-                '';
-              })
-              (lib.optionalAttrs (isColorscheme && (colorscheme != null)) {
-                colorscheme = lib.mkDefault colorscheme;
-              })
-              (extraConfig cfg)
-            ]
-          );
+                  (extraConfig cfg)
+                ]
+                ++ (lib.optionals (!hasConfigAttrs) [
+                  (lib.optionalAttrs callSetup (setLuaConfig setupCode))
+                ])
+                ++ (lib.optionals hasConfigAttrs [
+                  (lib.optionalAttrs callSetup { ${namespace}.${name}.luaConfig.content = setupCode; })
+                  (lib.optionalAttrs (configLocation != null) (setLuaConfig cfg.luaConfig.content))
+                ])
+              )
+            );
         };
     in
     {
diff --git a/lib/types.nix b/lib/types.nix
index a7d57f08ff..0b8deb6c06 100644
--- a/lib/types.nix
+++ b/lib/types.nix
@@ -113,4 +113,38 @@ rec {
         types.optionDescriptionPhrase (class: class == "noun" || class == "composite") elemType
       }";
     };
+
+  pluginLuaConfig = types.submodule (
+    { config, ... }:
+    {
+      options = {
+        pre = lib.mkOption {
+          type = types.lines;
+          default = "";
+          description = ''
+            Lua code inserted at the start of the plugin's configuration.
+            This is the same as using `lib.nixvim.utils.mkBeforeSection` when defining `content`.
+          '';
+        };
+        post = lib.mkOption {
+          type = types.lines;
+          default = "";
+          description = ''
+            Lua code inserted at the end of the plugin's configuration.
+            This is the same as using `lib.nixvim.utils.mkAfterSection` when defining `content`.
+          '';
+        };
+        content = lib.mkOption {
+          type = types.lines;
+          default = "";
+          description = "Configuration of the plugin";
+        };
+      };
+
+      config.content = lib.mkMerge [
+        (lib.nixvim.utils.mkBeforeSection config.pre)
+        (lib.nixvim.utils.mkAfterSection config.post)
+      ];
+    }
+  );
 }
diff --git a/lib/utils.nix b/lib/utils.nix
index f36b1b7067..fae7093e7a 100644
--- a/lib/utils.nix
+++ b/lib/utils.nix
@@ -77,6 +77,12 @@ rec {
     in
     lib.concatStrings (map processWord words);
 
+  /**
+    Those helpers control the lua sections split in `pre, content, post`
+  */
+  mkBeforeSection = lib.mkOrder 300;
+  mkAfterSection = lib.mkOrder 2000;
+
   /**
     Capitalize a string by making the first character uppercase.
 
diff --git a/plugins/by-name/ccc/default.nix b/plugins/by-name/ccc/default.nix
index 964d2e2b6b..68a08a58d7 100644
--- a/plugins/by-name/ccc/default.nix
+++ b/plugins/by-name/ccc/default.nix
@@ -246,7 +246,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
     # ccc requires `termguicolors` to be enabled.
     opts.termguicolors = lib.mkDefault true;
 
-    extraConfigLua = ''
+    plugins.ccc.luaConfig.content = ''
       ccc = require('ccc')
       ccc.setup(${helpers.toLuaObject cfg.settings})
     '';
diff --git a/plugins/by-name/coq-nvim/default.nix b/plugins/by-name/coq-nvim/default.nix
index 8fb064ceb1..fa9d697252 100644
--- a/plugins/by-name/coq-nvim/default.nix
+++ b/plugins/by-name/coq-nvim/default.nix
@@ -82,7 +82,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
       coq_settings = cfg.settings;
     };
 
-    extraConfigLua = "require('coq')";
+    plugins.coq-nvim.luaConfig.content = "require('coq')";
 
     plugins.lsp = {
       preConfig = ''
diff --git a/plugins/by-name/hydra/default.nix b/plugins/by-name/hydra/default.nix
index bb2f26bbf4..a3387ed87c 100644
--- a/plugins/by-name/hydra/default.nix
+++ b/plugins/by-name/hydra/default.nix
@@ -36,7 +36,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
 
   callSetup = false;
   extraConfig = cfg: {
-    extraConfigLua = ''
+    plugins.hydra.luaConfig.content = ''
       hydra = require('hydra')
 
       hydra.setup(${helpers.toLuaObject cfg.settings})
diff --git a/plugins/by-name/mini/default.nix b/plugins/by-name/mini/default.nix
index 713d2241c9..89e7e727bf 100644
--- a/plugins/by-name/mini/default.nix
+++ b/plugins/by-name/mini/default.nix
@@ -99,7 +99,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
         '';
       }
     ];
-    extraConfigLua =
+    plugins.mini.luaConfig.content =
       lib.foldlAttrs (lines: name: config: ''
         ${lines}
         require(${lib.nixvim.toLuaObject "mini.${name}"}).setup(${lib.nixvim.toLuaObject config})
diff --git a/plugins/by-name/rustaceanvim/default.nix b/plugins/by-name/rustaceanvim/default.nix
index 5bce6dad8e..6fff2b325e 100644
--- a/plugins/by-name/rustaceanvim/default.nix
+++ b/plugins/by-name/rustaceanvim/default.nix
@@ -48,6 +48,8 @@ helpers.neovim-plugin.mkNeovimPlugin {
   };
 
   callSetup = false;
+  hasConfigAttrs = false;
+  configLocation = null;
   extraConfig =
     cfg:
     mkMerge [
diff --git a/plugins/by-name/telescope/default.nix b/plugins/by-name/telescope/default.nix
index c7f6d2c2e7..366de27878 100644
--- a/plugins/by-name/telescope/default.nix
+++ b/plugins/by-name/telescope/default.nix
@@ -122,7 +122,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
       }
     ) cfg.keymaps;
 
-    extraConfigLua = ''
+    plugins.telescope.luaConfig.content = ''
       require('telescope').setup(${toLuaObject cfg.settings})
 
       local __telescopeExtensions = ${toLuaObject cfg.enabledExtensions}
diff --git a/plugins/by-name/treesitter/default.nix b/plugins/by-name/treesitter/default.nix
index 2601a92fae..7265a4c0cc 100644
--- a/plugins/by-name/treesitter/default.nix
+++ b/plugins/by-name/treesitter/default.nix
@@ -431,7 +431,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
   installPackage = false;
 
   extraConfig = cfg: {
-    extraConfigLua =
+    plugins.treesitter.luaConfig.content =
       # NOTE: Upstream state that the parser MUST be at the beginning of runtimepath.
       # Otherwise the parsers from Neovim takes precedent, which may be incompatible with some queries.
       (optionalString (cfg.settings.parser_install_dir != null) ''
diff --git a/plugins/by-name/which-key/default.nix b/plugins/by-name/which-key/default.nix
index db3ee803aa..92822b1fbf 100644
--- a/plugins/by-name/which-key/default.nix
+++ b/plugins/by-name/which-key/default.nix
@@ -596,7 +596,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
         ''
       ];
 
-      extraConfigLua = lib.optionalString opt.registrations.isDefined ''
+      plugins.which-key.luaConfig.content = lib.optionalString opt.registrations.isDefined ''
         require("which-key").register(${toLuaObject cfg.registrations})
       '';
     };
diff --git a/plugins/by-name/yanky/default.nix b/plugins/by-name/yanky/default.nix
index 0bef4f7aef..3d2d18f238 100644
--- a/plugins/by-name/yanky/default.nix
+++ b/plugins/by-name/yanky/default.nix
@@ -373,7 +373,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
       }
     ];
 
-    extraConfigLua = ''
+    plugins.yanky.luaConfig.content = ''
       do
         local utils = require('yanky.utils')
         ${optionalString config.plugins.telescope.enable "local mapping = require('yanky.telescope.mapping')"}
diff --git a/plugins/cmp/default.nix b/plugins/cmp/default.nix
index 89050c190d..06db25ea15 100644
--- a/plugins/cmp/default.nix
+++ b/plugins/cmp/default.nix
@@ -70,7 +70,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
 
   callSetup = false;
   extraConfig = cfg: {
-    extraConfigLua =
+    plugins.cmp.luaConfig.content =
       ''
         local cmp = require('cmp')
         cmp.setup(${helpers.toLuaObject cfg.settings})
diff --git a/plugins/colorschemes/ayu.nix b/plugins/colorschemes/ayu.nix
index b0e576fca7..d3a8756c3f 100644
--- a/plugins/colorschemes/ayu.nix
+++ b/plugins/colorschemes/ayu.nix
@@ -39,7 +39,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
   };
 
   extraConfig = cfg: {
-    extraConfigLuaPre = ''
+    colorschemes.ayu.luaConfig.content = ''
       local ayu = require("ayu")
       ayu.setup(${toLuaObject cfg.settings})
       ayu.colorscheme()
diff --git a/plugins/colorschemes/base16/default.nix b/plugins/colorschemes/base16/default.nix
index f7541a51da..9063decdc6 100644
--- a/plugins/colorschemes/base16/default.nix
+++ b/plugins/colorschemes/base16/default.nix
@@ -179,7 +179,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
     # `settings` can either be passed to `with_config` before calling `setup`,
     # or it can be passed as `setup`'s 2nd argument.
     # See https://github.com/RRethy/base16-nvim/blob/6ac181b5733518040a33017dde654059cd771b7c/lua/base16-colorscheme.lua#L107-L125
-    extraConfigLuaPre = ''
+    colorschemes.base16.luaConfig.content = ''
       do
         local base16 = require('${luaName}')
         base16.with_config(${toLuaObject cfg.settings})
diff --git a/plugins/colorschemes/onedark.nix b/plugins/colorschemes/onedark.nix
index aa5ac73f44..ac8f7f7203 100644
--- a/plugins/colorschemes/onedark.nix
+++ b/plugins/colorschemes/onedark.nix
@@ -34,7 +34,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
   callSetup = false;
   colorscheme = null;
   extraConfig = cfg: {
-    extraConfigLuaPre = ''
+    colorschemes.onedark.luaConfig.content = ''
       _onedark = require('onedark')
       _onedark.setup(${lib.nixvim.toLuaObject cfg.settings})
       _onedark.load()
diff --git a/plugins/colorschemes/vscode.nix b/plugins/colorschemes/vscode.nix
index 11757cbbc0..f87cf84148 100644
--- a/plugins/colorschemes/vscode.nix
+++ b/plugins/colorschemes/vscode.nix
@@ -31,7 +31,7 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
   };
 
   extraConfig = cfg: {
-    extraConfigLuaPre = ''
+    colorschemes.vscode.luaConfig.content = ''
       local _vscode = require("vscode")
       _vscode.setup(${toLuaObject cfg.settings})
       _vscode.load()
diff --git a/plugins/none-ls/default.nix b/plugins/none-ls/default.nix
index b6ed75caf2..4f55e3a77a 100644
--- a/plugins/none-ls/default.nix
+++ b/plugins/none-ls/default.nix
@@ -147,7 +147,7 @@ helpers.neovim-plugin.mkNeovimPlugin {
       ];
 
       # We only do this here because of enableLspFormat
-      extraConfigLua = ''
+      plugins.none-ls.luaConfig.content = ''
         require("null-ls").setup(${helpers.toLuaObject setupOptions})
       '';
     };
diff --git a/plugins/pluginmanagers/lz-n.nix b/plugins/pluginmanagers/lz-n.nix
index 81c55d4cc5..c573693c59 100644
--- a/plugins/pluginmanagers/lz-n.nix
+++ b/plugins/pluginmanagers/lz-n.nix
@@ -190,7 +190,7 @@ nixvim.neovim-plugin.mkNeovimPlugin {
 
   extraConfig = cfg: {
     globals.lz_n = modules.mkAliasAndWrapDefsWithPriority id options.plugins.lz-n.settings;
-    extraConfigLua = mkIf (cfg.plugins != [ ]) ''
+    plugins.lz-n.luaConfig.content = mkIf (cfg.plugins != [ ]) ''
       require('lz.n').load( ${nixvim.toLuaObject cfg.plugins})
     '';
   };
diff --git a/plugins/utils/rest.nix b/plugins/utils/rest.nix
index bb969abc2d..7456ca2a44 100644
--- a/plugins/utils/rest.nix
+++ b/plugins/utils/rest.nix
@@ -434,7 +434,8 @@ lib.nixvim.neovim-plugin.mkNeovimPlugin {
       }
     ];
 
-    extraConfigLua = lib.mkIf cfg.enableTelescope ''require("telescope").load_extension("rest")'';
+    # TODO: There may be some interactions between this & telescope, maybe requiring #2292
+    plugins.rest.luaConfig.post = lib.mkIf cfg.enableTelescope ''require("telescope").load_extension("rest")'';
 
     extraPackages = [ cfg.curlPackage ];
 
diff --git a/tests/test-sources/plugins/lua-config.nix b/tests/test-sources/plugins/lua-config.nix
new file mode 100644
index 0000000000..54bbabd7b6
--- /dev/null
+++ b/tests/test-sources/plugins/lua-config.nix
@@ -0,0 +1,26 @@
+{
+  lua-config-pre-post = {
+    extraConfigLuaPre = ''
+      list = {}
+    '';
+    plugins.cmp = {
+      enable = true;
+      luaConfig = {
+        pre = ''
+          table.insert(list, "pre")
+        '';
+        content = ''
+          table.insert(list, "init")
+        '';
+        post = ''
+          table.insert(list, "post")
+        '';
+      };
+    };
+    extraConfigLuaPost = ''
+      if not vim.deep_equal(list, {"pre", "init", "post"}) then
+        vim.print("Unexpected list: " .. vim.inspect(list))
+      end
+    '';
+  };
+}