feat(exm): adapt design of sha256 for StepCircuit

Our task is to customize the design, namely to accept cell inputs and return them as well WIP
snarkify · Oct 5, 2023 · 03fb9f3 · 03fb9f3
1 parent 530ea90
commit 03fb9f3
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 15 deletions.
diff --git a/examples/sha256/main.rs b/examples/sha256/main.rs
@@ -65,6 +65,7 @@ pub mod sha256 {
             layouter: &mut impl Layouter<F>,
             initialized_state: &Self::State,
             input: [Self::BlockWord; BLOCK_SIZE],
+            cells: Option<[AssignedCell<F, F>; BLOCK_SIZE]>,
         ) -> Result<Self::State, Error>;
 
         /// Converts the given state into a message digest.
@@ -134,6 +135,7 @@ pub mod sha256 {
                 self.cur_block[..]
                     .try_into()
                     .expect("cur_block.len() == BLOCK_SIZE"),
+                None,
             )?;
             self.cur_block.clear();
 
@@ -145,6 +147,7 @@ pub mod sha256 {
                     &mut layouter,
                     &self.state,
                     chunk.try_into().expect("chunk.len() == BLOCK_SIZE"),
+                    None,
                 )?;
             }
 
@@ -172,6 +175,7 @@ pub mod sha256 {
                     self.cur_block[..]
                         .try_into()
                         .expect("cur_block.len() == BLOCK_SIZE"),
+                    None,
                 )?;
             }
             self.chip
@@ -195,6 +199,7 @@ pub mod sha256 {
                     self.cur_block[..]
                         .try_into()
                         .expect("cur_block.len() == BLOCK_SIZE"),
+                    None,
                 )?;
             }
             self.chip.digest_cells(&mut layouter, &self.state)
@@ -288,7 +293,7 @@ impl Circuit<pallas::Base> for TestSha256Circuit {
 
 type B = pallas::Base;
 // TODO
-const ARITY: usize = 31 * BLOCK_SIZE;
+const ARITY: usize = 8;
 
 impl StepCircuit<ARITY, B> for TestSha256Circuit {
     type StepConfig = Table16Config;
@@ -306,15 +311,11 @@ impl StepCircuit<ARITY, B> for TestSha256Circuit {
         Table16Chip::load(config.clone(), layouter)?;
         let table16_chip = Table16Chip::construct(config);
 
-        Ok(iter::repeat(
-            Sha256::digest_cells(table16_chip, layouter.namespace(|| "'abc' * 2"), &z_in)?
-                .into_iter(),
-        )
-        .take(ARITY / BLOCK_SIZE)
-        .flatten()
-        .collect::<Vec<_>>()
-        .try_into()
-        .unwrap())
+        Ok(Sha256::digest_cells(
+            table16_chip,
+            layouter.namespace(|| "'abc' * 2"),
+            &[],
+        )?)
     }
 }
 

diff --git a/examples/sha256/table16/message_schedule.rs b/examples/sha256/table16/message_schedule.rs
@@ -1,6 +1,8 @@
 use std::convert::TryInto;
 
-use super::{super::BLOCK_SIZE, AssignedBits, BlockWord, SpreadInputs, Table16Assignment, ROUNDS};
+use crate::BlockWord;
+
+use super::{super::BLOCK_SIZE, AssignedBits, SpreadInputs, Table16Assignment, ROUNDS};
 use halo2_proofs::{
     circuit::Layouter,
     plonk::{Advice, Column, ConstraintSystem, Error, Selector},
@@ -21,7 +23,7 @@ use schedule_util::*;
 pub use schedule_util::msg_schedule_test_input;
 
 #[derive(Clone, Debug)]
-pub(super) struct MessageWord(AssignedBits<32>);
+pub(super) struct MessageWord(pub AssignedBits<32>);
 
 impl std::ops::Deref for MessageWord {
     type Target = AssignedBits<32>;

diff --git a/examples/sha256/table16/mod.rs b/examples/sha256/table16/mod.rs
@@ -113,7 +113,7 @@ impl From<u32> for Bits<32> {
 }
 
 #[derive(Clone, Debug)]
-pub struct AssignedBits<const LEN: usize>(AssignedCell<Bits<LEN>, pallas::Base>);
+pub struct AssignedBits<const LEN: usize>(pub AssignedCell<Bits<LEN>, pallas::Base>);
 
 impl<const LEN: usize> std::ops::Deref for AssignedBits<LEN> {
     type Target = AssignedCell<Bits<LEN>, pallas::Base>;
@@ -354,9 +354,24 @@ impl Sha256Instructions<pallas::Base> for Table16Chip {
         layouter: &mut impl Layouter<pallas::Base>,
         initialized_state: &Self::State,
         input: [Self::BlockWord; super::BLOCK_SIZE],
+        input_cells: Option<[AssignedCell<pallas::Base, pallas::Base>; super::BLOCK_SIZE]>,
     ) -> Result<Self::State, Error> {
         let config = self.config();
-        let (_, w_halves) = config.message_schedule.process(layouter, input)?;
+
+        let (word_cells, w_halves) = config.message_schedule.process(layouter, input)?;
+
+        if let Some(input_cells) = input_cells {
+            word_cells
+                .into_iter()
+                .zip(input_cells.iter())
+                .try_for_each(|(cell, input)| {
+                    layouter.assign_region(
+                        || "check input word equality",
+                        |mut region| region.constrain_equal(cell.cell(), input.cell()),
+                    )
+                })?;
+        }
+
         config
             .compression
             .compress(layouter, initialized_state.clone(), w_halves)
@@ -379,7 +394,9 @@ impl Sha256Instructions<pallas::Base> for Table16Chip {
     ) -> Result<[AssignedCell<pallas::Base, pallas::Base>; super::DIGEST_SIZE], Error> {
         // Copy the dense forms of the state variable chunks down to this gate.
         // Reconstruct the 32-bit dense words.
-        self.config().compression.digest_cells(layouter, state.clone())
+        self.config()
+            .compression
+            .digest_cells(layouter, state.clone())
     }
 }