[PATCH v3 6/8] gpu: nova-core: add NVKV typed encoding

[Eliot Courtney]

For struct-like GMCAPI messages encoding field by field manually is
noisy. Add some type machinery and a macro to automate encoding of
struct-like messages. The `Encodeable` trait can be implemented by any
type to say that it can be encoded into an NVKV `Encoder`. Add a simple
`nvkv_encode!` macro that works on structs and encodes each field in
order. Provide some base types, such as `Key` which statically
associates a NVKV key with some value, to avoid having to make a lot of
newtypes and implement `Encodeable` on them.

Signed-off-by: Eliot Courtney <ecourtney@xxxxxxxxxx>
---
 drivers/gpu/nova-core/gsp/nvkv.rs        |  56 ++++++-
 drivers/gpu/nova-core/gsp/nvkv/encode.rs | 241 +++++++++++++++++++++++++++++++
 2 files changed, 296 insertions(+), 1 deletion(-)

diff --git a/drivers/gpu/nova-core/gsp/nvkv.rs b/drivers/gpu/nova-core/gsp/nvkv.rs
index 10f7a16ffc23..7ac3a459a98b 100644
--- a/drivers/gpu/nova-core/gsp/nvkv.rs
+++ b/drivers/gpu/nova-core/gsp/nvkv.rs
@@ -10,8 +10,15 @@
 //! naturally maps to storing a &str with the GPU name.
 
 #![expect(unused_imports)]
+#![cfg_attr(not(CONFIG_KUNIT), expect(unused_macros))]
 
-use core::ops::Deref;
+use core::{
+    marker::PhantomData,
+    ops::{
+        Deref,
+        DerefMut, //
+    }, //
+};
 
 use kernel::{
     alloc::{
@@ -94,6 +101,53 @@ fn deref(&self) -> &Self::Target {
 /// The index of an NVKV value.
 pub(crate) type Index = Bounded<u64, 12>;
 
+/// A static association between an NVKV key `KEY_ID` and the storage of its value.
+///
+/// Use with the encoder or decoder macros `nvkv_encode!` and `nvkv_decode!` to let them know how to
+/// map the value `Key<T, KEY_ID, As>` to/from encoded data. For brevity, `As` inserts an additional
+/// conversion (`From`) to avoid having to implement [`Encodable`] for many types. For example,
+/// enums that are easily convertible to a u32 can have `As = u32` and rely on the existing encoding
+/// for u32.
+#[repr(transparent)]
+pub(crate) struct Key<T, const KEY_ID: KeyId, As = T>(T, PhantomData<As>);
+
+impl<T, const KEY_ID: KeyId, As> From<T> for Key<T, KEY_ID, As> {
+    #[inline]
+    fn from(value: T) -> Self {
+        Self(value, PhantomData)
+    }
+}
+
+impl<'a, T, const KEY_ID: KeyId, As, const N: usize> From<&'a [T; N]> for Key<&'a [T], KEY_ID, As> {
+    #[inline]
+    fn from(value: &'a [T; N]) -> Self {
+        Self(&value[..], PhantomData)
+    }
+}
+
+impl<T, const KEY_ID: KeyId, As> Deref for Key<T, KEY_ID, As> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T, const KEY_ID: KeyId, As> DerefMut for Key<T, KEY_ID, As> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<T: Default, const KEY_ID: KeyId, As> Default for Key<T, KEY_ID, As> {
+    #[inline]
+    fn default() -> Self {
+        Self(T::default(), PhantomData)
+    }
+}
+
 bitfield! {
     /// The op word that starts each NVKV operation.
     struct Op(u64) {
diff --git a/drivers/gpu/nova-core/gsp/nvkv/encode.rs b/drivers/gpu/nova-core/gsp/nvkv/encode.rs
index 8b66ddb15189..2cc0cfbad814 100644
--- a/drivers/gpu/nova-core/gsp/nvkv/encode.rs
+++ b/drivers/gpu/nova-core/gsp/nvkv/encode.rs
@@ -8,11 +8,152 @@
 use super::{
     EncodedStream,
     Index,
+    Key,
     KeyId,
     Op,
     Opcode, //
 };
 
+/// A type that can encode itself into an [`Encoder`].
+pub(crate) trait Encodable {
+    /// Encodes `self` into `encoder`.
+    fn encode(&self, encoder: &mut Encoder) -> Result;
+}
+
+/// Defines a struct together with its [`Encodable`] implementation.
+///
+/// The implementation encodes each field in declaration order. Each field type must implement
+/// [`Encodable`], which is done already for types like `Key<T, KEY_ID>`.
+///
+/// # Examples
+///
+/// ```
+/// nvkv_encode! {
+///     struct Request {
+///         id: Key<u32, 0x0001>,
+///         name: Key<&'static [u8], 0x0002>,
+///     }
+/// }
+/// ```
+macro_rules! nvkv_encode {
+    (
+        $(#[$attr:meta])*
+        $vis:vis struct $name:ident {
+            $(
+                $(#[$field_attr:meta])*
+                $field_vis:vis $field:ident : $ty:ty
+            ),* $(,)?
+        }
+    ) => {
+        $(#[$attr])*
+        $vis struct $name {
+            $(
+                $(#[$field_attr])*
+                $field_vis $field: $ty,
+            )*
+        }
+
+        impl $crate::gsp::nvkv::Encodable for $name {
+            #[inline]
+            fn encode(&self, encoder: &mut $crate::gsp::nvkv::Encoder) -> ::kernel::error::Result {
+                $( $crate::gsp::nvkv::Encodable::encode(&self.$field, encoder)?; )*
+                Ok(())
+            }
+        }
+    };
+}
+
+/// A value with a specific index that encodes under the NVKV key `KEY_ID`.
+struct IndexedKey<T, const KEY_ID: KeyId> {
+    index: Index,
+    value: T,
+}
+
+impl<T, const KEY_ID: KeyId> IndexedKey<T, KEY_ID> {
+    /// Creates a key with the given index and value.
+    fn new(index: Index, value: T) -> Self {
+        Self { index, value }
+    }
+}
+
+impl<const KEY_ID: KeyId> Encodable for IndexedKey<u32, KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_u32(KEY_ID, self.index, self.value)
+    }
+}
+
+impl<const KEY_ID: KeyId> Encodable for IndexedKey<u64, KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_u64(KEY_ID, self.index, self.value)
+    }
+}
+
+impl<const KEY_ID: KeyId> Encodable for IndexedKey<&[u8], KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_array8(KEY_ID, self.index, self.value)
+    }
+}
+
+impl<const KEY_ID: KeyId> Encodable for IndexedKey<&[u32], KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_array32(KEY_ID, self.index, self.value)
+    }
+}
+
+impl<const KEY_ID: KeyId> Encodable for IndexedKey<&[u64], KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_array64(KEY_ID, self.index, self.value)
+    }
+}
+
+impl<const N: usize, const KEY_ID: KeyId> Encodable for IndexedKey<[u8; N], KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_array8(KEY_ID, self.index, &self.value)
+    }
+}
+
+impl<const N: usize, const KEY_ID: KeyId> Encodable for IndexedKey<[u32; N], KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_array32(KEY_ID, self.index, &self.value)
+    }
+}
+
+impl<const N: usize, const KEY_ID: KeyId> Encodable for IndexedKey<[u64; N], KEY_ID> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        encoder.encode_array64(KEY_ID, self.index, &self.value)
+    }
+}
+
+impl<T, const KEY_ID: KeyId, As> Encodable for Key<T, KEY_ID, As>
+where
+    IndexedKey<As, KEY_ID>: Encodable,
+    As: From<T>,
+    T: Copy,
+{
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        IndexedKey::new(Index::new::<0>(), As::from(self.0)).encode(encoder)
+    }
+}
+
+impl<T: Encodable> Encodable for Option<T> {
+    #[inline]
+    fn encode(&self, encoder: &mut Encoder) -> Result {
+        if let Some(value) = self {
+            value.encode(encoder)?;
+        }
+        Ok(())
+    }
+}
+
 /// An encoder for an NVKV stream.
 pub(crate) struct Encoder {
     stream: EncodedStream,
@@ -207,4 +348,104 @@ fn encode_all_value_kinds() -> Result {
 
         Ok(())
     }
+
+    // Tests that encoding via the `nvkv_encode!` macro works correctly.
+    #[test]
+    fn encode_typed_struct() -> Result {
+        const U32_KEY: KeyId = 0x0001;
+        const U64_KEY: KeyId = 0x0002;
+        const NAME_KEY: KeyId = 0x0003;
+        const FIXED_KEY: KeyId = 0x0004;
+        const OPT_KEY: KeyId = 0x0005;
+        const INDEXED_KEY: KeyId = 0x0006;
+
+        nvkv_encode! {
+            struct TypedRequest {
+                a: Key<u32, U32_KEY>,
+                b: Key<u64, U64_KEY>,
+                name: Key<&'static [u8], NAME_KEY>,
+                fixed: Key<[u8; 4], FIXED_KEY>,
+                opt: Option<Key<u32, OPT_KEY>>,
+                indexed: IndexedKey<u32, INDEXED_KEY>,
+            }
+        }
+
+        const U32_VALUE: u32 = 0x89ab_cdef;
+        const U64_VALUE: u64 = 0x0123_4567_89ab_cdef;
+        const INDEXED_VALUE: u32 = 0x1234_5678;
+
+        let index5 = Index::new::<5>();
+        let request = TypedRequest {
+            a: U32_VALUE.into(),
+            b: U64_VALUE.into(),
+            name: b"name\0".into(),
+            fixed: [1u8, 2, 3, 4].into(),
+            opt: None,
+            indexed: IndexedKey::new(index5, INDEXED_VALUE),
+        };
+
+        let mut encoder = Encoder::new();
+        request.encode(&mut encoder)?;
+        let encoded = encoder.finish();
+
+        // `opt` is `None` and is not encoded.
+        assert_eq!(encoded.len(), 8);
+
+        let index0 = Index::new::<0>();
+        // `a`: IMM32 with the value in the `Op`.
+        assert_eq!(
+            encoded[0],
+            Op::zeroed()
+                .with_key(U32_KEY)
+                .with_index(index0)
+                .with_opcode(Opcode::Imm32)
+                .with_value(U32_VALUE)
+                .into_raw()
+        );
+        // `b`: a single-key SEQ64 followed by the value.
+        assert_eq!(
+            encoded[1],
+            Op::zeroed()
+                .with_key(U64_KEY)
+                .with_index(index0)
+                .with_opcode(Opcode::Seq64)
+                .with_value(1u32)
+                .into_raw()
+        );
+        assert_eq!(encoded[2], U64_VALUE);
+        // `name`: ARRAY8 of 5 bytes, zero-padded to one `u64`.
+        assert_eq!(
+            encoded[3],
+            Op::zeroed()
+                .with_key(NAME_KEY)
+                .with_index(index0)
+                .with_opcode(Opcode::Array8)
+                .with_value(5u32)
+                .into_raw()
+        );
+        assert_eq!(encoded[4], u64::from_le_bytes(*b"name\0\0\0\0"));
+        // `fixed`: ARRAY8 of 4 bytes, zero-padded to one `u64`.
+        assert_eq!(
+            encoded[5],
+            Op::zeroed()
+                .with_key(FIXED_KEY)
+                .with_index(index0)
+                .with_opcode(Opcode::Array8)
+                .with_value(4u32)
+                .into_raw()
+        );
+        assert_eq!(encoded[6], u64::from_le_bytes([1, 2, 3, 4, 0, 0, 0, 0]));
+        // `indexed`: IMM32 with a non-zero index in the `Op`.
+        assert_eq!(
+            encoded[7],
+            Op::zeroed()
+                .with_key(INDEXED_KEY)
+                .with_index(index5)
+                .with_opcode(Opcode::Imm32)
+                .with_value(INDEXED_VALUE)
+                .into_raw()
+        );
+
+        Ok(())
+    }
 }

-- 
2.55.0