darkfi/src/zk/gadget/less_than.rs

/* This file is part of DarkFi (https://dark.fi)
 *
 * Copyright (C) 2020-2023 Dyne.org foundation
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use halo2_proofs::{
    arithmetic::Field,
    circuit::{AssignedCell, Chip, Layouter, Region, Value},
    pasta::pallas,
    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector, TableColumn},
    poly::Rotation,
};

use super::native_range_check::{NativeRangeCheckChip, NativeRangeCheckConfig};

#[derive(Clone, Debug)]
pub struct LessThanConfig<
    const WINDOW_SIZE: usize,
    const NUM_OF_BITS: usize,
    const NUM_OF_WINDOWS: usize,
> {
    pub s_lt: Selector,
    pub a: Column<Advice>,
    pub b: Column<Advice>,
    pub a_offset: Column<Advice>,
    pub range_a_config: NativeRangeCheckConfig<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>,
    pub range_a_offset_config: NativeRangeCheckConfig<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>,
    pub k_values_table: TableColumn,
}

#[derive(Clone, Debug)]
pub struct LessThanChip<
    const WINDOW_SIZE: usize,
    const NUM_OF_BITS: usize,
    const NUM_OF_WINDOWS: usize,
> {
    config: LessThanConfig<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>,
}

impl<const WINDOW_SIZE: usize, const NUM_OF_BITS: usize, const NUM_OF_WINDOWS: usize>
    Chip<pallas::Base> for LessThanChip<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>
{
    type Config = LessThanConfig<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>;
    type Loaded = ();

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

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

impl<const WINDOW_SIZE: usize, const NUM_OF_BITS: usize, const NUM_OF_WINDOWS: usize>
    LessThanChip<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>
{
    pub fn construct(config: LessThanConfig<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>) -> Self {
        Self { config }
    }

    pub fn configure(
        meta: &mut ConstraintSystem<pallas::Base>,
        a: Column<Advice>,
        b: Column<Advice>,
        a_offset: Column<Advice>,
        z1: Column<Advice>,
        z2: Column<Advice>,
        k_values_table: TableColumn,
    ) -> LessThanConfig<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS> {
        let s_lt = meta.selector();

        meta.enable_equality(a);
        meta.enable_equality(b);
        meta.enable_equality(a_offset);
        meta.enable_equality(z1);
        meta.enable_equality(z2);

        // configure range check for `a` and `offset`
        let range_a_config =
            NativeRangeCheckChip::<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>::configure(
                meta,
                z1,
                k_values_table,
            );

        let range_a_offset_config =
            NativeRangeCheckChip::<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>::configure(
                meta,
                z2,
                k_values_table,
            );

        let config = LessThanConfig {
            s_lt,
            a,
            b,
            a_offset,
            range_a_config,
            range_a_offset_config,
            k_values_table,
        };

        meta.create_gate("a_offset - 2^m + b - a", |meta| {
            let s_lt = meta.query_selector(config.s_lt);
            let a = meta.query_advice(config.a, Rotation::cur());
            let b = meta.query_advice(config.b, Rotation::cur());
            let a_offset = meta.query_advice(config.a_offset, Rotation::cur());
            let two_pow_m =
                Expression::Constant(pallas::Base::from(2).pow([NUM_OF_BITS as u64, 0, 0, 0]));
            // a_offset - 2^m + b - a = 0
            vec![s_lt * (a_offset - two_pow_m + b - a)]
        });

        config
    }

    pub fn witness_less_than(
        &self,
        mut layouter: impl Layouter<pallas::Base>,
        a: Value<pallas::Base>,
        b: Value<pallas::Base>,
        offset: usize,
        strict: bool,
    ) -> Result<(), Error> {
        let (a, _, a_offset) = layouter.assign_region(
            || "a less than b",
            |mut region: Region<'_, pallas::Base>| {
                let a = region.assign_advice(|| "a", self.config.a, offset, || a)?;
                let b = region.assign_advice(|| "b", self.config.b, offset, || b)?;
                let a_offset = self.less_than(region, a.clone(), b.clone(), offset)?;
                Ok((a, b, a_offset))
            },
        )?;

        self.less_than_range_check(layouter, a, a_offset, strict)?;

        Ok(())
    }

    /*
    pub fn witness_less_than2(
        &self,
        mut layouter: impl Layouter<pallas::Base>,
        a: Value<pallas::Base>,
        b: Value<pallas::Base>,
        offset: usize,
        strict: bool,
    ) -> Result<AssignedCell<pallas::Base, pallas::Base>, Error> {
        let (a, _, a_offset) = layouter.assign_region(
            || "a less than b",
            |mut region: Region<'_, pallas::Base>| {
                let a = region.assign_advice(|| "a", self.config.a, offset, || a)?;
                let b = region.assign_advice(|| "b", self.config.b, offset, || b)?;
                let a_offset = self.less_than(region, a.clone(), b.clone(), offset)?;
                Ok((a, b, a_offset))
            },
        )?;

        self.less_than_range_check(layouter, a, a_offset.clone(), strict)?;

        Ok(a_offset)
    }
    */

    pub fn copy_less_than(
        &self,
        mut layouter: impl Layouter<pallas::Base>,
        a: AssignedCell<pallas::Base, pallas::Base>,
        b: AssignedCell<pallas::Base, pallas::Base>,
        offset: usize,
        strict: bool,
    ) -> Result<(), Error> {
        let (a, _, a_offset) = layouter.assign_region(
            || "a less than b",
            |mut region: Region<'_, pallas::Base>| {
                let a = a.copy_advice(|| "a", &mut region, self.config.a, offset)?;
                let b = b.copy_advice(|| "b", &mut region, self.config.b, offset)?;
                let a_offset = self.less_than(region, a.clone(), b.clone(), offset)?;
                Ok((a, b, a_offset))
            },
        )?;

        self.less_than_range_check(layouter, a, a_offset, strict)?;

        Ok(())
    }

    pub fn less_than_range_check(
        &self,
        mut layouter: impl Layouter<pallas::Base>,
        a: AssignedCell<pallas::Base, pallas::Base>,
        a_offset: AssignedCell<pallas::Base, pallas::Base>,
        strict: bool,
    ) -> Result<(), Error> {
        let range_a_chip =
            NativeRangeCheckChip::<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>::construct(
                self.config.range_a_config.clone(),
            );
        let range_a_offset_chip =
            NativeRangeCheckChip::<WINDOW_SIZE, NUM_OF_BITS, NUM_OF_WINDOWS>::construct(
                self.config.range_a_offset_config.clone(),
            );

        range_a_chip.copy_range_check(layouter.namespace(|| "a copy_range_check"), a, strict)?;

        range_a_offset_chip.copy_range_check(
            layouter.namespace(|| "a_offset copy_range_check"),
            a_offset,
            strict,
        )?;

        Ok(())
    }

    pub fn less_than(
        &self,
        mut region: Region<'_, pallas::Base>,
        a: AssignedCell<pallas::Base, pallas::Base>,
        b: AssignedCell<pallas::Base, pallas::Base>,
        offset: usize,
    ) -> Result<AssignedCell<pallas::Base, pallas::Base>, Error> {
        // enable `less_than` selector
        self.config.s_lt.enable(&mut region, offset)?;

        // assign `a + offset`
        let two_pow_m = pallas::Base::from(2).pow([NUM_OF_BITS as u64, 0, 0, 0]);
        let a_offset = a.value().zip(b.value()).map(|(a, b)| *a + (two_pow_m - b));
        let a_offset =
            region.assign_advice(|| "a_offset", self.config.a_offset, offset, || a_offset)?;

        Ok(a_offset)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use darkfi_sdk::crypto::pasta_prelude::PrimeField;
    use halo2_proofs::{
        circuit::{floor_planner, Value},
        dev::{CircuitLayout, MockProver},
        pasta::pallas,
        plonk::Circuit,
    };

    macro_rules! test_circuit {
        ($window_size:expr, $num_bits:expr, $num_windows:expr) => {
            #[derive(Default)]
            struct LessThanCircuit {
                a: Value<pallas::Base>,
                b: Value<pallas::Base>,
            }

            impl Circuit<pallas::Base> for LessThanCircuit {
                type Config =
                    (LessThanConfig<$window_size, $num_bits, $num_windows>, Column<Advice>);
                type FloorPlanner = floor_planner::V1;

                fn without_witnesses(&self) -> Self {
                    Self { a: Value::unknown(), b: Value::unknown() }
                }

                fn configure(meta: &mut ConstraintSystem<pallas::Base>) -> Self::Config {
                    let w = meta.advice_column();
                    meta.enable_equality(w);

                    let a = meta.advice_column();
                    let b = meta.advice_column();
                    let a_offset = meta.advice_column();
                    let z1 = meta.advice_column();
                    let z2 = meta.advice_column();

                    let k_values_table = meta.lookup_table_column();

                    let constants = meta.fixed_column();
                    meta.enable_constant(constants);

                    (
                        LessThanChip::<$window_size, $num_bits, $num_windows>::configure(
                            meta,
                            a,
                            b,
                            a_offset,
                            z1,
                            z2,
                            k_values_table,
                        ),
                        w,
                    )
                }

                fn synthesize(
                    &self,
                    config: Self::Config,
                    mut layouter: impl Layouter<pallas::Base>,
                ) -> Result<(), Error> {
                    let less_than_chip =
                        LessThanChip::<$window_size, $num_bits, $num_windows>::construct(
                            config.0.clone(),
                        );

                    NativeRangeCheckChip::<$window_size, $num_bits, $num_windows>::load_k_table(
                        &mut layouter,
                        config.0.k_values_table,
                    )?;

                    less_than_chip.witness_less_than(
                        layouter.namespace(|| "a < b"),
                        self.a,
                        self.b,
                        0,
                        true,
                    )?;

                    Ok(())
                }
            }
        };
    }

    #[test]
    fn less_than_64() {
        test_circuit!(3, 64, 22);
        let k = 5;

        let valid_pairs = [
            (pallas::Base::from(13), pallas::Base::from(15)),
            (pallas::Base::ZERO, pallas::Base::from(u64::MAX)),
            (pallas::Base::ONE, pallas::Base::from(rand::random::<u64>())),
        ];

        let invalid_pairs = [
            (pallas::Base::from(14), pallas::Base::from(11)),
            (pallas::Base::from(u64::MAX), pallas::Base::ZERO),
            (pallas::Base::ZERO, pallas::Base::ZERO),
            (pallas::Base::ONE, pallas::Base::ONE),
            (pallas::Base::ONE, pallas::Base::ZERO),
            (pallas::Base::from(u64::MAX), pallas::Base::from(u64::MAX)),
            (pallas::Base::from(u64::MAX), pallas::Base::from(u64::MAX) + pallas::Base::ONE),
        ];

        use plotters::prelude::*;
        let circuit = LessThanCircuit {
            a: Value::known(pallas::Base::zero()),
            b: Value::known(pallas::Base::one()),
        };
        let root = BitMapBackend::new("target/lessthan_circuit_layout.png", (3840, 2160))
            .into_drawing_area();
        CircuitLayout::default().render(k, &circuit, &root).unwrap();

        for (a, b) in valid_pairs {
            println!("64 bit (valid) {:?} < {:?} check", a, b);
            let circuit = LessThanCircuit { a: Value::known(a), b: Value::known(b) };
            let prover = MockProver::run(k, &circuit, vec![]).unwrap();
            prover.assert_satisfied();
        }

        for (a, b) in invalid_pairs {
            println!("64 bit (invalid) {:?} < {:?} check", a, b);
            let circuit = LessThanCircuit { a: Value::known(a), b: Value::known(b) };
            let prover = MockProver::run(k, &circuit, vec![]).unwrap();
            assert!(prover.verify().is_err())
        }
    }

    #[test]
    fn less_than_253() {
        test_circuit!(3, 253, 85);
        let k = 7;

        const P_MINUS_1: pallas::Base = pallas::Base::from_raw([
            0x992d30ed00000000,
            0x224698fc094cf91b,
            0x0000000000000000,
            0x4000000000000000,
        ]);

        // 2^253 - 1. This is the maximum we can check.
        const MAX_253: pallas::Base = pallas::Base::from_raw([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0x1FFFFFFFFFFFFFFF,
        ]);

        let valid_pairs = [
            (pallas::Base::ZERO, pallas::Base::ONE),
            (pallas::Base::from(u64::MAX), pallas::Base::from(u64::MAX) + pallas::Base::ONE),
            (
                pallas::Base::from_u128(u128::MAX),
                pallas::Base::from_u128(u128::MAX) + pallas::Base::ONE,
            ),
            (MAX_253 - pallas::Base::from(2), MAX_253 - pallas::Base::ONE),
            (MAX_253 - pallas::Base::ONE, MAX_253),
            (MAX_253, MAX_253 + pallas::Base::ONE), // <-- Is this legit??
        ];

        let invalid_pairs = [
            (pallas::Base::ZERO, pallas::Base::ZERO),
            (pallas::Base::ONE, pallas::Base::ZERO),
            (P_MINUS_1 - pallas::Base::ONE, P_MINUS_1),
            (P_MINUS_1, P_MINUS_1),
            (P_MINUS_1, pallas::Base::ZERO),
            (MAX_253, MAX_253),
            (MAX_253, pallas::Base::ZERO),
            (MAX_253, pallas::Base::ONE),
            (MAX_253 + pallas::Base::ONE, pallas::Base::ZERO),
            (MAX_253 + pallas::Base::ONE, pallas::Base::ONE),
            (MAX_253 + pallas::Base::ONE, MAX_253 + pallas::Base::ONE),
            (MAX_253 + pallas::Base::ONE, MAX_253 + pallas::Base::from(2)),
        ];

        for (a, b) in valid_pairs {
            println!("253 bit (valid) {:?} < {:?} check", a, b);
            let circuit = LessThanCircuit { a: Value::known(a), b: Value::known(b) };
            let prover = MockProver::run(k, &circuit, vec![]).unwrap();
            prover.assert_satisfied();
        }

        for (a, b) in invalid_pairs {
            println!("253 bit (invalid) {:?} < {:?} check", a, b);
            let circuit = LessThanCircuit { a: Value::known(a), b: Value::known(b) };
            let prover = MockProver::run(k, &circuit, vec![]).unwrap();
            assert!(prover.verify().is_err())
        }
    }
}