From 404e56066f1b7d774f593fa780e8d38e530aa766 Mon Sep 17 00:00:00 2001
From: Stanislav Asunkin <1353637+stasjok@users.noreply.github.com>
Date: Sat, 3 May 2025 18:25:49 +0300
Subject: [PATCH] modules/performance: add ability to byte compile
 extraLuaPackages

This commit adds `performance.byteCompileLua.luaLib` options. When
enabled it byte-compiles lua packages from extraLuaPackages option.
---
 modules/performance.nix                       |  3 +
 modules/top-level/output.nix                  |  9 ++-
 .../modules/performance/byte-compile-lua.nix  | 61 ++++++++++++++++---
 3 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/modules/performance.nix b/modules/performance.nix
index d88222ba..c12c7226 100644
--- a/modules/performance.nix
+++ b/modules/performance.nix
@@ -57,6 +57,9 @@ in
       nvimRuntime = lib.mkEnableOption "nvimRuntime" // {
         description = "Whether to byte compile lua files in Nvim runtime.";
       };
+      luaLib = lib.mkEnableOption "luaLib" // {
+        description = "Whether to byte compile lua library.";
+      };
     };
 
     combinePlugins = {
diff --git a/modules/top-level/output.nix b/modules/top-level/output.nix
index 204e1cd3..8f1d1332 100644
--- a/modules/top-level/output.nix
+++ b/modules/top-level/output.nix
@@ -171,11 +171,18 @@ in
 
   config =
     let
+      # Optionally byte compile lua library
+      extraLuaPackages =
+        if config.performance.byteCompileLua.enable && config.performance.byteCompileLua.luaLib then
+          ps: map builders.byteCompileLuaDrv (config.extraLuaPackages ps)
+        else
+          config.extraLuaPackages;
+
       neovimConfig = pkgs.neovimUtils.makeNeovimConfig (
         {
+          inherit extraLuaPackages;
           inherit (config)
             extraPython3Packages
-            extraLuaPackages
             viAlias
             vimAlias
             withRuby
diff --git a/tests/test-sources/modules/performance/byte-compile-lua.nix b/tests/test-sources/modules/performance/byte-compile-lua.nix
index fa4818c0..a4b88ae0 100644
--- a/tests/test-sources/modules/performance/byte-compile-lua.nix
+++ b/tests/test-sources/modules/performance/byte-compile-lua.nix
@@ -22,16 +22,29 @@ let
         end
       end
 
+      local function test_lualib_file(func, is_compiled)
+        -- Get the source of the func
+        local info = debug.getinfo(func, "S")
+        -- The source returned by debug.getinfo is prefixed with '@'
+        local file = info.source:sub(2)
+        assert(vim.uv.fs_stat(file))
+        if is_compiled then
+          assert(is_byte_compiled(file), file .. " is expected to be byte compiled, but it's not")
+        else
+          assert(not is_byte_compiled(file), file .. " is not expected to be byte compiled, but it is")
+        end
+      end
+
       local function assert_g_var(varname, filename)
-          assert(
-              vim.g[varname] == true or vim.g[varname] == 1,
-              string.format(
-                  "expected vim.g.%s to be truthy, got %s. File %s isn't executed?",
-                  varname,
-                  vim.g[varname],
-                  filename
-              )
+        assert(
+          vim.g[varname] == true or vim.g[varname] == 1,
+          string.format(
+            "expected vim.g.%s to be truthy, got %s. File %s isn't executed?",
+            varname,
+            vim.g[varname],
+            filename
           )
+        )
       end
     '';
 
@@ -130,6 +143,8 @@ in
 
       extraPlugins = [ stubDrvPlugin ];
 
+      extraLuaPackages = ps: [ ps.say ];
+
       extraConfigLua = ''
         -- The test will search for the next string in nixvim-print-init's output: VALIDATING_STRING.
         -- Since this is the comment, it won't appear in byte compiled file.
@@ -180,6 +195,9 @@ in
 
           -- Nvim runtime isn't byte compiled by default
           test_rtp_file("lua/vim/lsp.lua", false)
+
+          -- Lua library isn't byte compiled by default
+          test_lualib_file(require("say").set, false)
         '';
 
     };
@@ -193,6 +211,8 @@ in
 
     extraPlugins = [ stubDrvPlugin ];
 
+    extraLuaPackages = ps: [ ps.say ];
+
     extraConfigLua = ''
       ${isByteCompiledFun}
 
@@ -208,6 +228,8 @@ in
       test_rtp_file("lua/stub_drv_plugin/init.lua", false)
       -- Neovim runtime
       test_rtp_file("lua/vim/lsp.lua", false)
+      -- Lua library
+      test_lualib_file(require("say").set, false)
     '';
   };
 
@@ -276,6 +298,29 @@ in
       vim.cmd.py3("import yaml")
     '';
   };
+
+  lua-lib = {
+    performance.byteCompileLua = {
+      enable = true;
+      luaLib = true;
+    };
+
+    extraLuaPackages =
+      ps: with ps; [
+        say
+        argparse
+      ];
+
+    extraConfigLuaPost = ''
+      ${isByteCompiledFun}
+
+      -- Lua modules are importable and byte compiled
+      local say = require("say")
+      test_lualib_file(say.set, true)
+      local argparse = require("argparse")
+      test_lualib_file(argparse("test").parse, true)
+    '';
+  };
 }
 //
   # Two equal tests, one with combinePlugins.enable = true