Impl StepCircuit for sha256 gadget in example

examples/sha256/table16/compression.rs

@@ -1,16 +1,18 @@
-use super::{
-    super::DIGEST_SIZE,
-    util::{i2lebsp, lebs2ip},
-    AssignedBits, BlockWord, SpreadInputs, SpreadVar, Table16Assignment, ROUNDS, STATE,
-};
+use std::convert::TryInto;
+use std::ops::Range;
+
 use halo2_proofs::{
     circuit::{Layouter, Value},
     plonk::{Advice, Column, ConstraintSystem, Error, Selector},
     poly::Rotation,
 };
 use halo2curves::pasta::pallas;
-use std::convert::TryInto;
-use std::ops::Range;
+
+use super::{
+    util::{i2lebsp, lebs2ip},
+    AssignedBits, BlockWord, SpreadInputs, SpreadVar, Table16Assignment, ROUNDS, STATE,
+};
+use crate::{AssignedCell, DIGEST_SIZE};
 
 mod compression_gates;
 mod compression_util;
@@ -933,6 +935,18 @@ impl CompressionConfig {
         )?;
         Ok(digest)
     }
+
+    /// After the final round, convert the state into the final digest.
+    pub(super) fn digest_cells(
+        &self,
+        layouter: &mut impl Layouter<pallas::Base>,
+        state: State,
+    ) -> Result<[AssignedCell<pallas::Base, pallas::Base>; DIGEST_SIZE], Error> {
+        layouter.assign_region(
+            || "digest",
+            |mut region| self.assign_digest_cells(&mut region, state.clone()),
+        )
+    }
 }
 
 #[cfg(test)]

examples/sha256/table16/compression/compression_util.rs

@@ -2,9 +2,7 @@ use super::{
     AbcdVar, CompressionConfig, EfghVar, RoundWord, RoundWordA, RoundWordDense, RoundWordE,
     RoundWordSpread, State, UpperSigmaVar,
 };
-use crate::table16::{
-    util::*, AssignedBits, SpreadVar, SpreadWord, StateWord, Table16Assignment,
-};
+use crate::table16::{util::*, AssignedBits, SpreadVar, SpreadWord, StateWord, Table16Assignment};
 use halo2_proofs::{
     circuit::{Region, Value},
     plonk::{Advice, Column, Error},

examples/sha256/table16/compression/subregion_digest.rs

@@ -1,11 +1,17 @@
-use super::super::{super::DIGEST_SIZE, BlockWord, RoundWordDense};
-use super::{compression_util::*, CompressionConfig, State};
 use halo2_proofs::{
     circuit::{Region, Value},
     plonk::{Advice, Column, Error},
 };
 use halo2curves::pasta::pallas;
 
+use crate::AssignedCell;
+
+use super::{
+    super::{BlockWord, RoundWordDense},
+    {compression_util::*, CompressionConfig, State},
+};
+use crate::DIGEST_SIZE;
+
 impl CompressionConfig {
     #[allow(clippy::many_single_char_names)]
     pub fn assign_digest(
@@ -77,6 +83,93 @@ impl CompressionConfig {
         ])
     }
 
+    pub fn assign_digest_cells(
+        &self,
+        region: &mut Region<'_, pallas::Base>,
+        state: State,
+    ) -> Result<[AssignedCell<pallas::Base, pallas::Base>; DIGEST_SIZE], Error> {
+        let a_3 = self.extras[0];
+        let a_4 = self.extras[1];
+        let a_5 = self.message_schedule;
+        let a_6 = self.extras[2];
+        let a_7 = self.extras[3];
+        let a_8 = self.extras[4];
+
+        let (a, b, c, d, e, f, g, h) = match_state(state);
+
+        let abcd_row = 0;
+        self.s_digest.enable(region, abcd_row)?;
+        let efgh_row = abcd_row + 2;
+        self.s_digest.enable(region, efgh_row)?;
+
+        // Assign digest for A, B, C, D
+        a.dense_halves
+            .0
+            .copy_advice(|| "a_lo", region, a_3, abcd_row)?;
+        a.dense_halves
+            .1
+            .copy_advice(|| "a_hi", region, a_4, abcd_row)?;
+        let a = a.dense_halves.value();
+        let a_cell = region.assign_advice(
+            || "a",
+            a_5,
+            abcd_row,
+            || a.map(|a| pallas::Base::from(a as u64)),
+        )?;
+
+        let b_cell =
+            self.assign_digest_word_cell(region, abcd_row, a_6, a_7, a_8, b.dense_halves)?;
+        let c_cell =
+            self.assign_digest_word_cell(region, abcd_row + 1, a_3, a_4, a_5, c.dense_halves)?;
+        let d_cell = self.assign_digest_word_cell(region, abcd_row + 1, a_6, a_7, a_8, d)?;
+
+        // Assign digest for E, F, G, H
+        e.dense_halves
+            .0
+            .copy_advice(|| "e_lo", region, a_3, efgh_row)?;
+        e.dense_halves
+            .1
+            .copy_advice(|| "e_hi", region, a_4, efgh_row)?;
+        let e = e.dense_halves.value();
+        let e_cell = region.assign_advice(
+            || "e",
+            a_5,
+            efgh_row,
+            || e.map(|e| pallas::Base::from(e as u64)),
+        )?;
+
+        let f_cell =
+            self.assign_digest_word_cell(region, efgh_row, a_6, a_7, a_8, f.dense_halves)?;
+        let g_cell =
+            self.assign_digest_word_cell(region, efgh_row + 1, a_3, a_4, a_5, g.dense_halves)?;
+        let h_cell = self.assign_digest_word_cell(region, efgh_row + 1, a_6, a_7, a_8, h)?;
+
+        Ok([
+            a_cell, b_cell, c_cell, d_cell, e_cell, f_cell, g_cell, h_cell,
+        ])
+    }
+
+    fn assign_digest_word_cell(
+        &self,
+        region: &mut Region<'_, pallas::Base>,
+        row: usize,
+        lo_col: Column<Advice>,
+        hi_col: Column<Advice>,
+        word_col: Column<Advice>,
+        dense_halves: RoundWordDense,
+    ) -> Result<AssignedCell<pallas::Base, pallas::Base>, Error> {
+        dense_halves.0.copy_advice(|| "lo", region, lo_col, row)?;
+        dense_halves.1.copy_advice(|| "hi", region, hi_col, row)?;
+
+        let val = dense_halves.value();
+        region.assign_advice(
+            || "word",
+            word_col,
+            row,
+            || val.map(|val| pallas::Base::from(val as u64)),
+        )
+    }
+
     fn assign_digest_word(
         &self,
         region: &mut Region<'_, pallas::Base>,

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>;

examples/sha256/table16/message_schedule/schedule_util.rs

@@ -44,10 +44,7 @@ pub fn get_word_row(word_idx: usize) -> usize {
     } else if (14..=48).contains(&word_idx) {
         SUBREGION_0_ROWS + SUBREGION_1_ROWS + SUBREGION_2_WORD * (word_idx - 14) + 1
     } else if (49..=61).contains(&word_idx) {
-        SUBREGION_0_ROWS
-            + SUBREGION_1_ROWS
-            + SUBREGION_2_ROWS
-            + SUBREGION_3_WORD * (word_idx - 49)
+        SUBREGION_0_ROWS + SUBREGION_1_ROWS + SUBREGION_2_ROWS + SUBREGION_3_WORD * (word_idx - 49)
     } else {
         SUBREGION_0_ROWS
             + SUBREGION_1_ROWS

examples/sha256/table16/message_schedule/subregion1.rs

@@ -7,7 +7,6 @@ use halo2_proofs::{
 use halo2curves::pasta::pallas;
 use std::convert::TryInto;
 
-#[allow(dead_code)]
 // A word in subregion 1
 // (3, 4, 11, 14)-bit chunks
 #[derive(Debug)]

examples/sha256/table16/message_schedule/subregion2.rs

@@ -7,7 +7,6 @@ use halo2_proofs::{
 use halo2curves::pasta::pallas;
 use std::convert::TryInto;
 
-#[allow(dead_code)]
 /// A word in subregion 2
 /// (3, 4, 3, 7, 1, 1, 13)-bit chunks
 #[derive(Clone, Debug)]

examples/sha256/table16.rs → examples/sha256/table16/mod.rs

@@ -1,7 +1,7 @@
 use std::convert::TryInto;
 use std::marker::PhantomData;
 
-use super::Sha256Instructions;
+use super::sha256::Sha256Instructions;
 use halo2_proofs::{
     circuit::{AssignedCell, Chip, Layouter, Region, Value},
     plonk::{Advice, Any, Assigned, Column, ConstraintSystem, Error},
@@ -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,22 @@ impl Sha256Instructions<pallas::Base> for Table16Chip {
         layouter: &mut impl Layouter<pallas::Base>,
         initialized_state: &Self::State,
         input: [Self::BlockWord; super::BLOCK_SIZE],
+        input_cells: [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)?;
+
+        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)
@@ -371,6 +384,18 @@ impl Sha256Instructions<pallas::Base> for Table16Chip {
         // Reconstruct the 32-bit dense words.
         self.config().compression.digest(layouter, state.clone())
     }
+
+    fn digest_cells(
+        &self,
+        layouter: &mut impl Layouter<pallas::Base>,
+        state: &Self::State,
+    ) -> 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())
+    }
 }
 
 /// Common assignment patterns used by Table16 regions.

examples/sha256/table16/spread_table.rs

@@ -66,7 +66,6 @@ impl<const DENSE: usize, const SPREAD: usize> SpreadWord<DENSE, SPREAD> {
     }
 }
 
-#[allow(dead_code)]
 /// A variable stored in advice columns corresponding to a row of [`SpreadTableConfig`].
 #[derive(Clone, Debug)]
 pub(super) struct SpreadVar<const DENSE: usize, const SPREAD: usize> {

src/step_circuit.rs

@@ -1,7 +1,7 @@
 use ff::PrimeField;
 use halo2_proofs::{
-    circuit::{AssignedCell, Region},
-    plonk::{Circuit, ConstraintSystem},
+    circuit::{AssignedCell, Layouter},
+    plonk::ConstraintSystem,
 };
 
 #[derive(Debug, thiserror::Error)]
@@ -34,7 +34,7 @@ pub enum SynthesisError {
 /// - For a detailed understanding of IVC and the context in which a trait
 ///   `StepCircuit` might be used, refer to the 'Section 5' of
 ///   [Nova Whitepaper](https://eprint.iacr.org/2023/969.pdf).
-pub trait StepCircuit<const ARITY: usize, F: PrimeField>: Circuit<F> {
+pub trait StepCircuit<const ARITY: usize, F: PrimeField> {
     type StepConfig: Clone;
 
     /// Configure the step circuit. This method initializes necessary
@@ -51,8 +51,8 @@ pub trait StepCircuit<const ARITY: usize, F: PrimeField>: Circuit<F> {
     /// Return `z_out` result
     fn synthesize(
         &self,
-        config: Self::Config,
-        region: Region<'_, F>,
+        config: Self::StepConfig,
+        layouter: &mut impl Layouter<F>,
         z_in: &[AssignedCell<F, F>; ARITY],
     ) -> Result<[AssignedCell<F, F>; ARITY], SynthesisError>;