Skip to content

Commit

Permalink
migrate canonical representation
Browse files Browse the repository at this point in the history
  • Loading branch information
Mason Liang committed Oct 24, 2023
1 parent 1eb88e6 commit dc46a3e
Show file tree
Hide file tree
Showing 4 changed files with 122 additions and 76 deletions.
14 changes: 13 additions & 1 deletion src/constraint_builder/binary_column.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::{BinaryQuery, ConstraintBuilder, Query};
use super::{AdviceColumn, BinaryQuery, ConstraintBuilder, Query};
use crate::assignment_map::Column as ColumnEnum;
use halo2_proofs::{
arithmetic::FieldExt,
circuit::{Region, Value},
Expand Down Expand Up @@ -43,4 +44,15 @@ impl BinaryColumn {
.assign_advice(|| "binary", self.0, offset, || Value::known(F::from(value)))
.expect("failed assign_advice");
}

pub fn assignment<F: FieldExt>(
&self,
offset: usize,
value: bool,
) -> ((ColumnEnum, usize), Value<F>) {
(
(ColumnEnum::Advice(AdviceColumn(self.0)), offset),
Value::known(if value { F::one() } else { F::zero() }),
)
}
}
163 changes: 95 additions & 68 deletions src/gadgets/canonical_representation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use super::super::constraint_builder::{
SelectorColumn,
};
use super::{byte_bit::RangeCheck256Lookup, is_zero::IsZeroGadget, rlc_randomness::RlcRandomness};
use crate::assignment_map::Column;
use ethers_core::types::U256;
use halo2_proofs::{
arithmetic::{Field, FieldExt},
Expand All @@ -12,6 +13,7 @@ use halo2_proofs::{
};
use itertools::Itertools;
use num_traits::Zero;
use rayon::prelude::*;

pub trait CanonicalRepresentationLookup {
fn lookup<F: FieldExt>(&self) -> [Query<F>; 3];
Expand Down Expand Up @@ -134,82 +136,102 @@ impl CanonicalRepresentationConfig {
&self,
region: &mut Region<'_, Fr>,
randomness: Value<Fr>,
values: &[Fr],
values: Vec<Fr>,
n_rows: usize,
) {
let modulus = U256::from_str_radix(Fr::MODULUS, 16).unwrap();
let mut modulus_bytes = [0u8; 32];
modulus.to_big_endian(&mut modulus_bytes);

let mut offset = 1;
for value in values.iter() {
let mut bytes = value.to_bytes();
bytes.reverse();
let mut differences_are_zero_so_far = true;
let mut rlc = Value::known(Fr::zero());
for (index, (byte, modulus_byte)) in bytes.iter().zip_eq(&modulus_bytes).enumerate() {
self.byte.assign(region, offset, u64::from(*byte));
self.modulus_byte
.assign(region, offset, u64::from(*modulus_byte));

self.index
.assign(region, offset, u64::try_from(index).unwrap());
if index.is_zero() {
self.index_is_zero.enable(region, offset);
} else if index == 31 {
self.index_is_31.enable(region, offset);
let assignments: Vec<_> = self.assignments(values, n_rows, randomness).collect();
for ((column, offset), value) in assignments.into_iter() {
match column {
Column::Selector(s) => region.assign_fixed(|| "selector", s.0, offset, || value),
Column::Fixed(s) => region.assign_fixed(|| "fixed", s.0, offset, || value),
Column::Advice(s) => region.assign_advice(|| "advice", s.0, offset, || value),
Column::SecondPhaseAdvice(s) => {
region.assign_advice(|| "second phase advice", s.0, offset, || value)
}

let difference = Fr::from(u64::from(*modulus_byte)) - Fr::from(u64::from(*byte));
self.difference.assign(region, offset, difference);
self.difference_is_zero.assign(region, offset, difference);

self.differences_are_zero_so_far.assign(
region,
offset,
differences_are_zero_so_far,
);
differences_are_zero_so_far &= difference.is_zero_vartime();

self.value.assign(region, offset, *value);

rlc = rlc * randomness + Value::known(Fr::from(u64::from(*byte)));
self.rlc.assign(region, offset, rlc);

offset += 1
}
};
}
}

let expected_offset = Self::n_rows_required(values);
debug_assert!(
offset == expected_offset,
"assign used {offset} rows but {expected_offset} rows expected from `n_rows_required`",
);
pub fn assignments(
&self,
values: Vec<Fr>,
n_rows: usize,
randomness: Value<Fr>,
) -> impl ParallelIterator<Item = ((Column, usize), Value<Fr>)> + '_ {
let modulus = U256::from_str_radix(Fr::MODULUS, 16).unwrap();
let mut modulus_bytes = [0u8; 32];
modulus.to_big_endian(&mut modulus_bytes);

let n_padding_values = n_rows / 32 - values.len();
for _ in 0..n_padding_values {
for (index, modulus_byte) in modulus_bytes.iter().enumerate() {
self.modulus_byte
.assign(region, offset, u64::from(*modulus_byte));
let n_values = values.len();
values
.into_par_iter()
.enumerate()
.flat_map_iter(move |(i, value)| {
let mut assignments = vec![];
let mut offset = 1 + 32 * i;

let mut bytes = value.to_bytes();
bytes.reverse();
let mut differences_are_zero_so_far = true;
let mut rlc = Value::known(Fr::zero());
for (index, (byte, modulus_byte)) in bytes.iter().zip_eq(&modulus_bytes).enumerate()
{
let difference =
Fr::from(u64::from(*modulus_byte)) - Fr::from(u64::from(*byte));
rlc = rlc * randomness + Value::known(Fr::from(u64::from(*byte)));

assignments.extend([
self.byte.assignment(offset, u64::from(*byte)),
self.modulus_byte
.assignment(offset, u64::from(*modulus_byte)),
self.index.assignment(offset, u64::try_from(index).unwrap()),
self.differences_are_zero_so_far
.assignment(offset, differences_are_zero_so_far),
self.value.assignment(offset, value),
self.rlc.assignment(offset, rlc),
]);
assignments.extend(self.difference_is_zero.assignments(offset, difference));
if index.is_zero() {
assignments.push(self.index_is_zero.assignment(offset, true));
} else if index == 31 {
assignments.push(self.index_is_31.assignment(offset, true));
}

self.index
.assign(region, offset, u64::try_from(index).unwrap());
if index.is_zero() {
self.index_is_zero.enable(region, offset);
} else if index == 31 {
self.index_is_31.enable(region, offset);
differences_are_zero_so_far &= difference.is_zero_vartime();
offset += 1
}

let difference = Fr::from(u64::from(*modulus_byte));
self.difference.assign(region, offset, difference);
self.difference_is_zero.assign(region, offset, difference);

self.differences_are_zero_so_far
.assign(region, offset, index == 0);

offset += 1
}
}
assignments.into_iter()
})
.chain(
(n_values..n_rows / 32)
.into_par_iter()
.flat_map_iter(move |i| {
let mut assignments = vec![];
for (index, modulus_byte) in modulus_bytes.iter().enumerate() {
let offset = 1 + 32 * i + index;
assignments.extend([
self.modulus_byte
.assignment(offset, u64::from(*modulus_byte)),
self.index.assignment(offset, u64::try_from(index).unwrap()),
]);
assignments.extend(
self.difference_is_zero
.assignments(offset, u64::from(*modulus_byte)),
);

if index.is_zero() {
assignments.extend([
self.index_is_zero.assignment(offset, true),
self.differences_are_zero_so_far.assignment(offset, true),
]);
} else if index == 31 {
assignments.push(self.index_is_31.assignment(offset, true));
}
}
assignments.into_iter()
}),
)
}

pub fn n_rows_required(values: &[Fr]) -> usize {
Expand Down Expand Up @@ -290,7 +312,12 @@ mod test {
selector.enable(&mut region, offset);
}
byte_bit.assign(&mut region);
canonical_representation.assign(&mut region, randomness, &self.values, 256);
canonical_representation.assign(
&mut region,
randomness,
self.values.clone(),
256,
);
Ok(())
},
)
Expand Down
2 changes: 1 addition & 1 deletion src/gadgets/key_bit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ mod test {

key_bit.assign(&mut region, &self.lookups);
byte_bit.assign(&mut region);
canonical_representation.assign(&mut region, randomness, &keys, 256);
canonical_representation.assign(&mut region, randomness, keys.clone(), 256);
Ok(())
},
)
Expand Down
19 changes: 13 additions & 6 deletions src/mpt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,8 +143,6 @@ impl MptCircuitConfig {
keys.len()
);

self.canonical_representation
.assign(&mut region, randomness, &keys, n_rows);
self.key_bit.assign(&mut region, &key_bit_lookups(proofs));

let n_assigned_rows = self.mpt_update.assign(&mut region, proofs, randomness);
Expand All @@ -164,13 +162,22 @@ impl MptCircuitConfig {
},
);

let mut keys = mpt_update_keys(proofs);
keys.sort();
keys.dedup();
layouter.assign_regions(
|| "mpt circuit parallel assignment",
AssignmentMap::new(
self.byte_bit.assignments().chain(
self.byte_representation
.assignments(u32s, u64s, u128s, frs, randomness),
),
self.byte_bit
.assignments()
.chain(
self.byte_representation
.assignments(u32s, u64s, u128s, frs, randomness),
)
.chain(
self.canonical_representation
.assignments(keys, n_rows, randomness),
),
)
.assignments(),
)?;
Expand Down

0 comments on commit dc46a3e

Please sign in to comment.