From ea63a1ca6ac3af999f1c37f123d7e8dcb03d10ad 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

---
 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 f56555fc9d..f24a9a43e6 100644
--- a/modules/performance.nix
+++ b/modules/performance.nix
@@ -44,6 +44,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.";
+      };
     };
   };
 
diff --git a/modules/top-level/output.nix b/modules/top-level/output.nix
index 5352770d5f..98825d4377 100644
--- a/modules/top-level/output.nix
+++ b/modules/top-level/output.nix
@@ -261,7 +261,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 af3a60d646..fc508a5184 100644
--- a/tests/test-sources/modules/performance/byte-compile-lua.nix
+++ b/tests/test-sources/modules/performance/byte-compile-lua.nix
@@ -140,6 +140,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