darkfi/src/zk/gadget/range_check.rs

use std::marker::PhantomData;

use group::ff::PrimeFieldBits;
use halo2_proofs::{
    arithmetic::FieldExt,
    circuit::{Chip, Layouter, Value},
    plonk::{Advice, Column, ConstraintSystem, Error, Selector, TableColumn},
    poly::Rotation,
};

#[derive(Clone, Debug)]
pub struct RangeCheckConfig {
    pub z: Column<Advice>,
    pub s_rc: Selector,
    pub k_values_table: TableColumn,
}

pub struct RangeCheckChip<F: FieldExt + PrimeFieldBits, const WINDOW_SIZE: usize> {
    config: RangeCheckConfig,
    _marker: PhantomData<F>,
}

impl<F: FieldExt + PrimeFieldBits, const WINDOW_SIZE: usize> Chip<F>
    for RangeCheckChip<F, WINDOW_SIZE>
{
    type Config = RangeCheckConfig;
    type Loaded = ();

    fn config(&self) -> &Self::Config {
        &self.config
    }

    fn loaded(&self) -> &Self::Loaded {
        &()
    }
}

impl<F: FieldExt + PrimeFieldBits, const WINDOW_SIZE: usize> RangeCheckChip<F, WINDOW_SIZE> {
    pub fn configure(
        meta: &mut ConstraintSystem<F>,
        k_values_table: TableColumn,
    ) -> <Self as Chip<F>>::Config {
        let z = meta.advice_column();
        let s_rc = meta.complex_selector();

        let config = RangeCheckConfig { z, s_rc, k_values_table };

        meta.lookup(|meta| {
            let s_rc = meta.query_selector(config.s_rc);
            let z_curr = meta.query_advice(config.z, Rotation::cur());
            let z_next = meta.query_advice(config.z, Rotation::next());

            //    z_next = (z_curr - k_i) / 2^K
            // => k_i = z_curr - (z_next * 2^K)
            vec![(s_rc * (z_curr - z_next * F::from(1 << WINDOW_SIZE)), config.k_values_table)]
        });

        config
    }

    /// `k_values_table` should be reused across different chips
    /// which is why we don't limit it to a specific instance.
    pub fn load_k_table(
        layouter: &mut impl Layouter<F>,
        k_values_table: TableColumn,
    ) -> Result<(), Error> {
        layouter.assign_table(
            || format!("{} window table", WINDOW_SIZE),
            |mut table| {
                for index in 0..(1 << WINDOW_SIZE) {
                    table.assign_cell(
                        || "table",
                        k_values_table,
                        index,
                        || Value::known(F::from(index as u64)),
                    )?;
                }
                Ok(())
            },
        )
    }
}