From 14d9ed09a396d892880a7e3b569ebd259d9f688d Mon Sep 17 00:00:00 2001
From: Gerald Pinder <gmpinder@gmail.com>
Date: Mon, 4 Nov 2024 11:26:00 -0500
Subject: [PATCH] chore: Update akmods image ref gen to handle open drivers

---
 recipe/src/module.rs | 70 +++++++++++++++++++++++++-------------------
 1 file changed, 40 insertions(+), 30 deletions(-)

diff --git a/recipe/src/module.rs b/recipe/src/module.rs
index 39b575b4..98ee8910 100644
--- a/recipe/src/module.rs
+++ b/recipe/src/module.rs
@@ -4,7 +4,7 @@ use blue_build_utils::syntax_highlighting::highlight_ser;
 use bon::Builder;
 use colored::Colorize;
 use indexmap::IndexMap;
-use log::{trace, warn};
+use log::trace;
 use miette::{bail, Result};
 use serde::{Deserialize, Serialize};
 use serde_yaml::Value;
@@ -93,10 +93,30 @@ impl<'a> ModuleRequiredFields<'a> {
 
     #[must_use]
     pub fn generate_akmods_info(&'a self, os_version: &u64) -> AkmodsInfo {
-        #[derive(Debug, Copy, Clone)]
+        #[derive(Debug, Default, Copy, Clone)]
         enum NvidiaAkmods {
-            Nvidia(bool),
-            Version(u64),
+            #[default]
+            Disabled,
+            Enabled,
+            Open,
+            Proprietary,
+        }
+
+        impl From<&Value> for NvidiaAkmods {
+            fn from(value: &Value) -> Self {
+                match value.get("nvidia") {
+                    Some(enabled) if enabled.is_bool() => match enabled.as_bool() {
+                        Some(true) => Self::Enabled,
+                        _ => Self::Disabled,
+                    },
+                    Some(driver_type) if driver_type.is_string() => match driver_type.as_str() {
+                        Some("open") => Self::Open,
+                        Some("proprietary") => Self::Proprietary,
+                        _ => Self::Disabled,
+                    },
+                    _ => Self::Disabled,
+                }
+            }
         }
 
         trace!("generate_akmods_base({self:#?}, {os_version})");
@@ -105,60 +125,50 @@ impl<'a> ModuleRequiredFields<'a> {
             .config
             .get("base")
             .map(|b| b.as_str().unwrap_or_default());
-        let nvidia = self.config.get("nvidia-version").map_or_else(
-            || {
-                self.config
-                    .get("nvidia")
-                    .map_or_else(|| NvidiaAkmods::Nvidia(false), |v| NvidiaAkmods::Nvidia(v.as_bool().unwrap_or_default()))
-            },
-            |v| {
-                warn!(
-                    "The `nvidia-version` property is deprecated as upstream images may no longer exist, replace it with `nvidia: true`"
-                );
-                NvidiaAkmods::Version(v.as_u64().unwrap_or_default())
-            },
-        );
+        let nvidia = self
+            .config
+            .get("nvidia")
+            .map_or_else(Default::default, NvidiaAkmods::from);
 
         AkmodsInfo::builder()
             .images(match (base, nvidia) {
-                (Some(b), NvidiaAkmods::Nvidia(nv)) if !b.is_empty() && nv => (
+                (Some(b), NvidiaAkmods::Enabled | NvidiaAkmods::Proprietary) if !b.is_empty() => (
                     format!("akmods:{b}-{os_version}"),
                     format!("akmods-extra:{b}-{os_version}"),
                     Some(format!("akmods-nvidia:{b}-{os_version}")),
                 ),
-                (Some(b), NvidiaAkmods::Version(nv)) if !b.is_empty() && nv > 0 => (
+                (Some(b), NvidiaAkmods::Disabled) if !b.is_empty() => (
                     format!("akmods:{b}-{os_version}"),
                     format!("akmods-extra:{b}-{os_version}"),
-                    Some(format!("akmods-nvidia:{b}-{os_version}-{nv}")),
+                    None,
                 ),
-                (Some(b), _) if !b.is_empty() => (
+                (Some(b), NvidiaAkmods::Open) if !b.is_empty() => (
                     format!("akmods:{b}-{os_version}"),
                     format!("akmods-extra:{b}-{os_version}"),
-                    None,
+                    Some(format!("akmods-nvidia-open:{b}-{os_version}")),
                 ),
-                (_, NvidiaAkmods::Nvidia(nv)) if nv => (
+                (_, NvidiaAkmods::Enabled | NvidiaAkmods::Proprietary) => (
                     format!("akmods:main-{os_version}"),
                     format!("akmods-extra:main-{os_version}"),
                     Some(format!("akmods-nvidia:main-{os_version}")),
                 ),
-                (_, NvidiaAkmods::Version(nv)) if nv > 0 => (
+                (_, NvidiaAkmods::Disabled) => (
                     format!("akmods:main-{os_version}"),
                     format!("akmods-extra:main-{os_version}"),
-                    Some(format!("akmods-nvidia:main-{os_version}-{nv}")),
+                    None,
                 ),
-                _ => (
+                (_, NvidiaAkmods::Open) => (
                     format!("akmods:main-{os_version}"),
                     format!("akmods-extra:main-{os_version}"),
-                    None,
+                    Some(format!("akmods-nvidia-open:main-{os_version}")),
                 ),
             })
             .stage_name(format!(
                 "{}{}",
                 base.unwrap_or("main"),
                 match nvidia {
-                    NvidiaAkmods::Nvidia(nv) if nv => "-nvidia".to_string(),
-                    NvidiaAkmods::Version(nv) if nv > 0 => format!("-nvidia-{nv}"),
-                    _ => String::default(),
+                    NvidiaAkmods::Disabled => String::default(),
+                    _ => "-nvidia".to_string(),
                 }
             ))
             .build()