fixed 1 test

Author: ananas-block

sdk-libs/macros/src/hasher.rs

@@ -82,7 +82,7 @@ pub(crate) fn hasher(input: ItemStruct) -> Result<TokenStream> {
             .any(|attr| attr.path().is_ident("flatten"))
     });
 
-    let mut flattned_fields_added = vec![quote! { Self::NUM_FIELDS as usize }];
+    let mut flattned_fields_added = vec![quote! { 0usize }];
     let mut truncate_code = Vec::new();
 
     // Process each field
@@ -206,21 +206,37 @@ pub(crate) fn hasher(input: ItemStruct) -> Result<TokenStream> {
             flattned_fields_added.push(quote! {
                 + #field_type::NUM_FIELDS as usize
             });
+
+            // For flattened fields, we need to use their individual byte arrays directly
+            // rather than hashing the whole struct first
             code.push(quote! {
                 {
-                    for (j, element) in <#field_type as ::light_hasher::to_byte_array::ToByteArray>::to_byte_arrays::<{#field_type::NUM_FIELDS}>(&self.#field_name)?.iter().enumerate() {
-                        field_array[#i + j + num_flattned_fields ] = *element;
-                        num_flattned_fields +=1;
+                    // Get individual byte arrays from the flattened field
+                    let flattened_arrays = <#field_type as ::light_hasher::to_byte_array::ToByteArray>::to_byte_arrays::<{#field_type::NUM_FIELDS}>(&self.#field_name)?;
+                    // Add each element individually to the field_array
+                    for element in flattened_arrays.iter() {
+                        field_array[num_flattned_fields] = *element;
+                        num_flattned_fields += 1;
                     }
                 }
             });
         } else {
+            if flatten_field_exists {
+                flattned_fields_added.push(quote! {
+                    + 1
+                });
+            }
             to_byte_arrays_fields.push(quote! {
                 arrays[#i ] =  self.#field_name.to_byte_array()?;
             });
             if flatten_field_exists {
+                // Store field index in the field_assignments for later non-flattened field processing
                 field_assignments.push(quote! {
-                    field_array[#i + num_flattned_fields ] = self.#field_name.to_byte_array()?;
+                    #i
+                });
+                code.push(quote! {
+                    field_array[num_flattned_fields] = self.#field_name.to_byte_array()?;
+                    num_flattned_fields += 1;
                 });
             } else {
                 field_assignments.push(quote! {
@@ -242,8 +258,9 @@ pub(crate) fn hasher(input: ItemStruct) -> Result<TokenStream> {
             },
         );
         code.push(quote! {
-            for element in field_array.iter() {
-                slices[num_flattned_fields] = element.as_slice();
+            // Set all slices properly for both flattened and non-flattened fields
+            for i in 0..num_flattned_fields {
+                slices[i] = field_array[i].as_slice();
             }
         });
         quote! {
@@ -310,9 +327,35 @@ pub(crate) fn hasher(input: ItemStruct) -> Result<TokenStream> {
         }
     };
 
+    // Calculate the total number of fields, accounting for flattened fields
+    let total_field_count = if flatten_field_exists {
+        // When there are flattened fields, we need to adjust the total field count
+        let mut sum = quote! { 0 };
+
+        for field in fields.named.iter() {
+            let flatten = field
+                .attrs
+                .iter()
+                .any(|attr| attr.path().is_ident("flatten"));
+
+            if flatten {
+                // Use the field type's NUM_FIELDS instead of counting as one field
+                let field_type = &field.ty;
+                sum = quote! { #sum + #field_type::NUM_FIELDS };
+            } else {
+                // Regular fields count as one
+                sum = quote! { #sum + 1 };
+            }
+        }
+        sum
+    } else {
+        // Without flattened fields, just use the regular field count
+        quote! { #field_count }
+    };
+
     Ok(quote! {
         impl #impl_gen ::light_hasher::to_byte_array::ToByteArray for #struct_name #type_gen #where_clause {
-            const NUM_FIELDS: usize = #field_count;
+            const NUM_FIELDS: usize = #total_field_count;
 
             fn to_byte_array(&self) -> ::std::result::Result<[u8; 32], ::light_hasher::HasherError> {
                 #to_byte_array

sdk-libs/macros/tests/flatten.rs

@@ -43,7 +43,9 @@ fn flatten() {
         array[31] = 4;
         array
     };
-    let manual_hash =
-        Poseidon::hashv(&[one.as_ref(), two.as_ref(), three.as_ref(), four.as_ref()]).unwrap();
+    let hash = Poseidon::hashv(&[three.as_slice(), four.as_slice()]).unwrap();
+    let manual_slices = [one.as_ref(), two.as_ref(), hash.as_ref()];
+    println!("manual_slices {:?}", manual_slices);
+    let manual_hash = Poseidon::hashv(&manual_slices).unwrap();
     assert_eq!(test.hash::<Poseidon>().unwrap(), manual_hash);
 }