miden-precompiles-prover 0.33.0

use std::{format, string::String, vec, vec::Vec};

use k256::{ProjectivePoint, elliptic_curve::sec1::ToSec1Point};
use miden_air::lookup::Challenges;
use miden_core::{
    Felt,
    deferred::{Digest, Node as VmNode, PrecompileWitness, TRUE_DIGEST as VM_TRUE_DIGEST},
    field::QuadFelt,
    proof::{HashFunction, StarkProof},
    serde::{Deserializable, Serializable},
};
use miden_precompiles::{
    CurveId, CurvePrecompile, Keccak256Precompile, UintDomain, UintPrecompile,
};
use miden_precompiles_verifier::{VerifyError, verify_deferred};
use rand::{Rng, RngExt, SeedableRng, rngs::StdRng};

use crate::{
    PrecompileProvingError,
    deferred::session::session_from_witnesses,
    hash::{
        chunk_node_sponge::SPONGE_COL_OFFSET,
        keccak::sponge::{COL_ACT as SPONGE_COL_ACT, SPONGE_PERIOD, trace::keccak_oracle},
    },
    math::{U256, from_hex, to_limbs32},
    prove_precompiles,
    relations::{MAX_MESSAGE_WIDTH, NUM_BUS_IDS},
    session::{Session, SessionTraces},
    tests::{
        SessionTracesTestExt, batch_witness::WitnessFixture, bus_balance::session_stack_residual,
        verify_deferred as verify_session,
    },
    transcript::poseidon2::P2Digest,
};

/// A raw Keccak-only portable fixture and the prover-typed view of its root.
#[derive(Debug)]
struct SyntheticKeccakWitness {
    state: WitnessFixture,
    input_digest: Digest,
    expected_digest: Digest,
    assertion_digest: Digest,
    vm_root: Digest,
    root: P2Digest,
}

/// Builds the Keccak-only committed graph for `input`:
/// `AND(TRUE_DIGEST, Keccak256Assert(chunks(input), chunks(keccak256(input))))`.
fn synthetic_keccak_state(input: &[u8]) -> SyntheticKeccakWitness {
    let mut state = WitnessFixture::new();

    let input_digest = state
        .register(VmNode::chunks_from_bytes(input))
        .expect("VM should register input chunks");
    let expected_digest = state
        .register(VmNode::chunks(keccak_digest_chunks(input)).expect("digest chunks are non-empty"))
        .expect("VM should register expected digest chunks");
    let assertion_digest = state
        .register(Keccak256Precompile::assert_node(
            len_bytes(input),
            input_digest,
            expected_digest,
        ))
        .expect("Keccak assertion should register structurally");

    let vm_root = state
        .log_statement(assertion_digest)
        .expect("true Keccak assertion should log into the deferred root");
    debug_assert_eq!(vm_root, VmNode::and(VM_TRUE_DIGEST, assertion_digest).digest());
    debug_assert_eq!(state.root(), vm_root);

    SyntheticKeccakWitness {
        state,
        input_digest,
        expected_digest,
        assertion_digest,
        vm_root,
        root: P2Digest::from(vm_root),
    }
}

fn keccak_session_traces(input: &[u8]) -> SessionTraces {
    let mut session = Session::new();
    let (_, claim) = session.keccak(input);
    let root = session.assert_and_fold([claim]);
    session.finish(root)
}

fn keccak_digest_chunks(input: &[u8]) -> Vec<[Felt; 8]> {
    vec![keccak_oracle(input).to_u32s().map(Felt::from_u32)]
}

fn len_bytes(input: &[u8]) -> u32 {
    u32::try_from(input.len()).expect("Keccak MVP inputs fit in a VM u32 length tag")
}

fn register_keccak_assertion(state: &mut WitnessFixture, input: &[u8]) -> Digest {
    let input_digest = state
        .register(VmNode::chunks_from_bytes(input))
        .expect("register Keccak input chunks");
    let expected_digest = state
        .register(VmNode::chunks(keccak_digest_chunks(input)).expect("digest chunks are non-empty"))
        .expect("register Keccak expected digest chunks");
    state
        .register(Keccak256Precompile::assert_node(
            u32::try_from(input.len()).expect("test input length fits u32"),
            input_digest,
            expected_digest,
        ))
        .expect("matching Keccak assertion registers")
}

fn register_uint_value(state: &mut WitnessFixture, domain: UintDomain, value: U256) -> Digest {
    state
        .register(UintPrecompile::value_node(domain, to_limbs32(value)))
        .expect("register uint value node")
}

fn register_uint_op(state: &mut WitnessFixture, op_id: u64, lhs: Digest, rhs: Digest) -> Digest {
    state
        .register(VmNode::join(UintPrecompile::op_tag(op_id), lhs, rhs).expect("uint op tag"))
        .expect("register uint op node")
}

fn register_curve_point(state: &mut WitnessFixture, curve: CurveId, x: U256, y: U256) -> Digest {
    let x_digest = register_uint_value(state, curve.base_domain(), x);
    let y_digest = register_uint_value(state, curve.base_domain(), y);
    state
        .register(CurvePrecompile::affine_node_from_digests(curve, x_digest, y_digest))
        .expect("register curve point value node")
}

fn register_curve_identity(state: &mut WitnessFixture, curve: CurveId) -> Digest {
    state
        .register(CurvePrecompile::identity_node(curve))
        .expect("register curve identity node")
}

fn register_curve_generator(state: &mut WitnessFixture, curve: CurveId) -> Digest {
    state
        .register(CurvePrecompile::generator_node(curve))
        .expect("register curve generator node")
}

fn register_curve_op(state: &mut WitnessFixture, op_id: u64, lhs: Digest, rhs: Digest) -> Digest {
    state
        .register(VmNode::join(CurvePrecompile::op_tag(op_id), lhs, rhs).expect("curve op tag"))
        .expect("register curve op node")
}

fn register_curve_msm(state: &mut WitnessFixture, pairs: Vec<(Digest, Digest)>) -> Digest {
    state
        .register(
            VmNode::try_pair_list(CurvePrecompile::msm_tag(), pairs)
                .expect("curve msm pair list is non-empty"),
        )
        .expect("register curve msm node")
}

fn be_to_u256(bytes: impl AsRef<[u8]>) -> U256 {
    let hex: String = bytes.as_ref().iter().map(|b| format!("{b:02x}")).collect();
    from_hex(&hex)
}

fn k256_coords(point: &ProjectivePoint) -> (U256, U256) {
    let enc = point.to_affine().to_sec1_point(false);
    (
        be_to_u256(enc.x().expect("finite point")),
        be_to_u256(enc.y().expect("finite point")),
    )
}

fn k1_points() -> [(U256, U256); 3] {
    let g = ProjectivePoint::GENERATOR;
    let g2 = g + g;
    let g3 = g + g + g;
    [k256_coords(&g), k256_coords(&g2), k256_coords(&g3)]
}

fn all_node_vm_state() -> WitnessFixture {
    let mut state = WitnessFixture::new();

    let curve = CurveId::Secp256k1;
    let domain = UintDomain::K1Base;
    let scalar_domain = curve.scalar_domain();
    let [(gx, gy), (g2x, g2y), (g3x, g3y)] = k1_points();

    let mut claims = Vec::new();

    claims.push(register_keccak_assertion(&mut state, b"all-node synthetic dag"));

    let u11 = register_uint_value(&mut state, domain, U256::from(11u8));
    let u7 = register_uint_value(&mut state, domain, U256::from(7u8));

    let add = register_uint_op(&mut state, UintPrecompile::ADD_OP_ID, u11, u7);
    let add_expected = register_uint_value(&mut state, domain, U256::from(18u8));
    claims.push(register_uint_op(&mut state, UintPrecompile::EQ_OP_ID, add, add_expected));

    let sub = register_uint_op(&mut state, UintPrecompile::SUB_OP_ID, u11, u7);
    let sub_expected = register_uint_value(&mut state, domain, U256::from(4u8));
    claims.push(register_uint_op(&mut state, UintPrecompile::EQ_OP_ID, sub, sub_expected));

    let mul = register_uint_op(&mut state, UintPrecompile::MUL_OP_ID, u11, u7);
    let mul_expected = register_uint_value(&mut state, domain, U256::from(77u8));
    claims.push(register_uint_op(&mut state, UintPrecompile::EQ_OP_ID, mul, mul_expected));

    let g_digest = register_curve_point(&mut state, curve, gx, gy);
    let g2_digest = register_curve_point(&mut state, curve, g2x, g2y);
    let g3_digest = register_curve_point(&mut state, curve, g3x, g3y);
    let inf_digest = register_curve_identity(&mut state, curve);

    claims.push(register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, inf_digest, inf_digest));

    let add_digest = register_curve_op(&mut state, CurvePrecompile::ADD_OP_ID, g_digest, g2_digest);
    claims.push(register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, add_digest, g3_digest));

    let sub_digest = register_curve_op(&mut state, CurvePrecompile::SUB_OP_ID, g3_digest, g_digest);
    claims.push(register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, sub_digest, g2_digest));

    let one_digest = register_uint_value(&mut state, scalar_domain, from_hex("1"));
    let msm_digest =
        register_curve_msm(&mut state, vec![(g_digest, one_digest), (g2_digest, one_digest)]);
    claims.push(register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, msm_digest, g3_digest));

    for claim in claims {
        state.log_statement(claim).expect("truthy synthetic claim logs");
    }

    state
}

fn translated_traces_check(state: &WitnessFixture) {
    let traces = session_from_witnesses(vec![state.witness()]).unwrap().finish();
    assert_eq!(traces.public_root(), P2Digest::from(state.root()));
    traces.check();
}

fn prove_fixture(
    state: &WitnessFixture,
    hash_fn: HashFunction,
) -> Result<StarkProof, PrecompileProvingError> {
    Ok(prove_precompiles(vec![state.witness()], hash_fn)?.proof)
}

#[test]
fn shared_truthy_dag_from_wire_proves_and_verifies() {
    // Includes Keccak, uint and EC equality leaves, plus nested arithmetic and MSM claims.
    let mut state = all_node_vm_state();
    for _ in 0..8 {
        state.log_statement(state.root()).unwrap();
    }
    let bytes = state.witness().to_bytes();
    let witness = PrecompileWitness::read_from_bytes(&bytes).unwrap();
    let traces = session_from_witnesses(vec![witness]).unwrap().finish();
    assert_eq!(traces.public_root(), P2Digest::from(state.root()));
    traces.check();
    let verified = verify_session(&traces.prove()).expect("shared truthy DAG proof must verify");
    assert_eq!(verified, state.root());
}

#[test]
fn synthetic_keccak_deferred_state_reconstructs_root() {
    let input: Vec<u8> = (0u8..33).collect();
    let synthetic = synthetic_keccak_state(&input);

    assert_eq!(synthetic.state.root(), synthetic.vm_root);
    assert_eq!(
        synthetic.vm_root,
        VmNode::and(VM_TRUE_DIGEST, synthetic.assertion_digest).digest(),
    );
    assert_eq!(synthetic.root, P2Digest::from(synthetic.vm_root));
    assert!(synthetic.state.get_node(&synthetic.input_digest).is_some());
    assert!(synthetic.state.get_node(&synthetic.expected_digest).is_some());
    assert!(synthetic.state.get_node(&synthetic.assertion_digest).is_some());
}

#[test]
fn session_public_root_matches_synthetic_deferred_state_for_keccak_inputs() {
    let cases: [(&str, Vec<u8>); 9] = [
        ("empty", Vec::new()),
        ("short", b"abc".to_vec()),
        ("one_chunk_minus_one_limb", vec![0xa5; 31]),
        ("one_chunk", vec![0xa5; 32]),
        ("two_chunks", vec![0xa5; 33]),
        ("keccak_rate_boundary", vec![0xa5; 136]),
        ("post_keccak_rate_boundary", vec![0xa5; 137]),
        ("trailing_zero", b"abc\0".to_vec()),
        ("explicit_padding_zeroes", vec![0, 0, 0, 0, 0]),
    ];

    for (name, input) in cases {
        let synthetic = synthetic_keccak_state(&input);
        let traces = keccak_session_traces(&input);
        assert_eq!(traces.public_root(), synthetic.root, "case {name}");
    }
}

#[test]
fn session_public_root_matches_synthetic_deferred_state_for_all_supported_node_types() {
    translated_traces_check(&all_node_vm_state());
}

#[test]
fn deferred_session_translates_curve_claims_for_all_fixed_curves() {
    let mut state = WitnessFixture::new();

    for curve in CurveId::ALL {
        let identity = register_curve_identity(&mut state, curve);
        let generator = register_curve_generator(&mut state, curve);

        let identity_eq =
            register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, identity, identity);
        state.log_statement(identity_eq).expect("identity equality logs");

        let generator_eq =
            register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, generator, generator);
        state.log_statement(generator_eq).expect("generator equality logs");

        let sum = register_curve_op(&mut state, CurvePrecompile::ADD_OP_ID, generator, identity);
        let sum_eq = register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, sum, generator);
        state.log_statement(sum_eq).expect("generator plus identity logs");
    }

    translated_traces_check(&state);
}

#[test]
fn deferred_state_accepts_msm_with_all_zero_scalars() {
    // 0·P = 𝒪, checked against an independently committed identity point.
    let mut state = WitnessFixture::new();

    let curve = CurveId::Secp256k1;
    let point = register_curve_generator(&mut state, curve);
    let scalar = register_uint_value(&mut state, curve.scalar_domain(), U256::ZERO);
    let msm = register_curve_msm(&mut state, vec![(point, scalar)]);
    let identity = register_curve_identity(&mut state, curve);
    let msm_eq = register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, msm, identity);
    state.log_statement(msm_eq).expect("all-zero MSM equality logs");

    translated_traces_check(&state);
}

#[test]
fn deferred_state_accepts_msm_with_repeated_base() {
    // a·P + b·P = (a + b)·P, checked against an independently committed `5·G`.
    let mut state = WitnessFixture::new();

    let curve = CurveId::Secp256k1;
    let point = register_curve_generator(&mut state, curve);
    let two = register_uint_value(&mut state, curve.scalar_domain(), U256::from(2u64));
    let three = register_uint_value(&mut state, curve.scalar_domain(), U256::from(3u64));
    let msm = register_curve_msm(&mut state, vec![(point, two), (point, three)]);
    let five_g = ProjectivePoint::GENERATOR * k256::Scalar::from(5u64);
    let (fx, fy) = k256_coords(&five_g);
    let five_g_node = register_curve_point(&mut state, curve, fx, fy);
    let msm_eq = register_curve_op(&mut state, CurvePrecompile::EQ_OP_ID, msm, five_g_node);
    state.log_statement(msm_eq).expect("repeated-base MSM equality logs");

    translated_traces_check(&state);
}

#[test]
fn trailing_zero_input_changes_root() {
    let abc = synthetic_keccak_state(b"abc");
    let abc_zero = synthetic_keccak_state(b"abc\0");

    // Generic chunk nodes are lengthless: these inputs pack to the same single zero-padded chunk.
    // The Keccak assertion tag's `len_bytes` and digest child are what distinguish them.
    assert_eq!(abc.input_digest, abc_zero.input_digest);
    assert_ne!(abc.expected_digest, abc_zero.expected_digest);
    assert_ne!(abc.assertion_digest, abc_zero.assertion_digest);
    assert_ne!(abc.root, abc_zero.root);

    let abc_traces = keccak_session_traces(b"abc");
    let abc_zero_traces = keccak_session_traces(b"abc\0");
    assert_eq!(abc_traces.public_root(), abc.root);
    assert_eq!(abc_zero_traces.public_root(), abc_zero.root);
    assert_ne!(abc_traces.public_root(), abc_zero_traces.public_root());
}

#[test]
fn keccak_deferred_state_proof_verifies_and_rejects_trailing_bytes() {
    let input = b"abc";
    let synthetic = synthetic_keccak_state(input);
    let traces = session_from_witnesses(vec![synthetic.state.witness()]).unwrap().finish();
    assert_eq!(traces.public_root(), synthetic.root);

    let proof = traces.prove();
    assert_eq!(P2Digest::from(proof.1), synthetic.root);
    verify_session(&proof).expect("Keccak deferred-state proof should verify");

    // The proof encoding is exact: an otherwise-valid proof with a trailing byte is rejected.
    let stark = prove_fixture(&synthetic.state, HashFunction::Blake3_256)
        .expect("Keccak deferred state should prove");
    let mut proof_bytes = stark.bytes().to_vec();
    proof_bytes.push(0);
    let trailing = StarkProof::new(proof_bytes, stark.hash_fn());
    let err = verify_deferred(&trailing, synthetic.vm_root)
        .expect_err("trailing proof bytes must be rejected");
    assert!(matches!(
        err,
        VerifyError::Deserialization(wincode::error::ReadError::TrailingBytes)
    ));
}

#[test]
fn prove_deferred_state_proves_non_empty_root() {
    let synthetic = synthetic_keccak_state(b"abc");

    let proof = prove_fixture(&synthetic.state, HashFunction::Blake3_256)
        .expect("Keccak deferred state should prove");

    verify_deferred(&proof, synthetic.vm_root).expect("Keccak deferred-state proof should verify");
    assert!(
        verify_deferred(&proof, VM_TRUE_DIGEST).is_err(),
        "the proof must be bound to the state's exact root",
    );
}

/// Each `HashFunction` selects a distinct preprocessed-bundle cache slot
/// (`miden_precompiles_air::preprocessed`, keyed by LMCS type). Proving and
/// verifying twice per hash function exercises both the cold path (first
/// call in the process, builds and caches the bundle) and the warm path
/// (later calls, reused from cache) for every slot, guarding against a
/// mismatched or stale cached bundle being reused across hash functions.
#[test]
#[ignore = "full prove/verify round-trip; run explicitly"]
fn prove_deferred_state_round_trips_for_every_hash_function() {
    let synthetic = synthetic_keccak_state(b"abc");
    let hash_fns = [
        HashFunction::Blake3_256,
        HashFunction::Rpo256,
        HashFunction::Rpx256,
        HashFunction::Poseidon2,
        HashFunction::Keccak,
    ];

    for hash_fn in hash_fns {
        for pass in 0..2 {
            let proof = prove_fixture(&synthetic.state, hash_fn)
                .unwrap_or_else(|e| panic!("{hash_fn:?} pass {pass} should prove: {e}"));
            verify_deferred(&proof, synthetic.vm_root)
                .unwrap_or_else(|e| panic!("{hash_fn:?} pass {pass} should verify: {e}"));
        }
    }
}

/// Reconstruct the full ten-chiplet LogUp balance, including verifier-side fixed-boundary
/// consumes. This checks the generated traces against each AIR's lookup evaluator;
/// `eval_external` is tested separately in `session::prove`.
fn assert_session_balanced(traces: &SessionTraces, rng: &mut impl Rng) {
    let challenges = Challenges::new(
        QuadFelt::new([Felt::new(rng.random()).unwrap(), Felt::new(rng.random()).unwrap()]),
        QuadFelt::new([Felt::new(rng.random()).unwrap(), Felt::new(rng.random()).unwrap()]),
        MAX_MESSAGE_WIDTH,
        NUM_BUS_IDS,
    );
    let mains = traces.mains();
    let residual = session_stack_residual(&mains, &[], &challenges);
    assert!(
        residual.is_empty(),
        "session stack imbalance: {} unmatched denom(s); e.g. net {:?} on {}",
        residual.len(),
        residual.first().map(|(m, _)| *m),
        residual.first().map(|(_, s)| s.as_str()).unwrap_or(""),
    );
}

/// Exercises the merged sponge band on multi-block (`> 136`-byte) messages.
/// The default full-proof fixtures use the single-block input `b"abc"`; this
/// test covers cross-block state, invocation seams, overshoot lanes, padding,
/// and final squeezing through both constraint and bus-balance checks.
#[test]
fn merged_chunk_node_sponge_multi_block_checks_and_balances() {
    let mut rng = StdRng::seed_from_u64(0xc0de_5b09);
    // 137: first byte past the rate boundary (2 blocks, pad in block 2).
    // 271: rate boundary − 1 across two blocks. 300, 407: overshoot variety.
    for len in [137usize, 271, 300, 407] {
        let input: Vec<u8> = (0..len).map(|i| i as u8).collect();
        let traces = keccak_session_traces(&input);
        // Inspect the production merged band rather than inferring activity from the input. This
        // fails if trace construction silently truncates the sponge invocation.
        let merged = traces.mains()[0];
        let active_sponge_rows = merged
            .values
            .chunks_exact(merged.width)
            .filter(|row| row[SPONGE_COL_OFFSET + SPONGE_COL_ACT] == Felt::ONE)
            .count();
        assert!(
            active_sponge_rows > SPONGE_PERIOD,
            "case len={len} must activate more than one sponge block, got {active_sponge_rows} rows"
        );
        traces.check();
        assert_session_balanced(&traces, &mut rng);
    }
}

/// Explicit full prove+verify of a multi-block Keccak session — the
/// end-to-end counterpart to the fast check/balance guard above, closing
/// the merged-AIR multi-block gap through the real STARK path.
#[test]
#[ignore = "full prove/verify round-trip; run explicitly"]
fn prove_deferred_state_round_trips_for_multi_block_keccak() {
    let synthetic = synthetic_keccak_state(&(0u8..200).collect::<Vec<u8>>());
    let proof = prove_fixture(&synthetic.state, HashFunction::Blake3_256)
        .expect("multi-block keccak session should prove");
    verify_deferred(&proof, synthetic.vm_root).expect("multi-block keccak session should verify");
}