spl-discriminator generic arg support (#4601)

`spl discriminator` lifetimes & generics

Author: Joe C

Cargo.lock

@@ -7,7 +7,11 @@ repository = "https://github.com/solana-labs/solana-program-library"
 license = "Apache-2.0"
 edition = "2021"
 
+[features]
+borsh = ["dep:borsh"]
+
 [dependencies]
+borsh = { version = "0.10", optional = true }
 bytemuck = { version = "1.13.1", features = ["derive"] }
 solana-program = "1.16.1"
 spl-discriminator-derive = { version = "0.1.0", path = "./derive" }

libraries/discriminator/Cargo.toml

@@ -14,6 +14,10 @@ pub trait SplDiscriminate {
 }
 
 /// Array Discriminator type
+#[cfg_attr(
+    feature = "borsh",
+    derive(borsh::BorshSerialize, borsh::BorshDeserialize)
+)]
 #[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
 #[repr(transparent)]
 pub struct ArrayDiscriminator([u8; ArrayDiscriminator::LENGTH]);
@@ -66,3 +70,13 @@ impl TryFrom<&[u8]> for ArrayDiscriminator {
             .map_err(|_| ProgramError::InvalidAccountData)
     }
 }
+impl From<ArrayDiscriminator> for [u8; 8] {
+    fn from(from: ArrayDiscriminator) -> Self {
+        from.0
+    }
+}
+impl From<ArrayDiscriminator> for u64 {
+    fn from(from: ArrayDiscriminator) -> Self {
+        u64::from_le_bytes(from.0)
+    }
+}

libraries/discriminator/src/discriminator.rs

@@ -21,34 +21,80 @@ mod tests {
 
     #[allow(dead_code)]
     #[derive(SplDiscriminate)]
-    #[discriminator_hash_input("some_discriminator_hash_input")]
+    #[discriminator_hash_input("my_first_instruction")]
     pub struct MyInstruction1 {
         arg1: String,
         arg2: u8,
     }
 
     #[allow(dead_code)]
     #[derive(SplDiscriminate)]
-    #[discriminator_hash_input("yet_another_discriminator_hash_input")]
-    pub struct MyInstruction2 {
-        arg1: u64,
+    #[discriminator_hash_input("global:my_second_instruction")]
+    pub enum MyInstruction2 {
+        One,
+        Two,
+        Three,
     }
 
     #[allow(dead_code)]
     #[derive(SplDiscriminate)]
-    #[discriminator_hash_input("global:my_instruction_3")]
-    pub enum MyInstruction3 {
-        One,
-        Two,
-        Three,
+    #[discriminator_hash_input("global:my_instruction_with_lifetime")]
+    pub struct MyInstruction3<'a> {
+        data: &'a [u8],
+    }
+
+    #[allow(dead_code)]
+    #[derive(SplDiscriminate)]
+    #[discriminator_hash_input("global:my_instruction_with_one_generic")]
+    pub struct MyInstruction4<T> {
+        data: T,
+    }
+
+    #[allow(dead_code)]
+    #[derive(SplDiscriminate)]
+    #[discriminator_hash_input("global:my_instruction_with_one_generic_and_lifetime")]
+    pub struct MyInstruction5<'b, T> {
+        data: &'b [T],
+    }
+
+    #[allow(dead_code)]
+    #[derive(SplDiscriminate)]
+    #[discriminator_hash_input("global:my_instruction_with_multiple_generics_and_lifetime")]
+    pub struct MyInstruction6<'c, U, V> {
+        data1: &'c [U],
+        data2: &'c [V],
+    }
+
+    #[allow(dead_code)]
+    #[derive(SplDiscriminate)]
+    #[discriminator_hash_input(
+        "global:my_instruction_with_multiple_generics_and_lifetime_and_where"
+    )]
+    pub struct MyInstruction7<'c, U, V>
+    where
+        U: Clone + Copy,
+        V: Clone + Copy,
+    {
+        data1: &'c [U],
+        data2: &'c [V],
     }
 
     fn assert_discriminator<T: spl_discriminator::discriminator::SplDiscriminate>(
         hash_input: &str,
     ) {
         let discriminator = build_discriminator(hash_input);
-        assert_eq!(T::SPL_DISCRIMINATOR, discriminator);
-        assert_eq!(T::SPL_DISCRIMINATOR_SLICE, discriminator.as_slice());
+        assert_eq!(
+            T::SPL_DISCRIMINATOR,
+            discriminator,
+            "Discriminator mismatch: case: {}",
+            hash_input
+        );
+        assert_eq!(
+            T::SPL_DISCRIMINATOR_SLICE,
+            discriminator.as_slice(),
+            "Discriminator mismatch: case: {}",
+            hash_input
+        );
     }
 
     fn build_discriminator(hash_input: &str) -> ArrayDiscriminator {
@@ -60,10 +106,39 @@ mod tests {
 
     #[test]
     fn test_discrminators() {
-        assert_discriminator::<MyInstruction1>("some_discriminator_hash_input");
-        assert_discriminator::<MyInstruction2>("yet_another_discriminator_hash_input");
-        assert_discriminator::<MyInstruction3>("global:my_instruction_3");
-        let runtime_discrim = ArrayDiscriminator::new_with_hash_input("my_new_hash_input");
-        assert_eq!(runtime_discrim, build_discriminator("my_new_hash_input"),);
+        let runtime_discrim = ArrayDiscriminator::new_with_hash_input("my_runtime_hash_input");
+        assert_eq!(
+            runtime_discrim,
+            build_discriminator("my_runtime_hash_input"),
+        );
+
+        assert_discriminator::<MyInstruction1>("my_first_instruction");
+        assert_discriminator::<MyInstruction2>("global:my_second_instruction");
+        assert_discriminator::<MyInstruction3<'_>>("global:my_instruction_with_lifetime");
+        assert_discriminator::<MyInstruction4<u8>>("global:my_instruction_with_one_generic");
+        assert_discriminator::<MyInstruction5<'_, u8>>(
+            "global:my_instruction_with_one_generic_and_lifetime",
+        );
+        assert_discriminator::<MyInstruction6<'_, u8, u8>>(
+            "global:my_instruction_with_multiple_generics_and_lifetime",
+        );
+        assert_discriminator::<MyInstruction7<'_, u8, u8>>(
+            "global:my_instruction_with_multiple_generics_and_lifetime_and_where",
+        );
+    }
+}
+
+#[cfg(all(test, feature = "borsh"))]
+mod borsh_test {
+    use super::*;
+
+    #[test]
+    fn borsh_test() {
+        let my_discrim = ArrayDiscriminator::new_with_hash_input("my_discrim");
+        let mut buffer = [0u8; 8];
+        my_discrim.serialize(&mut buffer[..]).unwrap();
+        let my_discrim_again = ArrayDiscriminator::try_from_slice(&buffer).unwrap();
+        assert_eq!(my_discrim, my_discrim_again);
+        assert_eq!(buf, my_discrim.into());
     }
 }

libraries/discriminator/src/lib.rs

@@ -11,15 +11,18 @@ use {
     proc_macro2::{Span, TokenStream},
     quote::{quote, ToTokens},
     solana_program::hash,
-    syn::{parse::Parse, Ident, Item, ItemEnum, ItemStruct, LitByteStr},
+    syn::{parse::Parse, Generics, Ident, Item, ItemEnum, ItemStruct, LitByteStr, WhereClause},
 };
 
 /// "Builder" struct to implement the `SplDiscriminate` trait
 /// on an enum or struct
-#[derive(Debug)]
 pub struct SplDiscriminateBuilder {
     /// The struct/enum identifier
     pub ident: Ident,
+    /// The item's generic arguments (if any)
+    pub generics: Generics,
+    /// The item's where clause for generics (if any)
+    pub where_clause: Option<WhereClause>,
     /// The TLV hash_input
     pub hash_input: String,
 }
@@ -29,8 +32,15 @@ impl TryFrom<ItemEnum> for SplDiscriminateBuilder {
 
     fn try_from(item_enum: ItemEnum) -> Result<Self, Self::Error> {
         let ident = item_enum.ident;
+        let where_clause = item_enum.generics.where_clause.clone();
+        let generics = item_enum.generics;
         let hash_input = parse_hash_input(&item_enum.attrs)?;
-        Ok(Self { ident, hash_input })
+        Ok(Self {
+            ident,
+            generics,
+            where_clause,
+            hash_input,
+        })
     }
 }
 
@@ -39,8 +49,15 @@ impl TryFrom<ItemStruct> for SplDiscriminateBuilder {
 
     fn try_from(item_struct: ItemStruct) -> Result<Self, Self::Error> {
         let ident = item_struct.ident;
+        let where_clause = item_struct.generics.where_clause.clone();
+        let generics = item_struct.generics;
         let hash_input = parse_hash_input(&item_struct.attrs)?;
-        Ok(Self { ident, hash_input })
+        Ok(Self {
+            ident,
+            generics,
+            where_clause,
+            hash_input,
+        })
     }
 }
 
@@ -70,9 +87,11 @@ impl ToTokens for SplDiscriminateBuilder {
 impl From<&SplDiscriminateBuilder> for TokenStream {
     fn from(builder: &SplDiscriminateBuilder) -> Self {
         let ident = &builder.ident;
+        let generics = &builder.generics;
+        let where_clause = &builder.where_clause;
         let bytes = get_discriminator_bytes(&builder.hash_input);
         quote! {
-            impl spl_discriminator::discriminator::SplDiscriminate for #ident {
+            impl #generics spl_discriminator::discriminator::SplDiscriminate for #ident #generics #where_clause {
                 const SPL_DISCRIMINATOR: spl_discriminator::discriminator::ArrayDiscriminator
                     = spl_discriminator::discriminator::ArrayDiscriminator::new(*#bytes);
             }