feat: blake2b integration (#297)

Co-authored-by: Ratan Kaliani <[email protected]>

plonky2x/core/src/backend/circuit/serialization/hints.rs

@@ -1,9 +1,6 @@
 use core::fmt::Debug;
 use core::marker::PhantomData;
 
-use curta::chip::hash::blake::blake2b::generator::{
-    BLAKE2BAirParameters, BLAKE2BGenerator, BLAKE2BHintGenerator,
-};
 use curta::machine::hash::sha::sha256::SHA256;
 use curta::machine::hash::sha::sha512::SHA512;
 use curta::plonky2::cubic::arithmetic_gate::ArithmeticCubicGenerator;
@@ -58,7 +55,8 @@ use crate::frontend::eth::mpt::generators::LteGenerator;
 use crate::frontend::eth::storage::generators::{
     EthBlockGenerator, EthLogGenerator, EthStorageKeyGenerator, EthStorageProofHint,
 };
-use crate::frontend::hash::blake2::curta::MAX_NUM_CURTA_CHUNKS;
+use crate::frontend::hash::blake2::digest_hint::BLAKE2BDigestHint;
+use crate::frontend::hash::blake2::proof_hint::BLAKE2BProofHint;
 use crate::frontend::hash::keccak::keccak256::Keccak256Generator;
 use crate::frontend::hash::sha::curta::digest_hint::SHADigestHint;
 use crate::frontend::hash::sha::curta::proof_hint::SHAProofHint;
@@ -369,26 +367,6 @@ where
         let id = MulCubicGenerator::<L::Field, D>::id();
         r.register_simple::<MulCubicGenerator<L::Field, D>>(id);
 
-        let blake2b_hint_generator_id = BLAKE2BHintGenerator::id();
-        r.register_simple::<BLAKE2BHintGenerator>(blake2b_hint_generator_id);
-
-        let blake2b_generator = BLAKE2BGenerator::<
-            L::Field,
-            L::CubicParams,
-            L::CurtaConfig,
-            D,
-            BLAKE2BAirParameters<L::Field, L::CubicParams>,
-            MAX_NUM_CURTA_CHUNKS,
-        >::id();
-        r.register_simple::<BLAKE2BGenerator<
-            L::Field,
-            L::CubicParams,
-            L::CurtaConfig,
-            D,
-            BLAKE2BAirParameters<L::Field, L::CubicParams>,
-            MAX_NUM_CURTA_CHUNKS,
-        >>(blake2b_generator);
-
         r.register_hint::<SubArrayExtractorHint>();
 
         r.register_hint::<BeaconBlockRootsHint>();
@@ -413,6 +391,12 @@ where
         r.register_hint::<EcOpResultHint>();
         r.register_async_hint::<Async<EcOpResultHint>>();
 
+        r.register_hint::<BLAKE2BDigestHint>();
+        r.register_async_hint::<Async<BLAKE2BDigestHint>>();
+
+        r.register_hint::<BLAKE2BProofHint>();
+        r.register_async_hint::<Async<BLAKE2BProofHint>>();
+
         let dummy_proof_generator_id =
             DummyProofGenerator::<L::Field, L::Config, D>::default().id();
         r.register_simple::<DummyProofGenerator<L::Field, L::Config, D>>(dummy_proof_generator_id);

plonky2x/core/src/frontend/builder/mod.rs

@@ -22,7 +22,7 @@ use tokio::runtime::Runtime;
 
 pub use self::io::CircuitIO;
 use super::ecc::curve25519::curta::accelerator::EcOpAccelerator;
-use super::hash::blake2::curta::Blake2bAccelerator;
+use super::hash::blake2::accelerator::BLAKE2BAccelerator;
 use super::hash::sha::sha256::curta::SHA256Accelerator;
 use super::hash::sha::sha512::curta::SHA512Accelerator;
 use super::hint::HintGenerator;
@@ -46,7 +46,7 @@ pub struct CircuitBuilder<L: PlonkParameters<D>, const D: usize> {
     pub(crate) async_hints: Vec<AsyncHintDataRef<L, D>>,
     pub(crate) async_hints_indices: Vec<usize>,
 
-    pub blake2b_accelerator: Option<Blake2bAccelerator<L, D>>,
+    pub blake2b_accelerator: Option<BLAKE2BAccelerator>,
     pub sha256_accelerator: Option<SHA256Accelerator>,
     pub sha512_accelerator: Option<SHA512Accelerator>,
     pub ec_25519_ops_accelerator: Option<EcOpAccelerator>,
@@ -134,7 +134,7 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
     fn pre_build(&mut self) {
         let blake2b_accelerator = self.blake2b_accelerator.clone();
         if let Some(accelerator) = blake2b_accelerator {
-            accelerator.build(self);
+            self.curta_constrain_blake2b(accelerator);
         }
 
         let sha256_accelerator = self.sha256_accelerator.clone();

plonky2x/core/src/frontend/curta/builder.rs

@@ -152,6 +152,8 @@ mod tests {
             a_init.get(1),
             &Time::zero(),
             Some(num_rows_reg),
+            None,
+            None,
         );
 
         let clk = Time::from_element(air_builder.clk);
@@ -161,10 +163,10 @@ mod tests {
         let zero_trace = air_builder.alloc::<ElementRegister>();
         air_builder.set_to_expression(&zero_trace, GoldilocksField::ZERO.into());
         let a_0_trace = a_ptr.get_at(zero_trace);
-        let a = air_builder.load(&a_0_trace, &clk);
-        let b = air_builder.load(&a_ptr.get(1), &Time::zero());
+        let a = air_builder.load(&a_0_trace, &clk, None, None);
+        let b = air_builder.load(&a_ptr.get(1), &Time::zero(), None, None);
         let c = air_builder.and(&a, &b);
-        air_builder.store(&a_0_trace, c, &clk.advance(), None);
+        air_builder.store(&a_0_trace, c, &clk.advance(), None, None, None);
 
         let a_final = air_builder.api.alloc_public::<U32Register>();
 

plonky2x/core/src/frontend/ecc/curve25519/ed25519/eddsa.rs

@@ -18,8 +18,6 @@ use crate::prelude::{
     U256Variable, U32Variable, Variable,
 };
 
-const MAX_NUM_SIGS: usize = 256;
-
 #[derive(Clone, Debug, CircuitVariable)]
 pub struct EDDSASignatureVariable {
     pub r: CompressedEdwardsYVariable,
@@ -82,7 +80,7 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
         signatures: ArrayVariable<EDDSASignatureVariable, NUM_SIGS>,
         pubkeys: ArrayVariable<CompressedEdwardsYVariable, NUM_SIGS>,
     ) {
-        assert!(NUM_SIGS > 0 && NUM_SIGS <= MAX_NUM_SIGS);
+        assert!(NUM_SIGS > 0);
         assert!(is_active.len() == NUM_SIGS);
         assert!(messages.len() == NUM_SIGS);
         if let Some(ref msg_lens) = message_byte_lengths {
@@ -141,7 +139,7 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
         signatures: ArrayVariable<EDDSASignatureVariable, NUM_SIGS>,
         pubkeys: ArrayVariable<CompressedEdwardsYVariable, NUM_SIGS>,
     ) {
-        assert!(NUM_SIGS > 0 && NUM_SIGS <= MAX_NUM_SIGS);
+        assert!(NUM_SIGS > 0);
         assert!(messages.len() == NUM_SIGS);
         if let Some(ref msg_lens) = message_byte_lengths {
             assert!(msg_lens.len() == NUM_SIGS);

plonky2x/core/src/frontend/hash/blake2/accelerator.rs

@@ -0,0 +1,8 @@
+use super::request::BLAKE2BRequest;
+use crate::prelude::U64Variable;
+
+#[derive(Debug, Clone)]
+pub struct BLAKE2BAccelerator {
+    pub blake2b_requests: Vec<BLAKE2BRequest>,
+    pub blake2b_responses: Vec<[U64Variable; 4]>,
+}

plonky2x/core/src/frontend/hash/blake2/builder.rs

@@ -0,0 +1,57 @@
+use super::accelerator::BLAKE2BAccelerator;
+use super::digest_hint::BLAKE2BDigestHint;
+use super::proof_hint::BLAKE2BProofHint;
+use super::request::BLAKE2BRequest;
+use super::stark::{get_blake2b_data, stark};
+use crate::frontend::hint::synchronous::Async;
+use crate::prelude::*;
+
+impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
+    /// The constraints for an accelerated BLAKE2B computation using Curta.
+    pub(crate) fn curta_constrain_blake2b(&mut self, accelerator: BLAKE2BAccelerator) {
+        // Get all the digest values using the digest hint.
+        for (request, response) in accelerator
+            .blake2b_requests
+            .iter()
+            .zip(accelerator.blake2b_responses.iter())
+        {
+            let digest_hint = BLAKE2BDigestHint::new();
+            let mut input_stream = VariableStream::new();
+
+            match &request {
+                BLAKE2BRequest::Fixed(msg) => {
+                    let len = self.constant::<Variable>(L::Field::from_canonical_usize(msg.len()));
+                    input_stream.write(&len);
+                    input_stream.write_slice(msg);
+                }
+                BLAKE2BRequest::Variable(msg, len, _) => {
+                    input_stream.write(len);
+                    input_stream.write_slice(msg);
+                }
+            }
+
+            let output_stream = self.hint(input_stream, digest_hint);
+            let digest = output_stream.read::<[U64Variable; 4]>(self);
+            self.assert_is_equal(digest, *response);
+        }
+
+        // Prove correctness of the digest using the proof hint.
+
+        // Initialize the corresponding stark and hint.
+        let blake2b_data = get_blake2b_data(self, accelerator);
+        let parameters = blake2b_data.parameters();
+        let blake2b_stark = stark(parameters);
+        let proof_hint = BLAKE2BProofHint::new(parameters);
+        let mut input_stream = VariableStream::new();
+        input_stream.write_blake2b_input(&blake2b_data);
+
+        // Read the stark proof and public inputs from the hint's output stream.
+        let output_stream = self.async_hint(input_stream, Async(proof_hint));
+        let proof = output_stream.read_byte_stark_proof(self, &blake2b_stark.stark);
+        let num_public_inputs = blake2b_stark.stark.air_data.num_public_inputs;
+        let public_inputs = output_stream.read_vec(self, num_public_inputs);
+
+        // Verify the proof.
+        blake2b_stark.verify_proof(self, proof, &public_inputs, blake2b_data)
+    }
+}