From 8ca8406aecd5436ec4d6ec3b73fd677ab7c87b8a Mon Sep 17 00:00:00 2001 From: parazyd Date: Thu, 24 Mar 2022 18:59:14 +0100 Subject: [PATCH] zk: Implement EvenBits chip for bit decomposition. --- src/zk/even_bits.rs | 262 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 262 insertions(+) create mode 100644 src/zk/even_bits.rs diff --git a/src/zk/even_bits.rs b/src/zk/even_bits.rs new file mode 100644 index 000000000..595d92a3b --- /dev/null +++ b/src/zk/even_bits.rs @@ -0,0 +1,262 @@ +use std::{marker::PhantomData, ops::Deref}; + +use halo2_proofs::{ + arithmetic::FieldExt, + circuit::{AssignedCell, Chip, Layouter, Region}, + plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector, TableColumn}, + poly::Rotation, +}; + +/// Chip state is stored in a config struct. This is generated by the +/// chip during configuration, and then stored inside the chip +#[derive(Clone, Debug)] +pub struct EvenBitsConfig { + advice: [Column; 2], + even_bits: TableColumn, + + s_decompose: Selector, +} + +impl EvenBitsConfig { + pub fn load_private( + &self, + mut layouter: impl Layouter, + value: Option, + ) -> Result, Error> { + layouter.assign_region( + || "load private", + |mut region| { + region.assign_advice( + || "private input", + self.advice[0], + 0, + || value.ok_or(Error::Synthesis), + ) + }, + ) + } +} + +#[derive(Clone, Debug)] +pub struct EvenBitsChip { + config: EvenBitsConfig, + _marker: PhantomData, +} + +impl Chip for EvenBitsChip { + type Config = EvenBitsConfig; + type Loaded = (); + + fn config(&self) -> &Self::Config { + &self.config + } + + fn loaded(&self) -> &Self::Loaded { + &() + } +} + +impl EvenBitsChip { + pub fn construct(config: >::Config) -> Self { + Self { config, _marker: PhantomData } + } + + pub fn configure(meta: &mut ConstraintSystem) -> >::Config { + let advice = [meta.advice_column(), meta.advice_column()]; + for column in &advice { + meta.enable_equality(*column); + } + + let s_decompose = meta.complex_selector(); + let even_bits = meta.lookup_table_column(); + + meta.create_gate("decompose", |meta| { + let lhs = meta.query_advice(advice[0], Rotation::cur()); + let rhs = meta.query_advice(advice[1], Rotation::cur()); + let out = meta.query_advice(advice[0], Rotation::next()); + let s_decompose = meta.query_selector(s_decompose); + + // Finally, we return the polynomial expressions that constrain this gate. + // For our multiplication gate, we only need a single polynomial constraint. + // + // The polynomial expressions returned from `create_gate` will be + // constrained by the proving system to equal zero. + vec![s_decompose * (lhs + Expression::Constant(F::from(2)) * rhs - out)] + }); + + let _ = meta.lookup(|meta| { + let lookup = meta.query_selector(s_decompose); + let a = meta.query_advice(advice[0], Rotation::cur()); + + vec![(lookup * a, even_bits)] + }); + + let _ = meta.lookup(|meta| { + let lookup = meta.query_selector(s_decompose); + let b = meta.query_advice(advice[1], Rotation::cur()); + + vec![(lookup * b, even_bits)] + }); + + EvenBitsConfig { advice, even_bits, s_decompose } + } + + // Allocates all even bits in a table for the word size WORD_BITS. + // `2^(WORD_BITS/2)` rows of the constraint system + pub fn alloc_table(&self, layouter: &mut impl Layouter) -> Result<(), Error> { + layouter.assign_table( + || "even bits table", + |mut table| { + for i in 0..2usize.pow(WORD_BITS / 2) { + table.assign_cell( + || format!("even_bits row {}", i), + self.config.even_bits, + i, + || Ok(F::from(even_bits_at(i) as u64)), + )?; + } + Ok(()) + }, + ) + } +} + +fn even_bits_at(mut i: usize) -> usize { + let mut r = 0; + let mut c = 0; + + while i != 0 { + let lower_bit = i % 2; + r += lower_bit * 4usize.pow(c); + i >>= 1; + c += 1; + } + + r +} + +/// A newtype of a field element containing only bits that were in the +/// even position of the decomposed element. +/// All odd bits will be zero. +#[derive(Clone, Copy, Debug)] +pub struct EvenBits(pub W); + +impl Deref for EvenBits { + type Target = W; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +/// A newtype of a field element containing only bits thet were in the +/// odd position of the decomposed element. +/// All odd bits will be right shifted by 1 into even positions. +/// All odd bits will be zero. +#[derive(Clone, Copy, Debug)] +pub struct OddBits(pub W); + +impl Deref for OddBits { + type Target = W; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +pub trait EvenBitsLookup: Chip { + type Word; + + #[allow(clippy::type_complexity)] + fn decompose( + &self, + layouter: impl Layouter, + c: Self::Word, + ) -> Result<(EvenBits, OddBits), Error>; +} + +impl EvenBitsLookup for EvenBitsChip { + type Word = AssignedCell; + + fn decompose( + &self, + mut layouter: impl Layouter, + c: Self::Word, + ) -> Result<(EvenBits, OddBits), Error> { + let config = self.config(); + + layouter.assign_region( + || "decompose", + |mut region: Region<'_, F>| { + config.s_decompose.enable(&mut region, 0)?; + + let o_eo = c.value().cloned().map(decompose); + let e_cell = region + .assign_advice( + || "even bits", + config.advice[0], + 0, + || o_eo.map(|eo| *eo.0).ok_or(Error::Synthesis), + ) + .map(EvenBits)?; + + let o_cell = region + .assign_advice( + || "odd bits", + config.advice[1], + 0, + || o_eo.map(|eo| *eo.1).ok_or(Error::Synthesis), + ) + .map(OddBits)?; + + c.copy_advice(|| "out", &mut region, config.advice[0], 1)?; + Ok((e_cell, o_cell)) + }, + ) + } +} + +fn decompose(word: F) -> (EvenBits, OddBits) { + assert!(word <= F::from_u128(u128::MAX)); + + let mut even_only = word.to_repr(); + even_only.as_mut().iter_mut().for_each(|bits| { + *bits &= 0b01010101; + }); + + let mut odd_only = word.to_repr(); + odd_only.as_mut().iter_mut().for_each(|bits| { + *bits &= 0b10101010; + }); + + let even_only = EvenBits(F::from_repr(even_only).unwrap()); + let odd_only = F::from_repr(odd_only).unwrap(); + let odds_in_even = OddBits(F::from_u128(odd_only.get_lower_128() >> 1)); + (even_only, odds_in_even) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn even_bits_at_test() { + assert_eq!(0b0, even_bits_at(0)); + assert_eq!(0b1, even_bits_at(1)); + assert_eq!(0b100, even_bits_at(2)); + assert_eq!(0b101, even_bits_at(3)); + } + + #[test] + fn decompose_even_odd_test() { + use pasta_curves::pallas; + let odds = 0xAAAA; + let evens = 0x5555; + let (e, o) = decompose(pallas::Base::from_u128(odds)); + assert_eq!(e.get_lower_128(), 0); + assert_eq!(o.get_lower_128(), odds >> 1); + let (e, o) = decompose(pallas::Base::from_u128(evens)); + assert_eq!(e.get_lower_128(), evens); + assert_eq!(o.get_lower_128(), 0); + } +}