Hekate ZK Engine

use hekate_core::errors;
use hekate_gadgets::IntArithmeticChiplet;
use hekate_math::{Block128, TowerField};
use hekate_program::FixedShape;
use hekate_program::chiplet::ChipletDef;
use hekate_program::circuit::{Circuit, CircuitProgram};

type F = Block128;

fn build_program(num_rows: usize) -> errors::Result<CircuitProgram<F>> {
    let chiplet = IntArithmeticChiplet::new(32, num_rows, num_rows - 1)?;
    let layout = chiplet.layout().clone();

    // The arithmetic chiplet's columns
    // and AIR are the program's trace.
    let mut cx = Circuit::<F>::new("FibonacciInt", num_rows)?;
    let arith = cx.mount_unlinked(ChipletDef::from_air(&chiplet)?);

    let cs = cx.cs();

    let s_add = cs.col(layout.s_add);
    let val_b = cs.col(layout.val_b);
    let val_res = cs.col(layout.val_res);
    let next_val_a = cs.next(layout.val_a);
    let next_val_b = cs.next(layout.val_b);

    // The two Fibonacci transitions,
    // on top of the chiplet's AIR.
    cs.constrain(s_add * (next_val_a + val_b));
    cs.constrain(s_add * (next_val_b + val_res));

    cs.assert_zero_when(cs.one() + s_add, val_res);

    // Boundaries pin fib(0) and fib(1)
    cx.boundary(arith.col(layout.val_a), 0, F::ZERO);
    cx.boundary(arith.col(layout.val_b), 0, F::ONE);

    cx.fix(arith.col(layout.s_add), FixedShape::LastRow);

    // Publish exposes the last value
    cx.publish(arith.col(layout.val_b), num_rows - 1);

    cx.compile()
}

use hekate_core::errors;
use hekate_core::trace::{ColumnTrace, TraceColumn};
use hekate_gadgets::{
    ArithmeticOpcode, IntArithmeticLayout, IntArithmeticOp, generate_arithmetic_trace,
};
use hekate_math::{Block32, HardwareField};

const PHY_VAL_A: usize = 0;
const PHY_VAL_B: usize = 1;

fn generate_trace(num_rows: usize) -> errors::Result<(ColumnTrace, u32)> {
    let layout = IntArithmeticLayout::compute(32);

    let mut a: u32 = 0;
    let mut b: u32 = 1;
    let mut prev_b: u32 = 0;

    // N - 1 ADD ops become the chiplet trace
    let mut ops: Vec<IntArithmeticOp> = Vec::with_capacity(num_rows - 1);

    for _ in 0..num_rows - 1 {
        let sum = a.wrapping_add(b);
        ops.push(IntArithmeticOp::U32 {
            op: ArithmeticOpcode::ADD,
            a,
            b,
        });

        prev_b = b;
        a = b;
        b = sum;
    }

    let final_b = b;
    let final_a = prev_b;

    let mut trace = generate_arithmetic_trace(&ops, &layout, num_rows)?;

    // Patch the padding row: the transition
    // from row N - 2 to N - 1 stays valid.
    if let TraceColumn::B32(col) = &mut trace.columns[PHY_VAL_A] {
        col[num_rows - 1] = Block32::from(final_a).to_hardware();
    }

    if let TraceColumn::B32(col) = &mut trace.columns[PHY_VAL_B] {
        col[num_rows - 1] = Block32::from(final_b).to_hardware();
    }

    Ok((trace, final_b))
}

use hekate_core::config::Config;
use hekate_core::proofs::InnerProof;
use hekate_program::{ProgramInstance, ProgramWitness};
use hekate_prover_sys::prove;
use rand::{TryRng, rngs::SysRng};
use std::error::Error;

fn prove_fibonacci(
    num_rows: usize,
) -> Result<(ProgramInstance<F>, InnerProof<F>), Box<dyn Error>> {
    let (trace, final_b) = generate_trace(num_rows)?;

    // The instance is public;
    // the witness stays on the device.
    let instance = ProgramInstance::new(num_rows, vec![F::from(final_b as u128)]);
    let witness = ProgramWitness::new(trace);
    let program = build_program(num_rows)?;

    // A fresh blinding seed makes the proof zero-knowledge
    let mut seed = [0; 32];
    SysRng.try_fill_bytes(&mut seed)?;

    let proof = prove(
        b"FibonacciIntChiplet",
        &program,
        &instance,
        &witness,
        &Config::prod(),
        seed,
        None,
    )?;

    Ok((instance, proof))
}

use hekate_core::config::Config;
use hekate_core::proofs::InnerProof;
use hekate_crypto::DefaultHasher;
use hekate_crypto::transcript::Transcript;
use hekate_program::ProgramInstance;
use hekate_program::circuit::CircuitProgram;
use hekate_verifier::HekateVerifier;

fn verify_fibonacci(
    pinned_id: &[u8; 32],
    program: &CircuitProgram<F>,
    instance: &ProgramInstance<F>,
    proof: &InnerProof<F>,
) -> bool {
    // Same transcript label and Config as the prover
    let mut transcript = Transcript::<DefaultHasher>::new(b"FibonacciIntChiplet");

    // A proof for another circuit fails
    // against the pinned program id.
    let verdict = HekateVerifier::<F, DefaultHasher>::verify(
        pinned_id,
        program,
        instance,
        proof,
        &mut transcript,
        &Config::prod(),
    );

    matches!(verdict, Ok(true))
}