fix: blake2b digest targets (#83)

succinctlabs · Sep 22, 2023 · ce89331 · ce89331
1 parent 20cf588
commit ce89331
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/curta/src/chip/hash/blake/blake2b/builder_gadget.rs b/curta/src/chip/hash/blake/blake2b/builder_gadget.rs
@@ -20,6 +20,7 @@ pub struct BLAKE2BBuilderGadget<L: AirParameters + 'static + Clone + Debug + Sen
     pub padded_messages: Vec<Target>,
     pub msg_lengths: Vec<Target>,
     pub digests: Vec<Target>,
+    pub chunk_sizes: Vec<usize>,
     _phantom: PhantomData<L>,
 }
 
@@ -68,6 +69,7 @@ impl<
             padded_messages: Vec::new(),
             msg_lengths: Vec::new(),
             digests: Vec::new(),
+            chunk_sizes: Vec::new(),
             _phantom: PhantomData,
         }
     }
@@ -95,7 +97,11 @@ impl<
         gadget: Self::Gadget,
     ) {
         // Allocate public input targets
-        let public_blake2b_targets = BLAKE2BPublicData::add_virtual::<F, D, L>(self);
+        let public_blake2b_targets = BLAKE2BPublicData::add_virtual::<F, D, L>(
+            self,
+            gadget.digests.as_slice(),
+            gadget.chunk_sizes.as_slice(),
+        );
 
         let stark_data = BLAKE2BGenerator::<F, E, C, D, L>::stark_data();
         let BLAKE2BStarkData { stark, config, .. } = stark_data;

diff --git a/curta/src/chip/hash/blake/blake2b/generator.rs b/curta/src/chip/hash/blake/blake2b/generator.rs
@@ -132,7 +132,7 @@ impl<
     > SimpleGenerator<F, D> for BLAKE2BGenerator<F, E, C, D, L>
 {
     fn id(&self) -> String {
-        "BLAKE2B generator".to_string()
+        Self::id()
     }
 
     fn serialize(
@@ -238,6 +238,8 @@ impl BLAKE2BPublicData<Target> {
         L: AirParameters + 'static + Clone + Debug + Send + Sync,
     >(
         builder: &mut CircuitBuilder<F, D>,
+        digests: &[Target],
+        chunk_sizes: &[usize],
     ) -> Self {
         let num_chunks = L::num_rows() / NUM_MIX_ROUNDS;
 
@@ -253,9 +255,25 @@ impl BLAKE2BPublicData<Target> {
             .map(|_| builder.add_virtual_target_arr::<8>())
             .collect::<Vec<_>>();
 
-        let hash_state_targets = (0..num_chunks * HASH_ARRAY_SIZE)
-            .map(|_| builder.add_virtual_target_arr::<8>())
-            .collect::<Vec<_>>();
+        let mut hash_state_targets = Vec::new();
+        assert!(digests.len() / 8 <= num_chunks * HASH_ARRAY_SIZE);
+        assert!(digests.len() % 8 == 0);
+
+        for (digest, chunk_size) in digests.chunks_exact(32).zip_eq(chunk_sizes.iter()) {
+            hash_state_targets
+                .extend((0..8 * (chunk_size - 1)).map(|_| builder.add_virtual_target_arr::<8>()));
+
+            let u64_digest_byte = digest.chunks_exact(8).map(|arr| {
+                let array: [Target; 8] = arr.try_into().unwrap();
+                array
+            });
+            hash_state_targets.extend(u64_digest_byte);
+            hash_state_targets.extend((0..4).map(|_| builder.add_virtual_target_arr::<8>()));
+        }
+
+        for _ in hash_state_targets.len()..num_chunks * HASH_ARRAY_SIZE {
+            hash_state_targets.push(builder.add_virtual_target_arr::<8>());
+        }
 
         BLAKE2BPublicData {
             msg_chunks: msg_chunks_targets,