From 0ba8d21272415d332dcfc92134cdc119c7986a9e Mon Sep 17 00:00:00 2001
From: Stanislav Asunkin <1353637+stasjok@users.noreply.github.com>
Date: Wed, 17 Jul 2024 22:26:06 +0300
Subject: [PATCH] modules/performance: add ability to byte compile nvim runtime
 directory

This commit adds `performance.byteCompileLua.nvimRuntime` toggle that,
if enabled, byte compiles all lua files in Nvim runtime directory.
---
 modules/performance.nix                       |  3 ++
 modules/top-level/output.nix                  | 23 ++++++++++++-
 .../modules/performance/byte-compile-lua.nix  | 34 +++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/modules/performance.nix b/modules/performance.nix
index 9714b6101c..2d08615866 100644
--- a/modules/performance.nix
+++ b/modules/performance.nix
@@ -21,6 +21,9 @@ in
       plugins = lib.mkEnableOption "plugins" // {
         description = "Whether to byte compile lua plugins.";
       };
+      nvimRuntime = lib.mkEnableOption "nvimRuntime" // {
+        description = "Whether to byte compile lua files in Nvim runtime.";
+      };
     };
 
     combinePlugins = {
diff --git a/modules/top-level/output.nix b/modules/top-level/output.nix
index 4472ecfc87..d828cd080f 100644
--- a/modules/top-level/output.nix
+++ b/modules/top-level/output.nix
@@ -244,7 +244,28 @@ in
         ++ (optional config.wrapRc ''--add-flags -u --add-flags "${init}"'')
       );
 
-      wrappedNeovim = pkgs.wrapNeovimUnstable config.package (
+      package =
+        if config.performance.byteCompileLua.enable && config.performance.byteCompileLua.nvimRuntime then
+          # Using symlinkJoin to avoid rebuilding neovim
+          pkgs.symlinkJoin {
+            name = "neovim-byte-compiled-${lib.getVersion config.package}";
+            paths = [ config.package ];
+            # Required attributes from original neovim package
+            inherit (config.package) lua;
+            nativeBuildInputs = [ helpers.byteCompileLuaHook ];
+            postBuild = ''
+              # Replace Nvim's binary symlink with a regular file,
+              # or Nvim will use original runtime directory
+              rm $out/bin/nvim
+              cp ${config.package}/bin/nvim $out/bin/nvim
+
+              runHook postFixup
+            '';
+          }
+        else
+          config.package;
+
+      wrappedNeovim = pkgs.wrapNeovimUnstable package (
         neovimConfig
         // {
           wrapperArgs = lib.escapeShellArgs neovimConfig.wrapperArgs + " " + extraWrapperArgs;
diff --git a/tests/test-sources/modules/performance/byte-compile-lua.nix b/tests/test-sources/modules/performance/byte-compile-lua.nix
index e906c59623..3ed2945047 100644
--- a/tests/test-sources/modules/performance/byte-compile-lua.nix
+++ b/tests/test-sources/modules/performance/byte-compile-lua.nix
@@ -182,6 +182,40 @@ in
       test_rtp_file("plugin/test2.lua", false)
     '';
   };
+
+  nvim-runtime = {
+    performance.byteCompileLua = {
+      enable = true;
+      nvimRuntime = true;
+    };
+
+    extraPlugins = [
+      # Python 3 dependencies
+      (pkgs.vimPlugins.nvim-lspconfig.overrideAttrs { passthru.python3Dependencies = ps: [ ps.pyyaml ]; })
+    ];
+
+    extraConfigLuaPost = ''
+      ${isByteCompiledFun}
+
+      -- vim namespace is working
+      vim.opt.tabstop = 2
+      vim.api.nvim_get_runtime_file("init.lua", false)
+      vim.lsp.get_clients()
+      vim.treesitter.language.get_filetypes("nix")
+      vim.iter({})
+
+      test_rtp_file("lua/vim/lsp.lua", true)
+      test_rtp_file("lua/vim/iter.lua", true)
+      test_rtp_file("lua/vim/treesitter/query.lua", true)
+      test_rtp_file("lua/vim/lsp/buf.lua", true)
+      test_rtp_file("plugin/editorconfig.lua", true)
+      test_rtp_file("plugin/tutor.vim", false)
+      test_rtp_file("ftplugin/vim.vim", false)
+
+      -- Python3 packages are importable
+      vim.cmd.py3("import yaml")
+    '';
+  };
 }
 //
   # Two equal tests, one with combinePlugins.enable = true