From 5951ae6516f7aa3251abec68d74a7b8293be1877 Mon Sep 17 00:00:00 2001 From: therealyingtong Date: Fri, 8 Jan 2021 16:10:55 +0800 Subject: [PATCH] Add Compression --- src/gadget/sha256/table16.rs | 89 ++-- src/gadget/sha256/table16/compression.rs | 401 ++++++++++++++++++ src/gadget/sha256/table16/message_schedule.rs | 15 +- src/gadget/sha256/table16/spread_table.rs | 11 +- 4 files changed, 448 insertions(+), 68 deletions(-) create mode 100644 src/gadget/sha256/table16/compression.rs diff --git a/src/gadget/sha256/table16.rs b/src/gadget/sha256/table16.rs index 3542792..433ec83 100644 --- a/src/gadget/sha256/table16.rs +++ b/src/gadget/sha256/table16.rs @@ -7,11 +7,13 @@ use crate::{ plonk::{Advice, Column, ConstraintSystem, Error, Permutation}, }; +mod compression; mod gates; mod message_schedule; mod spread_table; mod util; +use compression::*; use gates::*; use message_schedule::*; use spread_table::*; @@ -93,71 +95,12 @@ impl Into for CellValue16 { } } -/// A variable that represents the `[A,B,C,D]` words of the SHA-256 internal state. -/// -/// The structure of this variable is influenced by the following factors: -/// - In `Σ_0(A)` we need `A` to be split into pieces `(a,b,c,d)` of lengths `(2,11,9,10)` -/// bits respectively (counting from the little end), as well as their spread forms. -/// - `Maj(A,B,C)` requires having the bits of each input in spread form. For `A` we can -/// reuse the pieces from `Σ_0(A)`. Since `B` and `C` are assigned from `A` and `B` -/// respectively in each round, we therefore also have the same pieces in earlier rows. -/// We align the columns to make it efficient to copy-constrain these forms where they -/// are needed. -#[derive(Copy, Clone, Debug)] -pub struct AbcdVar { - idx: i32, - val: u32, - a: SpreadVar, - b: SpreadVar, - c_lo: SpreadVar, - c_mid: SpreadVar, - c_hi: SpreadVar, - d: SpreadVar, -} - -/// A variable that represents the `[E,F,G,H]` words of the SHA-256 internal state. -/// -/// The structure of this variable is influenced by the following factors: -/// - In `Σ_1(E)` we need `E` to be split into pieces `(a,b,c,d)` of lengths `(6,5,14,7)` -/// bits respectively (counting from the little end), as well as their spread forms. -/// - `Ch(E,F,G)` requires having the bits of each input in spread form. For `E` we can -/// reuse the pieces from `Σ_1(E)`. Since `F` and `G` are assigned from `E` and `F` -/// respectively in each round, we therefore also have the same pieces in earlier rows. -/// We align the columns to make it efficient to copy-constrain these forms where they -/// are needed. -#[derive(Copy, Clone, Debug)] -pub struct EfghVar { - idx: i32, - val: u32, - a_lo: SpreadVar, - a_hi: SpreadVar, - b_lo: SpreadVar, - b_hi: SpreadVar, - c: SpreadVar, - d: SpreadVar, -} - -/// The internal state for SHA-256. -#[derive(Clone, Debug)] -pub struct State { - h_0: AbcdVar, - h_1: AbcdVar, - h_2: AbcdVar, - h_3: AbcdVar, - h_4: EfghVar, - h_5: EfghVar, - h_6: EfghVar, - h_7: EfghVar, -} - -#[derive(Clone, Debug)] -struct HPrime {} - /// Configuration for a [`Table16Chip`]. #[derive(Clone, Debug)] pub struct Table16Config { lookup_table: SpreadTable, message_schedule: MessageSchedule, + compression: Compression, } /// A chip that implements SHA-256 with a maximum lookup table size of $2^16$. @@ -213,12 +156,21 @@ impl Table16Chip { ], ); + let compression = Compression::configure( + meta, + lookup_inputs.clone(), + message_schedule, + extras, + perm.clone(), + ); + let message_schedule = MessageSchedule::configure(meta, lookup_inputs, message_schedule, extras, perm.clone()); Table16Config { lookup_table, message_schedule, + compression, } } } @@ -243,16 +195,22 @@ impl Sha256Instructions for Table16Chip { } fn initialization_vector(layouter: &mut impl Layouter) -> Result { - todo!() + let config = layouter.config().clone(); + config.compression.initialize_with_iv(layouter, IV) } fn initialization( layouter: &mut impl Layouter>, init_state: &Self::State, ) -> Result { - todo!() + let config = layouter.config().clone(); + config + .compression + .initialize_with_state(layouter, init_state.clone()) } + // Given an initialized state and an input message block, compress the + // message block and return the final state. fn compress( layouter: &mut impl Layouter, initialized_state: &Self::State, @@ -261,7 +219,9 @@ impl Sha256Instructions for Table16Chip { let config = layouter.config().clone(); let (_, w_halves) = config.message_schedule.process(layouter, input)?; - todo!() + config + .compression + .compress(layouter, initialized_state.clone(), w_halves) } fn digest( @@ -270,7 +230,8 @@ impl Sha256Instructions for Table16Chip { ) -> Result<[Self::BlockWord; super::DIGEST_SIZE], Error> { // Copy the dense forms of the state variable chunks down to this gate. // Reconstruct the 32-bit dense words. - todo!() + let config = layouter.config().clone(); + config.compression.digest(layouter, state.clone()) } } diff --git a/src/gadget/sha256/table16/compression.rs b/src/gadget/sha256/table16/compression.rs new file mode 100644 index 0000000..adbbf62 --- /dev/null +++ b/src/gadget/sha256/table16/compression.rs @@ -0,0 +1,401 @@ +use super::{ + super::DIGEST_SIZE, BlockWord, CellValue16, CellValue32, SpreadInputs, SpreadVar, + Table16Assignment, Table16Chip, ROUNDS, STATE, +}; +use crate::{ + arithmetic::FieldExt, + circuit::Layouter, + plonk::{Advice, Column, ConstraintSystem, Error, Fixed, Permutation}, +}; + +// mod compression_gates; +// mod compression_util; +// mod subregion_digest; +// mod subregion_initial; +// mod subregion_main; + +// use compression_gates::CompressionGate; + +/// A variable that represents the `[A,B,C,D]` words of the SHA-256 internal state. +/// +/// The structure of this variable is influenced by the following factors: +/// - In `Σ_0(A)` we need `A` to be split into pieces `(a,b,c,d)` of lengths `(2,11,9,10)` +/// bits respectively (counting from the little end), as well as their spread forms. +/// - `Maj(A,B,C)` requires having the bits of each input in spread form. For `A` we can +/// reuse the pieces from `Σ_0(A)`. Since `B` and `C` are assigned from `A` and `B` +/// respectively in each round, we therefore also have the same pieces in earlier rows. +/// We align the columns to make it efficient to copy-constrain these forms where they +/// are needed. +#[derive(Copy, Clone, Debug)] +pub struct AbcdVar { + idx: i32, + val: u32, + a: SpreadVar, + b: SpreadVar, + c_lo: SpreadVar, + c_mid: SpreadVar, + c_hi: SpreadVar, + d: SpreadVar, +} + +/// A variable that represents the `[E,F,G,H]` words of the SHA-256 internal state. +/// +/// The structure of this variable is influenced by the following factors: +/// - In `Σ_1(E)` we need `E` to be split into pieces `(a,b,c,d)` of lengths `(6,5,14,7)` +/// bits respectively (counting from the little end), as well as their spread forms. +/// - `Ch(E,F,G)` requires having the bits of each input in spread form. For `E` we can +/// reuse the pieces from `Σ_1(E)`. Since `F` and `G` are assigned from `E` and `F` +/// respectively in each round, we therefore also have the same pieces in earlier rows. +/// We align the columns to make it efficient to copy-constrain these forms where they +/// are needed. +#[derive(Copy, Clone, Debug)] +pub struct EfghVar { + idx: i32, + val: u32, + a_lo: SpreadVar, + a_hi: SpreadVar, + b_lo: SpreadVar, + b_hi: SpreadVar, + c: SpreadVar, + d: SpreadVar, +} + +#[derive(Clone, Debug)] +pub struct RoundWordDense { + dense_halves: (CellValue16, CellValue16), +} + +impl RoundWordDense { + pub fn new(dense_halves: (CellValue16, CellValue16)) -> Self { + RoundWordDense { dense_halves } + } +} + +#[derive(Clone, Debug)] +pub struct RoundWordSpread { + dense_halves: (CellValue16, CellValue16), + spread_halves: (CellValue32, CellValue32), +} + +impl RoundWordSpread { + pub fn new( + dense_halves: (CellValue16, CellValue16), + spread_halves: (CellValue32, CellValue32), + ) -> Self { + RoundWordSpread { + dense_halves, + spread_halves, + } + } +} + +impl Into for RoundWordSpread { + fn into(self) -> RoundWordDense { + RoundWordDense::new(self.dense_halves) + } +} + +#[derive(Clone, Debug)] +pub struct RoundWordA { + pieces: Option, + dense_halves: (CellValue16, CellValue16), + spread_halves: Option<(CellValue32, CellValue32)>, +} + +impl RoundWordA { + pub fn new( + pieces: AbcdVar, + dense_halves: (CellValue16, CellValue16), + spread_halves: (CellValue32, CellValue32), + ) -> Self { + RoundWordA { + pieces: Some(pieces), + dense_halves, + spread_halves: Some(spread_halves), + } + } + + pub fn new_dense(dense_halves: (CellValue16, CellValue16)) -> Self { + RoundWordA { + pieces: None, + dense_halves, + spread_halves: None, + } + } +} + +impl Into for RoundWordA { + fn into(self) -> RoundWordSpread { + RoundWordSpread::new(self.dense_halves, self.spread_halves.unwrap()) + } +} + +#[derive(Clone, Debug)] +pub struct RoundWordE { + pieces: Option, + dense_halves: (CellValue16, CellValue16), + spread_halves: Option<(CellValue32, CellValue32)>, +} + +impl RoundWordE { + pub fn new( + pieces: EfghVar, + dense_halves: (CellValue16, CellValue16), + spread_halves: (CellValue32, CellValue32), + ) -> Self { + RoundWordE { + pieces: Some(pieces), + dense_halves, + spread_halves: Some(spread_halves), + } + } + + pub fn new_dense(dense_halves: (CellValue16, CellValue16)) -> Self { + RoundWordE { + pieces: None, + dense_halves, + spread_halves: None, + } + } +} + +impl Into for RoundWordE { + fn into(self) -> RoundWordSpread { + RoundWordSpread::new(self.dense_halves, self.spread_halves.unwrap()) + } +} + +/// The internal state for SHA-256. +#[derive(Clone, Debug)] +pub struct State { + a: Option, + b: Option, + c: Option, + d: Option, + e: Option, + f: Option, + g: Option, + h: Option, +} + +impl State { + pub fn new( + a: StateWord, + b: StateWord, + c: StateWord, + d: StateWord, + e: StateWord, + f: StateWord, + g: StateWord, + h: StateWord, + ) -> Self { + State { + a: Some(a), + b: Some(b), + c: Some(c), + d: Some(d), + e: Some(e), + f: Some(f), + g: Some(g), + h: Some(h), + } + } + + pub fn empty_state() -> Self { + State { + a: None, + b: None, + c: None, + d: None, + e: None, + f: None, + g: None, + h: None, + } + } +} + +#[derive(Clone, Debug)] +pub enum StateWord { + A(RoundWordA), + B(RoundWordSpread), + C(RoundWordSpread), + D(RoundWordDense), + E(RoundWordE), + F(RoundWordSpread), + G(RoundWordSpread), + H(RoundWordDense), +} + +#[derive(Clone, Debug)] +pub(super) struct Compression { + lookup: SpreadInputs, + message_schedule: Column, + extras: [Column; 6], + + s_ch: Column, + s_ch_neg: Column, + s_maj: Column, + s_h_prime: Column, + s_a_new: Column, + s_e_new: Column, + + s_upper_sigma_0: Column, + s_upper_sigma_1: Column, + + // Decomposition gate for AbcdVar + s_decompose_abcd: Column, + // Decomposition gate for EfghVar + s_decompose_efgh: Column, + + s_digest: Column, + + perm: Permutation, +} + +impl Table16Assignment for Compression {} + +impl Compression { + pub(super) fn configure( + meta: &mut ConstraintSystem, + lookup: SpreadInputs, + message_schedule: Column, + extras: [Column; 6], + perm: Permutation, + ) -> Self { + let s_ch = meta.fixed_column(); + let s_ch_neg = meta.fixed_column(); + let s_maj = meta.fixed_column(); + let s_h_prime = meta.fixed_column(); + let s_a_new = meta.fixed_column(); + let s_e_new = meta.fixed_column(); + + let s_upper_sigma_0 = meta.fixed_column(); + let s_upper_sigma_1 = meta.fixed_column(); + + // Decomposition gate for AbcdVar + let s_decompose_abcd = meta.fixed_column(); + // Decomposition gate for EfghVar + let s_decompose_efgh = meta.fixed_column(); + + let s_digest = meta.fixed_column(); + + // Rename these here for ease of matching the gates to the specification. + let a_0 = lookup.tag; + let a_1 = lookup.dense; + let a_2 = lookup.spread; + let a_3 = extras[0]; + let a_4 = extras[1]; + let a_5 = message_schedule; + let a_6 = extras[2]; + let a_7 = extras[3]; + let a_8 = extras[4]; + let a_9 = extras[5]; + + // TODO: Create gates. + + Compression { + lookup, + message_schedule, + extras, + s_ch, + s_ch_neg, + s_maj, + s_h_prime, + s_a_new, + s_e_new, + s_upper_sigma_0, + s_upper_sigma_1, + s_decompose_abcd, + s_decompose_efgh, + s_digest, + perm, + } + } + + /// Initialize compression with a constant Initialization Vector of 32-byte words. + /// Returns an initialized state. + pub(super) fn initialize_with_iv( + &self, + layouter: &mut impl Layouter>, + init_state: [u32; STATE], + ) -> Result { + let mut new_state = State::empty_state(); + todo!() + } + + /// Initialize compression with some initialized state. This could be a state + /// output from a previous compression round. + pub(super) fn initialize_with_state( + &self, + layouter: &mut impl Layouter>, + init_state: State, + ) -> Result { + let mut new_state = State::empty_state(); + todo!() + } + + /// Given an initialized state and a message schedule, perform 64 compression rounds. + pub(super) fn compress( + &self, + layouter: &mut impl Layouter>, + initialized_state: State, + w_halves: [(CellValue16, CellValue16); ROUNDS], + ) -> Result { + let mut state = State::empty_state(); + todo!() + } + + /// After the final round, convert the state into the final digest. + pub(super) fn digest( + &self, + layouter: &mut impl Layouter>, + state: State, + ) -> Result<[BlockWord; DIGEST_SIZE], Error> { + let mut digest = [BlockWord::new(0); DIGEST_SIZE]; + todo!() + } + + pub(super) fn empty_configure( + meta: &mut ConstraintSystem, + lookup: SpreadInputs, + message_schedule: Column, + extras: [Column; 6], + perm: Permutation, + ) -> Self { + let s_ch = meta.fixed_column(); + let s_ch_neg = meta.fixed_column(); + let s_maj = meta.fixed_column(); + let s_h_prime = meta.fixed_column(); + let s_a_new = meta.fixed_column(); + let s_e_new = meta.fixed_column(); + + let s_upper_sigma_0 = meta.fixed_column(); + let s_upper_sigma_1 = meta.fixed_column(); + + // Decomposition gate for AbcdVar + let s_decompose_abcd = meta.fixed_column(); + // Decomposition gate for EfghVar + let s_decompose_efgh = meta.fixed_column(); + + let s_digest = meta.fixed_column(); + + Compression { + lookup, + message_schedule, + extras, + s_ch, + s_ch_neg, + s_maj, + s_h_prime, + s_a_new, + s_e_new, + s_upper_sigma_0, + s_upper_sigma_1, + s_decompose_abcd, + s_decompose_efgh, + s_digest, + perm, + } + } +} diff --git a/src/gadget/sha256/table16/message_schedule.rs b/src/gadget/sha256/table16/message_schedule.rs index febb0b3..e6feae2 100644 --- a/src/gadget/sha256/table16/message_schedule.rs +++ b/src/gadget/sha256/table16/message_schedule.rs @@ -474,7 +474,9 @@ impl MessageSchedule { #[cfg(test)] mod tests { - use super::super::{super::BLOCK_SIZE, BlockWord, SpreadTable, Table16Chip, Table16Config}; + use super::super::{ + super::BLOCK_SIZE, BlockWord, Compression, SpreadTable, Table16Chip, Table16Config, + }; use super::{schedule_util::*, MessageSchedule}; use crate::{ arithmetic::FieldExt, @@ -534,12 +536,21 @@ mod tests { ], ); + let compression = Compression::empty_configure( + meta, + lookup_inputs.clone(), + message_schedule, + extras, + perm.clone(), + ); + let message_schedule = MessageSchedule::configure(meta, lookup_inputs, message_schedule, extras, perm); Table16Config { lookup_table, message_schedule, + compression, } } @@ -554,8 +565,6 @@ mod tests { let table = layouter.config().lookup_table.clone(); table.load(&mut layouter)?; - // layouter.assign_region() - // Provide input // Test vector: "abc" let inputs: [BlockWord; BLOCK_SIZE] = get_msg_schedule_test_input(); diff --git a/src/gadget/sha256/table16/spread_table.rs b/src/gadget/sha256/table16/spread_table.rs index cce299d..41f1e63 100644 --- a/src/gadget/sha256/table16/spread_table.rs +++ b/src/gadget/sha256/table16/spread_table.rs @@ -245,7 +245,7 @@ mod tests { use std::marker::PhantomData; use super::{ - super::{util::*, MessageSchedule, Table16Chip, Table16Config}, + super::{util::*, Compression, MessageSchedule, Table16Chip, Table16Config}, SpreadInputs, SpreadTable, }; use crate::{ @@ -514,6 +514,14 @@ mod tests { ], ); + let compression = Compression::empty_configure( + meta, + lookup_inputs.clone(), + message_schedule, + extras, + perm.clone(), + ); + let message_schedule = MessageSchedule::empty_configure( meta, lookup_inputs.clone(), @@ -527,6 +535,7 @@ mod tests { sha256: Table16Config { lookup_table, message_schedule, + compression, }, } }