From 55ca9d235b114353af7921076aa3fb7d0acd37f5 Mon Sep 17 00:00:00 2001
From: Stanislav Asunkin <1353637+stasjok@users.noreply.github.com>
Date: Wed, 17 Jul 2024 21:06:30 +0300
Subject: [PATCH] modules/performance: add ability to byte compile lua plugins

This commit adds `performance.byteCompileLua.plugins` toggle that, if
enabled, byte compiles all lua files in plugins
---
 modules/performance.nix                       |  3 +
 modules/top-level/output.nix                  | 23 ++++++-
 .../modules/performance/byte-compile-lua.nix  | 68 +++++++++++++++++++
 3 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/modules/performance.nix b/modules/performance.nix
index 359f02dc..9714b610 100644
--- a/modules/performance.nix
+++ b/modules/performance.nix
@@ -18,6 +18,9 @@ in
         default = true;
         example = false;
       };
+      plugins = lib.mkEnableOption "plugins" // {
+        description = "Whether to byte compile lua plugins.";
+      };
     };
 
     combinePlugins = {
diff --git a/modules/top-level/output.nix b/modules/top-level/output.nix
index 35196599..4472ecfc 100644
--- a/modules/top-level/output.nix
+++ b/modules/top-level/output.nix
@@ -87,8 +87,27 @@ in
         defaultPlugin // (if p ? plugin then p else { plugin = p; });
       normalizePluginList = plugins: map normalize plugins;
 
-      # Normalized plugin list
-      normalizedPlugins = normalizePluginList config.extraPlugins;
+      # Byte compiling of normalized plugin list
+      byteCompilePlugins =
+        plugins:
+        let
+          byteCompile =
+            p:
+            (helpers.byteCompileLuaDrv p).overrideAttrs (
+              prev: lib.optionalAttrs (prev ? dependencies) { dependencies = map byteCompile prev.dependencies; }
+            );
+        in
+        map (p: p // { plugin = byteCompile p.plugin; }) plugins;
+
+      # Normalized and optionally byte compiled plugin list
+      normalizedPlugins =
+        let
+          normalized = normalizePluginList config.extraPlugins;
+        in
+        if config.performance.byteCompileLua.enable && config.performance.byteCompileLua.plugins then
+          byteCompilePlugins normalized
+        else
+          normalized;
 
       # Plugin list extended with dependencies
       allPlugins =
diff --git a/tests/test-sources/modules/performance/byte-compile-lua.nix b/tests/test-sources/modules/performance/byte-compile-lua.nix
index a0f2edd7..e906c596 100644
--- a/tests/test-sources/modules/performance/byte-compile-lua.nix
+++ b/tests/test-sources/modules/performance/byte-compile-lua.nix
@@ -183,3 +183,71 @@ in
     '';
   };
 }
+//
+  # Two equal tests, one with combinePlugins.enable = true
+  pkgs.lib.genAttrs
+    [
+      "plugins"
+      "plugins-combined"
+    ]
+    (name: {
+      performance = {
+        byteCompileLua = {
+          enable = true;
+          plugins = true;
+        };
+
+        combinePlugins.enable = pkgs.lib.hasSuffix "combined" name;
+      };
+
+      extraPlugins = with pkgs.vimPlugins; [
+        nvim-lspconfig
+        # Depends on plenary-nvim
+        telescope-nvim
+        # buildCommand plugin with python3 dependency
+        ((pkgs.writeTextDir "/plugin/test.lua" "vim.opt.tabstop = 2").overrideAttrs {
+          passthru.python3Dependencies = ps: [ ps.pyyaml ];
+        })
+        # Plugin with invalid lua file tests/indent/lua/cond.lua (should be ignored)
+        nvim-treesitter
+      ];
+
+      extraConfigLuaPost = ''
+        ${isByteCompiledFun}
+
+        -- Plugins are loadable
+        require("lspconfig")
+        require("telescope")
+        require("plenary")
+        require("nvim-treesitter")
+
+        -- Python modules are importable
+        vim.cmd.py3("import yaml")
+
+        -- nvim-lspconfig
+        test_rtp_file("lua/lspconfig.lua", true)
+        test_rtp_file("lua/lspconfig/server_configurations/nixd.lua", true)
+        test_rtp_file("plugin/lspconfig.lua", true)
+        test_rtp_file("doc/lspconfig.txt", false)
+
+        -- telescope-nvim
+        test_rtp_file("lua/telescope/init.lua", true)
+        test_rtp_file("lua/telescope/builtin/init.lua", true)
+        test_rtp_file("plugin/telescope.lua", true)
+        test_rtp_file("autoload/health/telescope.vim", false)
+        test_rtp_file("doc/telescope.txt", false)
+
+        -- Dependency of telescope-nvim (plenary-nvim)
+        test_rtp_file("lua/plenary/init.lua", true)
+        test_rtp_file("plugin/plenary.vim", false)
+
+        -- Test plugin
+        test_rtp_file("plugin/test.lua", true)
+
+        -- nvim-treesitter
+        test_rtp_file("lua/nvim-treesitter/health.lua", true)
+        test_rtp_file("lua/nvim-treesitter/install.lua", true)
+        test_rtp_file("plugin/nvim-treesitter.lua", true)
+        test_rtp_file("queries/nix/highlights.scm", false)
+      '';
+    })