feat: #[derive(LightHasher)] macro and AsByteVec trait (#1085)

* feat: `#[derive(LightHasher)]` macro and `AsByteVec` trait

- `AsByteVec`, providing `as_byte_vec()` method, guarantees consistent way
  of serializing types to 2D byte vectors before hashing. We provide default
  implementations for primitives. More complex types need to implement
  `IntoBytes` themselves.
  - By using 2D vectors, we make it possible to represent structs with
    multiple fields. This way, we can handle struct fields (nested structs)
    the same way as primitive fields and deal with all types by using
    one trait.
  - The reason behing using vectors instead of slices is that in some
    cases, we cannot just cast the type without defining custom bytes.
    Vectors, although they might introduce copies, make sure we always
    own the bytes we are creating.
- `#[derive(LightHasher)]` implements `AsByteVec` and `DataHasher` traits.
  The `DataHasher` implementation takes bytes returned by `AsByteVec` as
  an input.

* test: Add more test cases for the `AsByteVec` trait

- empty string
- `Option` with `Some` and `None` and assertion that `None != Some(0)`
- array (including an empty one)

* doc: Add documentation to `LightHasher` macro

Author: vadorovsky

Cargo.lock

@@ -14,5 +14,8 @@ syn = { version = "1.0", features = ["full"] }
 
 light-hasher = { path = "../../merkle-tree/hasher", version = "0.3.0" }
 
+[dev-dependencies]
+light-utils = { path = "../../utils", version = "0.3.0" }
+
 [lib]
 proc-macro = true

macros/light/Cargo.toml

@@ -0,0 +1,107 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{Error, Fields, ItemStruct, Result};
+
+pub(crate) fn hasher(input: ItemStruct) -> Result<TokenStream> {
+    let struct_name = &input.ident;
+
+    let (impl_gen, type_gen, where_clause) = input.generics.split_for_impl();
+
+    let fields = match input.fields {
+        Fields::Named(fields) => fields,
+        _ => {
+            return Err(Error::new_spanned(
+                input,
+                "Only structs with named fields are supported",
+            ))
+        }
+    };
+
+    let field_into_bytes_calls = fields
+        .named
+        .iter()
+        .filter(|field| {
+            !field.attrs.iter().any(|attr| {
+                if let Some(attr_ident) = attr.path.get_ident() {
+                    attr_ident == "skip"
+                } else {
+                    false
+                }
+            })
+        })
+        .map(|field| {
+            let field_name = &field.ident;
+            let truncate = field.attrs.iter().any(|attr| {
+                if let Some(attr_ident) = attr.path.get_ident() {
+                    attr_ident == "truncate"
+                } else {
+                    false
+                }
+            });
+            if truncate {
+                quote! {
+                    let truncated_bytes = self
+                        .#field_name
+                        .as_byte_vec()
+                        .iter()
+                        .map(|bytes| {
+                            let (bytes, _) = ::light_utils::hash_to_bn254_field_size_be(bytes).expect(
+                                "Could not truncate the field #field_name to the BN254 prime field"
+                            );
+                            bytes.to_vec()
+                        })
+                        .collect::<Vec<Vec<u8>>>();
+                    result.extend_from_slice(truncated_bytes.as_slice());
+                }
+            } else {
+                quote! {
+                    result.extend_from_slice(self.#field_name.as_byte_vec().as_slice());
+                }
+            }
+        })
+        .collect::<Vec<_>>();
+
+    Ok(quote! {
+        impl #impl_gen ::light_hasher::bytes::AsByteVec for #struct_name #type_gen #where_clause {
+            fn as_byte_vec(&self) -> Vec<Vec<u8>> {
+                use ::light_hasher::bytes::AsByteVec;
+
+                let mut result: Vec<Vec<u8>> = Vec::new();
+                #(#field_into_bytes_calls)*
+                result
+            }
+        }
+
+        impl #impl_gen ::light_hasher::DataHasher for #struct_name #type_gen #where_clause {
+            fn hash<H: light_hasher::Hasher>(&self) -> ::std::result::Result<[u8; 32], ::light_hasher::errors::HasherError> {
+                use ::light_hasher::bytes::AsByteVec;
+
+                H::hashv(self.as_byte_vec().iter().map(|v| v.as_slice()).collect::<Vec<_>>().as_slice())
+            }
+        }
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use syn::parse_quote;
+
+    #[test]
+    fn test_light_hasher() {
+        let input: ItemStruct = parse_quote! {
+            struct MyAccount {
+                a: u32,
+                b: i32,
+                c: u64,
+                d: i64,
+            }
+        };
+
+        let output = hasher(input).unwrap();
+        let output = output.to_string();
+
+        println!("{output}");
+    }
+}

macros/light/src/hasher.rs

@@ -7,6 +7,7 @@ use traits::process_light_traits;
 
 mod accounts;
 mod discriminator;
+mod hasher;
 mod pubkey;
 mod traits;
 
@@ -161,3 +162,119 @@ pub fn light_discriminator(input: TokenStream) -> TokenStream {
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+/// Makes the annotated struct hashable by implementing the following traits:
+///
+/// - [`AsByteVec`](light_hasher::bytes::AsByteVec), which makes the struct
+///   convertable to a 2D byte vector.
+/// - [`DataHasher`](light_hasher::DataHasher), which makes the struct hashable
+///   with the `hash()` method, based on the byte inputs from `AsByteVec`
+///   implementation.
+///
+/// This macro assumes that all the fields of the struct implement the
+/// `AsByteVec` trait. The trait is implemented by default for the most of
+/// standard Rust types (primitives, `String`, arrays and options carrying the
+/// former). If there is a field of a type not implementing the trait, there
+/// are two options:
+///
+/// 1. The most recommended one - annotating that type with the `light_hasher`
+///    macro as well.
+/// 2. Manually implementing the `AsByteVec` trait.
+///
+/// # Attributes
+///
+/// - `skip` - skips the given field, it doesn't get included neither in
+///   `AsByteVec` nor `DataHasher` implementation.
+/// - `truncate` - makes sure that the byte value does not exceed the BN254
+///   prime field modulus, by hashing it (with Keccak) and truncating it to 31
+///   bytes. It's generally a good idea to use it on any field which is
+///   expected to output more than 31 bytes.
+///
+/// # Examples
+///
+/// Compressed account with only primitive types as fields:
+///
+/// ```ignore
+/// #[derive(LightHasher)]
+/// pub struct MyCompressedAccount {
+///     a: i64,
+///     b: Option<u64>,
+/// }
+/// ```
+///
+/// Compressed account with fields which might exceed the BN254 prime field:
+///
+/// ```ignore
+/// #[derive(LightHasher)]
+/// pub struct MyCompressedAccount {
+///     a: i64
+///     b: Option<u64>,
+///     #[truncate]
+///     c: [u8; 32],
+///     #[truncate]
+///     d: String,
+/// }
+/// ```
+///
+/// Compressed account with fields we want to skip:
+///
+/// ```ignore
+/// #[derive(LightHasher)]
+/// pub struct MyCompressedAccount {
+///     a: i64
+///     b: Option<u64>,
+///     #[skip]
+///     c: [u8; 32],
+/// }
+/// ```
+///
+/// Compressed account with a nested struct:
+///
+/// ```ignore
+/// #[derive(LightHasher)]
+/// pub struct MyCompressedAccount {
+///     a: i64
+///     b: Option<u64>,
+///     c: MyStruct,
+/// }
+///
+/// #[derive(LightHasher)]
+/// pub struct MyStruct {
+///     a: i32
+///     b: u32,
+/// }
+/// ```
+///
+/// Compressed account with a type with a custom `AsByteVec` implementation:
+///
+/// ```ignore
+/// #[derive(LightHasher)]
+/// pub struct MyCompressedAccount {
+///     a: i64
+///     b: Option<u64>,
+///     c: RData,
+/// }
+///
+/// pub enum RData {
+///     A(Ipv4Addr),
+///     AAAA(Ipv6Addr),
+///     CName(String),
+/// }
+///
+/// impl AsByteVec for RData {
+///     fn as_byte_vec(&self) -> Vec<Vec<u8>> {
+///         match self {
+///             Self::A(ipv4_addr) => vec![ipv4_addr.octets().to_vec()],
+///             Self::AAAA(ipv6_addr) => vec![ipv6_addr.octets().to_vec()],
+///             Self::CName(cname) => cname.as_byte_vec(),
+///         }
+///     }
+/// }
+/// ```
+#[proc_macro_derive(LightHasher, attributes(skip, truncate))]
+pub fn light_hasher(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as ItemStruct);
+    hasher::hasher(input)
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}

macros/light/src/lib.rs

@@ -0,0 +1,78 @@
+use std::{cell::RefCell, marker::PhantomData, rc::Rc};
+
+use light_hasher::{DataHasher, Poseidon};
+use light_macros::LightHasher;
+
+#[derive(LightHasher)]
+pub struct MyAccount {
+    pub a: bool,
+    pub b: u64,
+    pub c: MyNestedStruct,
+    #[truncate]
+    pub d: [u8; 32],
+    #[skip]
+    pub e: MyNestedNonHashableStruct,
+    pub f: Option<usize>,
+}
+
+#[derive(LightHasher)]
+pub struct MyNestedStruct {
+    pub a: i32,
+    pub b: u32,
+    #[truncate]
+    pub c: String,
+}
+
+pub struct MyNestedNonHashableStruct {
+    pub a: PhantomData<()>,
+    pub b: Rc<RefCell<usize>>,
+}
+
+#[test]
+fn test_light_hasher() {
+    let my_account = MyAccount {
+        a: true,
+        b: u64::MAX,
+        c: MyNestedStruct {
+            a: i32::MIN,
+            b: u32::MAX,
+            c: "wao".to_string(),
+        },
+        d: [u8::MAX; 32],
+        e: MyNestedNonHashableStruct {
+            a: PhantomData,
+            b: Rc::new(RefCell::new(usize::MAX)),
+        },
+        f: None,
+    };
+    assert_eq!(
+        my_account.hash::<Poseidon>().unwrap(),
+        [
+            44, 62, 31, 169, 73, 125, 135, 126, 176, 7, 127, 96, 183, 224, 156, 140, 105, 77, 225,
+            230, 174, 196, 38, 92, 0, 44, 19, 25, 255, 109, 6, 168
+        ]
+    );
+
+    let my_account = MyAccount {
+        a: true,
+        b: u64::MAX,
+        c: MyNestedStruct {
+            a: i32::MIN,
+            b: u32::MAX,
+            c: "wao".to_string(),
+        },
+        d: [u8::MAX; 32],
+        e: MyNestedNonHashableStruct {
+            a: PhantomData,
+            b: Rc::new(RefCell::new(usize::MAX)),
+        },
+        f: Some(0),
+    };
+    assert_eq!(
+        my_account.hash::<Poseidon>().unwrap(),
+        [
+            32, 205, 141, 227, 236, 5, 28, 219, 24, 164, 215, 79, 151, 131, 162, 82, 224, 101, 171,
+            201, 4, 181, 26, 146, 6, 1, 95, 107, 239, 19, 233, 80
+        ]
+    );
+}

macros/light/tests/hasher.rs

@@ -18,3 +18,6 @@ thiserror = "1.0"
 ark-bn254 = "0.4.0"
 sha2 = "0.10"
 sha3 = "0.10"
+
+[dev-dependencies]
+rand = "0.8"