PortValhalla/lib/modules/partition.nix

{ lib, config, ... }:
  let
    inherit (lib) types mkOption;
    diskVarName = "myDisk";
    diskVar = ''''${${diskVarName}}'';
    isSwap = partition: builtins.elem partition.type ["swap" 19];

    fs = {
      ext4 = "ext4";
      swap = "swap";
      ntfs = "ntfs";
      fat32 = "fat32";
    };

    diskSelector = ''
      . ${./../../scripts/Common/Scripts/choose-disk.sh};
      chooseDisk ${diskVarName} "Which disk do you wish to install the OS on?";
    '';

    mkDiskType = osDisk: types.submodule (
      { config, name, ... }: {
        options = {
          wipe = mkOption {
            type = types.bool;
            description = "A value indicating whether the disk should be wiped.";
            default = !(lib.lists.any (_: _.keepExisting) (builtins.attrValues config.partitions));
          };

          deviceName = mkOption {
            type = types.nullOr types.str;
            description = "The name of the device.";
            default = if osDisk then null else name;
          };

          devicePath = mkOption {
            type =
              if osDisk
              then types.nullOr types.str
              else types.str;
            description = "The path to the device.";
            default =
              if osDisk && config.deviceName == null
              then null
              else "/dev/${config.deviceName}";
          };

          partitions = mkOption {
            type = types.attrsOf (types.nullOr partitionType);
            description = "The partitions of the disk.";
            default = {};
          };

          script = mkOption {
            type = types.str;
            description = "The script for formatting the disk.";
            internal = true;
          };
        };

        config =
        let
          partitions = builtins.foldl'
            (list: predicate: lib.lists.sortOn predicate list)
            (builtins.filter (_: _ != null) (builtins.attrValues config.partitions))
            [
              (_: _.index)
              (_: _.order)
              (_: if ((!isSwap _) && (_.size == null)) then 1 else 0)
            ];

          mkType = type:
            lib.strings.escapeShellArg (
              if builtins.isInt type
              then "${lib.trivial.toHexString type}"
              else type);

          fdiskCommand = arguments: "sudo sfdisk ${arguments}";
          fdiskScript = script: append: "echo ${script} | ${fdiskCommand "${if append then "--append" else ""} ${diskVar}"}";
          wipeScript = script: fdiskScript script false;
          appendScript = script: fdiskScript script true;

          cleanup = lib.strings.concatLines (builtins.map
            (partition: "${fdiskCommand "--delete ${diskVar} ${toString partition.index}"} || true")
            (builtins.filter (_: !_.keepExisting) partitions));

          fdiskCommands = lib.strings.concatLines (
            lib.optionals config.wipe [
              cleanup
              (wipeScript "label: gpt")
            ] ++
            (builtins.concatMap (
              partition:
                let
                  inherit (partition)
                    format
                    index
                    keepExisting
                    label
                    sizeScript
                    type
                  ;

                  partVarName = "myPartition";
                  partVar = ''''${${partVarName}}'';

                  sizeOption = ''
                    ${sizeScript} | sed -e "s/.*[^[:space:]]/size=\0/"
                  '';

                  formatScripts = {
                    ${fs.ext4} = "mkfs.ext4 -F ${partVar}";
                    ${fs.swap} = "mkswap ${partVar}";
                    ${fs.ntfs} = "mkfs.ntfs -F ${partVar}";
                    ${fs.fat32} = "mkfs.fat -F 32 ${partVar}";
                  };

                  labelScripts = {
                    ${fs.ext4} = label: "e2label ${partVar} ${label}";
                    ${fs.swap} = label: "swaplabel ${partVar} --label ${label}";
                    ${fs.ntfs} = label: "ntfslabel ${partVar} ${label}";
                    ${fs.fat32} = label: "fatlabel ${partVar} ${label}";
                  };

                  create = lib.strings.concatLines [
                    (appendScript "${toString index}: \"$(${sizeOption})\" type=${mkType type}")
                    # Wait for partition to be detected
                    ''
                      while true;
                      do
                        [ -b ${partVar} ] && break;
                      done;
                    ''
                    "sudo ${formatScripts.${format}}"
                  ];

                  fallback = ''
                    if ! { ls "${diskVar}"?(p)${toString index} 2>&1; } > /dev/null
                    then
                      ${create}
                    fi
                  '';
                in [
                  ''local diskPath="$(find -L /dev/disk/by-diskseq -samefile ${diskVar})"''
                  ''local ${partVarName}="$diskPath-part${toString index}"''
                  (if keepExisting then fallback else create)
                  "sudo ${labelScripts.${format} label}"
                ])
              partitions)
          );

          fixType = lib.strings.concatLines (
            builtins.concatMap
              (partition: lib.optional
                (partition.keepExisting && !(builtins.isNull partition.type))
                ''sudo sfdisk --part-type ${diskVar} ${toString partition.index} ${mkType partition.type}'')
              partitions);
        in {
          script = ''
            function partition() {
              ${
                if osDisk && config.devicePath == null then ''
                  ${diskSelector}
                ''
                else
                  "local ${diskVarName}=${config.devicePath}"}

              ${if (!config.wipe) then cleanup else ""}

              ${fdiskCommands}
              ${fixType}
            }

            partition
          '';
        };
      });

    partitionType = types.submodule (
      { name, config, ... }: {
        options = {
          index = mkOption {
            type = types.int;
            description = "The index of the partition.";
          };

          order = mkOption {
            type = types.int;
            description = "The sort order of the partition creation.";
            default = 0;
          };

          label = mkOption {
            type = types.str;
            description = "The label of the partition.";
            default = name;
          };

          keepExisting = mkOption {
            type = types.bool;
            description = "A value indicating whether the partition should be left untouched if it already exists.";
            default = false;
          };

          type = mkOption {
            type = types.nullOr (types.either types.str types.int);
            description = "The type of the partition.";
            default = null;
          };

          format = mkOption {
            type = types.enum (builtins.attrValues fs);
            description = "The file system format of the partition.";
            default =
              if (isSwap config)
              then fs.swap
              else throw("Partition format not specified.");
          };

          size = mkOption {
            type = types.nullOr types.str;
            description = "The size of the partition.";
            default = null;
          };

          sizeScript = mkOption {
            type = types.str;
            description = "A script for printing the size to the console.";
            internal = true;
          };

          useSwap = mkOption {
            type = types.bool;
            description = "A value indicating whether this partition should be used as swap.";
            default = isSwap config;
          };

          mountPoint = mkOption {
            type = types.nullOr types.str;
            description = "The mountpoint of the partition.";
            default = null;
          };

          mountOptions = mkOption {
            type = types.listOf types.str;
            description = "The options to apply to the mount.";
            default = [];
          };
        };

        config = {
          sizeScript = (
            if isSwap config
            then
              ''echo "$(cat /proc/meminfo | awk -F " " '/^MemTotal/ { print $2 }' | awk '{ print int((($1 / 1024 / 1024) * 0.75) + 0.5)}')"G''
            else
              "echo ${lib.strings.escapeShellArg (toString config.size)}"
          );
        };
      }
    );
  in {
    options = {
      partition = {
        rootDir = mkOption {
          type = types.str;
          description = "The root of the installation directory to mount disks into.";
          default = "/mnt";
        };

        os = mkOption {
          type = mkDiskType true;
          description = "The partition layout of the OS disk.";
        };

        disks = mkOption {
          type = types.attrsOf (mkDiskType false);
          description = "The additional disks to format.";
          default = {};
        };

        script = mkOption {
          type = types.lines;
          description = "The script for partitioning the system's disks.";
          internal = true;
        };
      };
    };

    config = {
      partition = {
        script =
          let
            inherit (config.partition) os rootDir;
            inherit (lib.strings) normalizePath;

            partPath = part: "/dev/disk/by-label/${part.label}";
            disks = ([os] ++ (builtins.attrValues config.partition.disks));
            partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);

            mountScript = lib.strings.concatLines (
              builtins.map
                (_: builtins.concatStringsSep " " [
                  "sudo"
                  "mount"
                  "--mkdir"
                  (builtins.concatStringsSep " " (builtins.map (_: "-o ${_}") _.mountOptions))
                  (partPath _)
                  (normalizePath "/${rootDir}/${_.mountPoint}")
                ])
                (lib.lists.sortOn
                  (_: builtins.stringLength _.mountPoint)
                  (builtins.filter (_: _.mountPoint != null) partitions)));

            swapScript = lib.strings.concatLines (
              builtins.map
                (_: "sudo swapon ${partPath _}")
                (builtins.filter (_: _.useSwap) partitions));
          in lib.strings.concatLines (
            (builtins.map (_: _.script) disks) ++ [
              mountScript
              swapScript
            ]
          );

        os = {
          partitions = {
            Boot = {
              index = 1;
              type = "uefi";
              size = "+1G";
              format = fs.fat32;
              mountPoint = "/boot";
            };

            Swap = {
              index = 2;
              type = "swap";
            };

            OS = {
              index = 3;
              type = "linux";
              format = fs.ext4;
              mountPoint = "/";
            };
          };
        };
      };
    };
  }