nf.rs

use anyhow::Result;
use halo2_proofs::{circuit::{SimpleFloorPlanner, Value}, dev::MockProver, pasta::Fp, plonk::{Advice, Circuit, Column, Instance}};
use orchard::{assign_free_advice, keys::{FullViewingKey, SpendingKey}, note::{NoteCommitment, RandomSeed, Rho}, value::NoteValue, AddChip, AddConfig, Address, EccChip, EccConfig, Note, OrchardFixedBases};
use halo2_gadgets::{ecc::Point, utilities::lookup_range_check::LookupRangeCheckConfig};
use halo2_gadgets::poseidon::{Pow5Chip as PoseidonChip, Pow5Config as PoseidonConfig};
use group::Curve;

#[derive(Default)]
struct NFCircuit {
    nk: Value<Fp>,
    cm: Value<NoteCommitment>,
    rho: Value<Fp>,
    psi: Value<Fp>,
}

#[derive(Clone)]
struct NFCircuitConfig {
    primary: Column<Instance>,
    advices: [Column<Advice>; 10],
    add_config: AddConfig,
    ecc_config: EccConfig<OrchardFixedBases>,
    poseidon_config: PoseidonConfig<Fp, 3, 2>,
}

impl Circuit<Fp> for NFCircuit {
    type Config = NFCircuitConfig;

    type FloorPlanner = SimpleFloorPlanner;

    fn without_witnesses(&self) -> Self {
        Self::default()
    }

    fn configure(meta: &mut halo2_proofs::plonk::ConstraintSystem<Fp>) -> Self::Config {
        let primary = meta.instance_column();
        meta.enable_equality(primary);
        let advices = [
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
            meta.advice_column(),
        ];
        for advice in advices.iter() {
            meta.enable_equality(*advice);
        }
        let lagrange_coeffs = [
            meta.fixed_column(),
            meta.fixed_column(),
            meta.fixed_column(),
            meta.fixed_column(),
            meta.fixed_column(),
            meta.fixed_column(),
            meta.fixed_column(),
            meta.fixed_column(),
        ];
        meta.enable_constant(lagrange_coeffs[0]);
        let table_idx = meta.lookup_table_column();
        let range_check = LookupRangeCheckConfig::configure(meta, advices[9], table_idx);
        let rc_a = lagrange_coeffs[2..5].try_into().unwrap();
        let rc_b = lagrange_coeffs[5..8].try_into().unwrap();
        
        let add_config = orchard::AddChip::configure(meta, advices[7], advices[8], advices[6]);
        let ecc_config =
            EccChip::<OrchardFixedBases>::configure(meta, advices, lagrange_coeffs, range_check);
        let poseidon_config = PoseidonChip::configure::<halo2_gadgets::poseidon::primitives::P128Pow5T3>(
            meta,
            // We place the state columns after the partial_sbox column so that the
            // pad-and-add region can be laid out more efficiently.
            advices[6..9].try_into().unwrap(),
            advices[5],
            rc_a,
            rc_b,
        );
    
        NFCircuitConfig {
            primary,
            advices,
            add_config,
            ecc_config,
            poseidon_config,
        }
    }

    fn synthesize(&self, config: Self::Config, mut layouter: impl halo2_proofs::circuit::Layouter<Fp>) -> std::prelude::v1::Result<(), halo2_proofs::plonk::Error> {
        let ecc_chip = EccChip::construct(config.ecc_config.clone());
        let rho = assign_free_advice(
            layouter.namespace(|| "witness rho_old"),
            config.advices[0],
            self.rho,
        )?;
        let psi = assign_free_advice(
            layouter.namespace(|| "witness psi_old"),
            config.advices[0],
            self.psi,
        )?;
        let cm = Point::new(
            ecc_chip.clone(),
            layouter.namespace(|| "cm_old"),
            self.cm.as_ref().map(|cm| cm.inner().to_affine()),
        )?;
        let nk = assign_free_advice(
            layouter.namespace(|| "witness nk"),
            config.advices[0],
            self.nk,
        )?;

        let nf = orchard::circuit::gadget::derive_nullifier(
            layouter.namespace(|| "nf_old = DeriveNullifier_nk(rho_old, psi_old, cm_old)"),
            PoseidonChip::construct(config.poseidon_config.clone()),
            AddChip::construct(config.add_config.clone()),
            ecc_chip,
            rho,
            &psi,
            &cm,
            nk,
        )?;

        // Constrain nf_old to equal public input
        layouter.constrain_instance(nf.inner().cell(), config.primary, 0)?;
        Ok(())
    }
}

fn main() -> Result<()> {
    let address = Address::from_raw_address_bytes(&[11; 43]).unwrap();
    let nv = NoteValue::from_raw(1_000_000);
    let rho = Rho::from_bytes(&[42; 32]).unwrap();
    let rseed = RandomSeed::from_bytes([11; 32], &rho).unwrap();
    let note = Note::from_parts(address, nv, rho,
    rseed).unwrap();

    let sk = SpendingKey::from_bytes([42; 32]).unwrap();
    let fvk = FullViewingKey::from(&sk);
    let nf = note.nullifier(&fvk);
    println!("{}", hex::encode(nf.to_bytes()));
    let nf = nf.0;
    
    let nk = fvk.nk().inner();
    let cm = note.commitment();
    let psi = rseed.psi(&rho);
    let rho = rho.into_inner();

    let nf_circuit = NFCircuit {
        nk: Value::known(nk),
        cm: Value::known(cm),
        rho: Value::known(rho),
        psi: Value::known(psi),
    };
    
    let prover = MockProver::run(8, &nf_circuit, vec![vec![nf]])?;
    prover.verify().unwrap();

    Ok(())
}