kernel-builder: add variants support (#420)

* kernel-builder: add variants support

`list-variants` now lists all supported build variants of a kernel.

The `devshell`, `testshell`, and `build` commands now accept the
`--variant` option to select a different variant than the default.

* kernel-builder: fix `build` subcommand

* Better naming

Author: danieldk

kernel-builder/src/build.rs

@@ -2,17 +2,18 @@ use std::path::PathBuf;
 
 use eyre::Result;
 
+use crate::list_variants::variants;
 use crate::nix::{Flake, Nix, NixSubcommand};
 use crate::pyproject::write_card;
 use crate::util::{check_or_infer_kernel_dir, parse_build};
 
-pub fn run_build(
+fn prepare_build(
     kernel_dir: Option<PathBuf>,
     max_jobs: Option<u32>,
     cores: Option<u32>,
     print_build_logs: bool,
-    target: &str,
-) -> Result<()> {
+    variant: Option<String>,
+) -> Result<(Flake, Option<String>, Nix)> {
     let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?;
 
     if let Ok(build) = parse_build(&kernel_dir) {
@@ -23,6 +24,16 @@ pub fn run_build(
 
     let flake = Flake::from_path(kernel_dir)?;
 
+    if let Some(ref variant) = variant {
+        let valid_variants = variants(&flake)?;
+        if !valid_variants.contains(variant) {
+            eyre::bail!(
+                "Unknown variant `{variant}`.\nValid variants are: {}",
+                valid_variants.join(", ")
+            );
+        }
+    }
+
     let mut nix = Nix::new();
     if let Some(jobs) = max_jobs {
         nix = nix.max_jobs(jobs);
@@ -32,8 +43,35 @@ pub fn run_build(
     }
     nix = nix.print_build_logs(print_build_logs);
 
+    Ok((flake, variant, nix))
+}
+
+pub fn run_build(
+    kernel_dir: Option<PathBuf>,
+    max_jobs: Option<u32>,
+    cores: Option<u32>,
+    print_build_logs: bool,
+    variant: Option<String>,
+) -> Result<()> {
+    let (flake, variant, nix) =
+        prepare_build(kernel_dir, max_jobs, cores, print_build_logs, variant)?;
+
+    nix.run(NixSubcommand::Build {
+        flake: &flake,
+        attribute: variant.map(|v| format!("redistributable.{v}")),
+    })
+}
+
+pub fn run_build_and_copy(
+    kernel_dir: Option<PathBuf>,
+    max_jobs: Option<u32>,
+    cores: Option<u32>,
+    print_build_logs: bool,
+) -> Result<()> {
+    let (flake, _, nix) = prepare_build(kernel_dir, max_jobs, cores, print_build_logs, None)?;
+
     nix.run(NixSubcommand::Run {
-        flake,
-        attribute: Some(target.to_owned()),
+        flake: &flake,
+        attribute: Some("build-and-copy".to_owned()),
     })
 }

kernel-builder/src/develop.rs

@@ -4,20 +4,44 @@ use std::path::PathBuf;
 
 use eyre::Result;
 
+use crate::list_variants::arch_variants;
 use crate::nix::{Flake, Nix, NixSubcommand};
 use crate::util::check_or_infer_kernel_dir;
 
+/// Validate a variant against the arch variants list.
+fn validate_arch_variant(flake: &Flake, variant: &str) -> Result<()> {
+    let valid_variants = arch_variants(flake)?;
+    if !valid_variants.contains(&variant.to_string()) {
+        eyre::bail!(
+            "Unknown variant `{variant}`.\nValid variants for this architecture are: {}",
+            valid_variants.join(", ")
+        );
+    }
+    Ok(())
+}
+
 /// Run a Nix development shell.
 fn run_develop(
     kernel_dir: Option<PathBuf>,
     max_jobs: Option<u32>,
     cores: Option<u32>,
     print_build_logs: bool,
-    attribute: Option<String>,
+    variant: Option<String>,
+    shell_prefix: &str,
+    default_attribute: Option<String>,
 ) -> Result<()> {
     let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?;
     let flake = Flake::from_path(kernel_dir)?;
 
+    if let Some(ref variant) = variant {
+        validate_arch_variant(&flake, variant)?;
+    }
+
+    let attribute = match variant {
+        Some(ref v) => Some(format!("{shell_prefix}.{v}")),
+        None => default_attribute,
+    };
+
     let mut nix = Nix::new();
 
     if let Some(max_jobs) = max_jobs {
@@ -30,7 +54,10 @@ fn run_develop(
 
     nix = nix.print_build_logs(print_build_logs);
 
-    nix.run(NixSubcommand::Develop { flake, attribute })
+    nix.run(NixSubcommand::Develop {
+        flake: &flake,
+        attribute,
+    })
 }
 
 /// Run a kernel development shell.
@@ -39,8 +66,17 @@ pub fn devshell(
     max_jobs: Option<u32>,
     cores: Option<u32>,
     print_build_logs: bool,
+    variant: Option<String>,
 ) -> Result<()> {
-    run_develop(kernel_dir, max_jobs, cores, print_build_logs, None)
+    run_develop(
+        kernel_dir,
+        max_jobs,
+        cores,
+        print_build_logs,
+        variant,
+        "devShells",
+        None,
+    )
 }
 
 /// Run a kernel test shell.
@@ -49,12 +85,15 @@ pub fn testshell(
     max_jobs: Option<u32>,
     cores: Option<u32>,
     print_build_logs: bool,
+    variant: Option<String>,
 ) -> Result<()> {
     run_develop(
         kernel_dir,
         max_jobs,
         cores,
         print_build_logs,
+        variant,
+        "testShells",
         Some("test".to_string()),
     )
 }

kernel-builder/src/list_variants.rs

@@ -0,0 +1,49 @@
+use std::path::PathBuf;
+
+use eyre::Result;
+
+use crate::nix::{Flake, Nix, NixSubcommand};
+use crate::util::check_or_infer_kernel_dir;
+
+/// Get the list of variants for a flake.
+pub fn variants(flake: &Flake) -> Result<Vec<String>> {
+    let nix = Nix::new().json(true);
+    let output = nix.run_capture_stdout(NixSubcommand::Eval {
+        flake,
+        attribute: "variants".to_string(),
+    })?;
+
+    let variants: Vec<String> = serde_json::from_slice(&output)?;
+
+    Ok(variants)
+}
+
+/// Get the list of architecture-specific variants for a flake.
+pub fn arch_variants(flake: &Flake) -> Result<Vec<String>> {
+    let nix = Nix::new().json(true);
+    let output = nix.run_capture_stdout(NixSubcommand::Eval {
+        flake,
+        attribute: "archVariants".to_string(),
+    })?;
+
+    let variants: Vec<String> = serde_json::from_slice(&output)?;
+
+    Ok(variants)
+}
+
+pub fn list_variants(kernel_dir: Option<PathBuf>, arch: bool) -> Result<()> {
+    let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?;
+    let flake = Flake::from_path(kernel_dir)?;
+
+    let variant_list = if arch {
+        arch_variants(&flake)?
+    } else {
+        variants(&flake)?
+    };
+
+    for variant in &variant_list {
+        println!("{}", variant);
+    }
+
+    Ok(())
+}

kernel-builder/src/main.rs

@@ -13,13 +13,16 @@ mod develop;
 use develop::{devshell, testshell};
 
 mod build;
-use build::run_build;
+use build::{run_build, run_build_and_copy};
 
 mod hf;
 
 mod init;
 use init::{run_init, InitArgs};
 
+mod list_variants;
+use list_variants::list_variants;
+
 mod upload;
 use upload::{run_upload, RepoTypeArg, UploadArgs};
 
@@ -69,6 +72,10 @@ enum Commands {
         #[arg(value_name = "KERNEL_DIR")]
         kernel_dir: Option<PathBuf>,
 
+        /// Build a specific variant.
+        #[arg(long)]
+        variant: Option<String>,
+
         #[command(flatten)]
         nix_args: NixArgs,
     },
@@ -89,6 +96,10 @@ enum Commands {
         #[arg(value_name = "KERNEL_DIR")]
         kernel_dir: Option<PathBuf>,
 
+        /// Build a specific variant.
+        #[arg(long)]
+        variant: Option<String>,
+
         #[command(flatten)]
         nix_args: NixArgs,
 
@@ -137,15 +148,33 @@ enum Commands {
         #[arg(name = "KERNEL_DIR")]
         kernel_dir: Option<PathBuf>,
 
+        /// Use a specific variant.
+        #[arg(long)]
+        variant: Option<String>,
+
         #[command(flatten)]
         nix_args: NixArgs,
     },
 
+    /// List build variants.
+    ListVariants {
+        #[arg(name = "KERNEL_DIR")]
+        kernel_dir: Option<PathBuf>,
+
+        /// Only list variants for the current architecture.
+        #[arg(long)]
+        arch: bool,
+    },
+
     /// Spawn a kernel test shell.
     Testshell {
         #[arg(name = "KERNEL_DIR")]
         kernel_dir: Option<PathBuf>,
 
+        /// Use a specific variant.
+        #[arg(long)]
+        variant: Option<String>,
+
         #[command(flatten)]
         nix_args: NixArgs,
     },
@@ -197,26 +226,28 @@ fn main() -> Result<()> {
         Commands::Upload(args) => run_upload(args),
         Commands::Build {
             kernel_dir,
+            variant,
             nix_args,
         } => run_build(
             kernel_dir,
             nix_args.max_jobs,
             nix_args.cores,
             nix_args.print_build_logs,
-            "build",
+            variant,
         ),
         Commands::BuildAndCopy {
             kernel_dir,
             nix_args,
-        } => run_build(
+            ..
+        } => run_build_and_copy(
             kernel_dir,
             nix_args.max_jobs,
             nix_args.cores,
             nix_args.print_build_logs,
-            "build-and-copy",
         ),
         Commands::BuildAndUpload {
             kernel_dir,
+            variant,
             nix_args,
             repo_id,
             branch,
@@ -228,7 +259,7 @@ fn main() -> Result<()> {
                 nix_args.max_jobs,
                 nix_args.cores,
                 nix_args.print_build_logs,
-                "build",
+                variant,
             )?;
             run_upload(UploadArgs {
                 kernel_dir,
@@ -246,21 +277,26 @@ fn main() -> Result<()> {
         } => create_pyproject(kernel_dir, target_dir, force, ops_id),
         Commands::Devshell {
             kernel_dir,
+            variant,
             nix_args,
         } => devshell(
             kernel_dir,
             nix_args.max_jobs,
             nix_args.cores,
             nix_args.print_build_logs,
+            variant,
         ),
+        Commands::ListVariants { kernel_dir, arch } => list_variants(kernel_dir, arch),
         Commands::Testshell {
             kernel_dir,
+            variant,
             nix_args,
         } => testshell(
             kernel_dir,
             nix_args.max_jobs,
             nix_args.cores,
             nix_args.print_build_logs,
+            variant,
         ),
         Commands::UpdateBuild { kernel_dir } => update_build(kernel_dir),
         Commands::Validate { kernel_dir } => {