modules/performance: add byteCompileLua option

flake-modules/wrappers.nix

@@ -15,6 +15,13 @@
               inherit (inputs) home-manager;
               nixvim = self;
             }).activationPackage;
+          home-manager-extra-files-byte-compiling =
+            import ../tests/modules/hm-extra-files-byte-compiling.nix
+              {
+                inherit pkgs;
+                inherit (inputs) home-manager;
+                nixvim = self;
+              };
         }
         // pkgs.lib.optionalAttrs (!pkgs.stdenv.isDarwin) {
           nixos-module =

lib/builders.nix

@@ -1,5 +1,5 @@
 { lib, pkgs }:
-{
+rec {
   /*
     Write a lua file to the nix store, formatted using stylua.
 
@@ -26,4 +26,110 @@
         --indent-width 4 \
         "$out"
     '';
+
+  /*
+    Write a byte compiled lua file to the nix store.
+
+    # Type
+
+    ```
+    writeByteCompiledLua :: String -> String -> Derivation
+    ```
+
+    # Arguments
+
+    - [name] The name of the derivation
+    - [text] The content of the lua file
+  */
+  writeByteCompiledLua =
+    name: text:
+    pkgs.runCommandLocal name { inherit text; } ''
+      echo -n "$text" > "$out"
+
+      ${lib.getExe' pkgs.luajit "luajit"} -bd -- "$out" "$out"
+    '';
+
+  /*
+    Get a source lua file and write a byte compiled copy of it
+    to the nix store.
+
+    # Type
+
+    ```
+    byteCompileLuaFile :: String -> String -> Derivation
+    ```
+
+    # Arguments
+
+    - [name] The name of the derivation
+    - [src] The path to the source lua file
+  */
+  byteCompileLuaFile =
+    name: src:
+    pkgs.runCommandLocal name { inherit src; } ''
+      ${lib.getExe' pkgs.luajit "luajit"} -bd -- "$src" "$out"
+    '';
+
+  # Setup hook to byte compile all lua files in the output directory.
+  # Invalid lua files are ignored.
+  byteCompileLuaHook = pkgs.makeSetupHook { name = "byte-compile-lua-hook"; } (
+    let
+      luajit = lib.getExe' pkgs.luajit "luajit";
+    in
+    pkgs.writeText "byte-compile-lua-hook.sh" # bash
+      ''
+        byteCompileLuaPostFixup() {
+            # Target is a single file
+            if [[ -f $out ]]; then
+                if [[ $out = *.lua ]]; then
+                    tmp=$(mktemp)
+                    ${luajit} -bd -- "$out" "$tmp"
+                    mv "$tmp" "$out"
+                fi
+                return
+            fi
+
+            # Target is a directory
+            while IFS= read -r -d "" file; do
+                tmp=$(mktemp -u "$file.XXXX")
+                # Ignore invalid lua files
+                if ${luajit} -bd -- "$file" "$tmp"; then
+                    mv "$tmp" "$file"
+                else
+                    echo "WARNING: Ignoring byte compiling error for '$file' lua file" >&2
+                fi
+            done < <(find "$out" -type f,l -name "*.lua" -print0)
+        }
+
+        postFixupHooks+=(byteCompileLuaPostFixup)
+      ''
+  );
+
+  /*
+    Returns an overridden derivation with all lua files byte compiled.
+
+    # Type
+
+    ```
+    byteCompileLuaDrv :: Derivation -> Derivation
+    ```
+
+    # Arguments
+
+    - [drv] Input derivation
+  */
+  byteCompileLuaDrv =
+    drv:
+    drv.overrideAttrs (
+      prev:
+      {
+        nativeBuildInputs = prev.nativeBuildInputs or [ ] ++ [ byteCompileLuaHook ];
+      }
+      // lib.optionalAttrs (prev ? buildCommand) {
+        buildCommand = ''
+          ${prev.buildCommand}
+          runHook postFixup
+        '';
+      }
+    );
 }

modules/files.nix

@@ -2,14 +2,16 @@
   lib,
   helpers,
   pkgs,
+  config,
   ...
 }:
 let
-  fileType = lib.types.submodule (
+  fileTypeModule =
     {
       name,
       config,
       options,
+      topConfig,
       ...
     }:
     {
@@ -41,6 +43,14 @@ let
           type = lib.types.path;
           description = "Path of the source file.";
         };
+
+        finalSource = lib.mkOption {
+          type = lib.types.path;
+          description = "Path to the final source file.";
+          readOnly = true;
+          visible = false;
+          internal = true;
+        };
       };
 
       config =
@@ -54,9 +64,29 @@ let
             # This means our `source` definition has the same priority as `text`.
             lib.mkDerivedConfig options.text (pkgs.writeText derivationName)
           );
+          finalSource =
+            # Byte compile lua files if performance.byteCompileLua option is enabled
+            if
+              lib.hasSuffix ".lua" config.target
+              && topConfig.performance.byteCompileLua.enable
+              && topConfig.performance.byteCompileLua.configs
+            then
+              if lib.isDerivation config.source then
+                # Source is a derivation
+                helpers.byteCompileLuaDrv config.source
+              else
+                # Source is a path or string
+                helpers.byteCompileLuaFile derivationName config.source
+            else
+              config.source;
         };
-    }
-  );
+    };
+
+  fileType = lib.types.submoduleWith {
+    shorthandOnlyDefinesConfig = true;
+    modules = [ fileTypeModule ];
+    specialArgs.topConfig = config;
+  };
 
   # TODO: Added 2024-07-07, remove after 24.11
   # Before we had a fileType, we used types.str.

modules/performance.nix

@@ -4,6 +4,28 @@ let
 in
 {
   options.performance = {
+    byteCompileLua = {
+      enable = lib.mkEnableOption "byte compiling of lua files";
+      initLua = lib.mkOption {
+        description = "Whether to byte compile init.lua.";
+        type = types.bool;
+        default = true;
+        example = false;
+      };
+      configs = lib.mkOption {
+        description = "Whether to byte compile lua configuration files.";
+        type = types.bool;
+        default = true;
+        example = false;
+      };
+      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 = {
       enable = lib.mkEnableOption "combinePlugins" // {
         description = ''

modules/top-level/files/default.nix

@@ -84,11 +84,11 @@ in
 
         mkdir -p "$out"
         ${lib.concatMapStringsSep "\n" (
-          { target, source, ... }:
+          { target, finalSource, ... }:
           lib.escapeShellArgs [
             "makeEntry"
             # Force local source paths to be added to the store
-            "${source}"
+            "${finalSource}"
             target
           ]
         ) extraFiles}

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 =
@@ -206,7 +225,17 @@ in
         config.content
       ];
 
-      init = helpers.writeLua "init.lua" customRC;
+      textInit = helpers.writeLua "init.lua" customRC;
+      byteCompiledInit = helpers.writeByteCompiledLua "init.lua" customRC;
+      init =
+        if
+          config.type == "lua"
+          && config.performance.byteCompileLua.enable
+          && config.performance.byteCompileLua.initLua
+        then
+          byteCompiledInit
+        else
+          textInit;
 
       extraWrapperArgs = builtins.concatStringsSep " " (
         (optional (
@@ -215,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;
@@ -232,7 +282,7 @@ in
         name = "nixvim-print-init";
         runtimeInputs = [ pkgs.bat ];
         text = ''
-          bat --language=lua "${init}"
+          bat --language=lua "${textInit}"
         '';
       };
 

tests/fetch-tests.nix

@@ -14,7 +14,7 @@ let
       file = "${root}/${relativePath}/${name}";
     in
     if type == "regular" then
-      [
+      lib.optional (lib.hasSuffix ".nix" name) [
         {
           namespace = namespace ++ [ (lib.strings.removeSuffix ".nix" name) ];
           cases = import file;