feat: hasher derivation for derived zero copy

Author: ananas-block

Cargo.lock

@@ -17,6 +17,7 @@ thiserror = { workspace = true }
 arrayvec = { workspace = true }
 num-bigint = { workspace = true }
 zerocopy = { workspace = true }
+light-zero-copy = { workspace = true }
 
 [target.'cfg(not(target_os = "solana"))'.dependencies]
 ark-bn254 = { workspace = true }

program-libs/hasher/Cargo.toml

@@ -1,4 +1,5 @@
 use arrayvec::ArrayVec;
+use light_zero_copy::ZeroCopyTraits;
 
 use crate::{keccak::Keccak, to_byte_array::ToByteArray, Hasher, HasherError};
 
@@ -41,7 +42,78 @@ where
         arrays.iter().for_each(|x| slices.push(x.as_slice()));
         let bump_seed = [HASH_TO_FIELD_SIZE_SEED];
         slices.push(bump_seed.as_slice());
-        Keccak::hashv(slices.as_slice())
+        let mut hash = Keccak::hashv(slices.as_slice())?;
+        hash[0] = 0;
+        Ok(hash)
+    }
+}
+
+impl HashToFieldSize for &mut [u8] {
+    fn hash_to_field_size(&self) -> Result<[u8; 32], HasherError> {
+        let mut slices: [&[u8]; 2] = [&[]; 2];
+        slices[0] = self;
+        let bump_seed = [HASH_TO_FIELD_SIZE_SEED];
+        slices[1] = bump_seed.as_slice();
+        let mut hash = Keccak::hashv(slices.as_slice())?;
+        hash[0] = 0;
+        Ok(hash)
+    }
+}
+
+impl HashToFieldSize for &[u8] {
+    fn hash_to_field_size(&self) -> Result<[u8; 32], HasherError> {
+        let mut slices: [&[u8]; 2] = [&[]; 2];
+        slices[0] = self;
+        let bump_seed = [HASH_TO_FIELD_SIZE_SEED];
+        slices[1] = bump_seed.as_slice();
+        let mut hash = Keccak::hashv(slices.as_slice())?;
+        hash[0] = 0;
+        Ok(hash)
+    }
+}
+
+impl<T, const PAD: bool, L> HashToFieldSize for light_zero_copy::slice::ZeroCopySlice<'_, L, T, PAD>
+where
+    T: ToByteArray + ZeroCopyTraits,
+    L: ZeroCopyTraits,
+    u64: From<L>,
+{
+    fn hash_to_field_size(&self) -> Result<[u8; 32], HasherError> {
+        let mut arrays = Vec::with_capacity(self.len());
+        for item in self.iter() {
+            let byte_array = item.to_byte_array()?;
+            arrays.push(byte_array);
+        }
+        let mut slices = Vec::with_capacity(self.len() + 1);
+        arrays.iter().for_each(|x| slices.push(x.as_slice()));
+        let bump_seed = [HASH_TO_FIELD_SIZE_SEED];
+        slices.push(bump_seed.as_slice());
+        let mut hash = Keccak::hashv(slices.as_slice())?;
+        hash[0] = 0;
+        Ok(hash)
+    }
+}
+
+impl<T, const PAD: bool, L> HashToFieldSize
+    for light_zero_copy::slice_mut::ZeroCopySliceMut<'_, L, T, PAD>
+where
+    T: ToByteArray + ZeroCopyTraits,
+    L: ZeroCopyTraits,
+    u64: From<L>,
+{
+    fn hash_to_field_size(&self) -> Result<[u8; 32], HasherError> {
+        let mut arrays = Vec::with_capacity(self.len());
+        for item in self.iter() {
+            let byte_array = item.to_byte_array()?;
+            arrays.push(byte_array);
+        }
+        let mut slices = Vec::with_capacity(self.len() + 1);
+        arrays.iter().for_each(|x| slices.push(x.as_slice()));
+        let bump_seed = [HASH_TO_FIELD_SIZE_SEED];
+        slices.push(bump_seed.as_slice());
+        let mut hash = Keccak::hashv(slices.as_slice())?;
+        hash[0] = 0;
+        Ok(hash)
     }
 }
 

program-libs/hasher/src/hash_to_field_size.rs

@@ -1,5 +1,3 @@
-use zerocopy::IntoBytes;
-
 use crate::{Hasher, HasherError, Poseidon};
 
 pub trait ToByteArray {
@@ -57,24 +55,35 @@ impl ToByteArray for solana_program::pubkey::Pubkey {
     }
 }
 
-impl ToByteArray for zerocopy::little_endian::U16 {
-    const NUM_FIELDS: usize = 1;
+macro_rules! impl_to_byte_array_zero_copy_primitive_types {
+    ($target:ty, $inner:ty) => {
+        impl ToByteArray for $target {
+            const NUM_FIELDS: usize = 1;
+            const IS_PRIMITIVE: bool = true;
 
-    fn to_byte_array(&self) -> Result<[u8; 32], HasherError> {
-        let bytes: u16 = (*self).into();
-        println!("U16 to byte array {}", bytes);
-        bytes.to_byte_array()
-    }
+            fn to_byte_array(&self) -> Result<[u8; 32], HasherError> {
+                let bytes: $inner = (*self).into();
+                bytes.to_byte_array()
+            }
 
-    fn to_byte_arrays<const NUM_FIELDS: usize>(
-        &self,
-    ) -> Result<[[u8; 32]; NUM_FIELDS], HasherError> {
-        if Self::NUM_FIELDS != NUM_FIELDS {
-            return Err(HasherError::InvalidNumFields);
+            fn to_byte_arrays<const NUM_FIELDS: usize>(
+                &self,
+            ) -> Result<[[u8; 32]; NUM_FIELDS], HasherError> {
+                if Self::NUM_FIELDS != NUM_FIELDS {
+                    return Err(HasherError::InvalidNumFields);
+                }
+                Ok([self.to_byte_array()?; NUM_FIELDS])
+            }
         }
-        Ok([self.to_byte_array()?; NUM_FIELDS])
-    }
+    };
 }
+use zerocopy::little_endian::{I16, I32, I64, U16, U32, U64};
+impl_to_byte_array_zero_copy_primitive_types!(U16, u16);
+impl_to_byte_array_zero_copy_primitive_types!(U32, u32);
+impl_to_byte_array_zero_copy_primitive_types!(U64, u64);
+impl_to_byte_array_zero_copy_primitive_types!(I16, i16);
+impl_to_byte_array_zero_copy_primitive_types!(I32, i32);
+impl_to_byte_array_zero_copy_primitive_types!(I64, i64);
 
 impl<T: ToByteArray> ToByteArray for Option<T> {
     const NUM_FIELDS: usize = 1;

program-libs/hasher/src/to_byte_array.rs

@@ -67,14 +67,15 @@ pub fn derive_zero_copy(input: TokenStream) -> TokenStream {
         }
         false
     });
+    let saved_hasher = hasher;
     println!("hasher {}", hasher);
 
     // Process the input to extract struct information
     let (name, z_struct_name, z_struct_meta_name, fields) = utils::process_input(&input);
 
     // Process the fields to separate meta fields and struct fields
     let (meta_fields, struct_fields) = utils::process_fields(fields);
-
+    // let hasher = false;
     // Generate each implementation part using the respective modules
     let meta_struct_def_mut =
         meta_struct::generate_meta_struct::<true>(&z_struct_meta_name, &meta_fields, hasher);
@@ -88,6 +89,7 @@ pub fn derive_zero_copy(input: TokenStream) -> TokenStream {
         &meta_fields,
         hasher,
     );
+    let hasher = saved_hasher;
     let z_struct_def = z_struct::generate_z_struct::<false>(
         &z_struct_name,
         &z_struct_meta_name,
@@ -159,20 +161,6 @@ pub fn derive_zero_copy_eq(input: TokenStream) -> TokenStream {
     // Parse the input DeriveInput
     let input = parse_macro_input!(input as DeriveInput);
 
-    // Check for LightHasher in derive attributes (same logic as in derive_zero_copy)
-    let hasher = input.attrs.iter().any(|attr| {
-        if attr.path().is_ident("poseidon_hasher") {
-            return true;
-        }
-
-        if attr.path().is_ident("derive") {
-            let derive_string = attr.to_token_stream().to_string();
-            return derive_string.contains("LightHasher");
-        }
-
-        false
-    });
-
     // Process the input to extract struct information
     let (name, z_struct_name, z_struct_meta_name, fields) = utils::process_input(&input);
 

program-libs/zero-copy-derive/src/lib.rs

@@ -19,8 +19,21 @@ pub fn generate_meta_struct<const MUT: bool>(
     // Generate the meta struct fields with converted types
     let meta_fields_with_converted_types = meta_fields.iter().map(|field| {
         let field_name = &field.ident;
+        let attributes = if hasher {
+            field
+                .attrs
+                .iter()
+                .map(|attr| {
+                    let path = attr;
+                    quote! { #path }
+                })
+                .collect::<Vec<_>>()
+        } else {
+            vec![quote! {}]
+        };
         let field_type = convert_to_zerocopy_type(&field.ty);
         quote! {
+            #(#attributes)*
             pub #field_name: #field_type
         }
     });