miden-air 0.32.1

//! STARK configuration factories for different hash functions.
//!
//! Each factory creates a [`StarkConfig`](miden_crypto::stark::StarkConfig) bundling the
//! PCS parameters, LMCS commitment scheme, and Fiat-Shamir challenger for proving and verification.

use alloc::vec;

use miden_core::{Felt, Word, crypto::hash::Poseidon2, field::QuadFelt};
use miden_crypto::{
    field::Field,
    hash::{
        blake::Blake3Hasher,
        keccak::{Keccak256Hash, KeccakF, VECTOR_LEN},
        poseidon2::Poseidon2Permutation256,
        rpo::RpoPermutation256,
        rpx::RpxPermutation256,
    },
    merkle::{MerklePath, MerkleTree, NodeIndex},
    stark::{
        GenericStarkConfig,
        challenger::{CanObserve, DuplexChallenger, HashChallenger, SerializingChallenger64},
        dft::Radix2DitParallel,
        hasher::{ChainingHasher, SerializingStatefulSponge, StatefulSponge},
        lmcs::config::LmcsConfig,
        pcs::PcsParams,
        symmetric::{
            CompressionFunctionFromHasher, CryptographicPermutation, PaddingFreeSponge,
            TruncatedPermutation,
        },
    },
};

use crate::{PROOF_ORDER_COUNT, PROOF_ORDER_REGISTRY_DEPTH, ProofOrder};

// SHARED TYPES
// ================================================================================================

/// Miden VM STARK configuration with pre-filled common type parameters.
///
/// All Miden configurations use `Felt` as the base field, `QuadFelt` as the extension field,
/// and `Radix2DitParallel<Felt>` as the DFT. Only the LMCS commitment scheme (`L`) and
/// Fiat-Shamir challenger (`Ch`) vary by hash function.
pub type MidenStarkConfig<L, Ch> =
    GenericStarkConfig<Felt, QuadFelt, L, Radix2DitParallel<Felt>, Ch>;

type PackedFelt = <Felt as Field>::Packing;

/// Number of inputs to the Merkle compression function.
const COMPRESSION_INPUTS: usize = 2;

// PCS PARAMETERS
// ================================================================================================

/// Log2 of the FRI blowup factor (blowup = 8).
pub(crate) const LOG_BLOWUP: u8 = 3;
/// Log2 of the FRI folding arity (arity = 4).
pub const LOG_FOLDING_ARITY: u8 = 2;
/// Log2 of the final polynomial degree (degree = 128).
const LOG_FINAL_DEGREE: u8 = 7;
/// Proof-of-work bits for FRI folding challenges.
pub const FOLDING_POW_BITS: usize = 4;
/// Proof-of-work bits for DEEP composition polynomial.
pub const DEEP_POW_BITS: usize = 12;
/// Number of FRI query repetitions.
const NUM_QUERIES: usize = 27;
/// Proof-of-work bits for the query phase. Together with 27 queries, this gives the query round
/// exactly 96 conjectured bits; lowering either value drops the preset below that target.
const QUERY_POW_BITS: usize = 17;

/// Default PCS parameters shared by all hash function configurations.
pub fn pcs_params() -> PcsParams {
    PcsParams::new(
        LOG_BLOWUP,
        LOG_FOLDING_ARITY,
        LOG_FINAL_DEGREE,
        FOLDING_POW_BITS,
        DEEP_POW_BITS,
        NUM_QUERIES,
        QUERY_POW_BITS,
    )
    .expect("invalid PCS parameters")
}

// DOMAIN-SEPARATED FIAT-SHAMIR TRANSCRIPT
// ================================================================================================

/// Relation digest absorbed into the Fiat-Shamir transcript domain separator.
pub type RelationDigest = [Felt; 4];

/// Bind a protocol version and ACE registry root into a relation digest.
pub fn relation_digest(protocol_id: u64, registry_root: &Word) -> RelationDigest {
    let input = [
        Felt::new_unchecked(protocol_id),
        registry_root[0],
        registry_root[1],
        registry_root[2],
        registry_root[3],
    ];
    let digest = Poseidon2::hash_elements(&input);
    let elements = digest.as_elements();
    [elements[0], elements[1], elements[2], elements[3]]
}

/// RELATION_DIGEST = Poseidon2::hash_elements([PROTOCOL_ID, ACE_CIRCUIT_REGISTRY_ROOT]).
///
/// Compile-time constant binding the Fiat-Shamir transcript to the Miden VM AIR.
/// Must match the constants in `crates/lib/core/asm/sys/vm/mod.masm`.
pub const RELATION_DIGEST: RelationDigest = [
    Felt::new_unchecked(8509919582315365814),
    Felt::new_unchecked(12865978328275266043),
    Felt::new_unchecked(12073117316737140237),
    Felt::new_unchecked(5647109555087128169),
];

/// Root of the accepted ACE circuit registry.
///
/// Active leaves are ACE circuit commitments indexed by `ProofOrder::tag()`.
pub const ACE_CIRCUIT_REGISTRY_ROOT: [Felt; 4] = [
    Felt::new_unchecked(8563712008625779037),
    Felt::new_unchecked(2346635758357605862),
    Felt::new_unchecked(17725435322360039753),
    Felt::new_unchecked(12106900609136090006),
];

/// Smallest ACE circuit registry depth covering every proof-order tag.
///
/// With `n` AIRs, proof-order tags range over the `n!` AIR permutations.
pub const ACE_CIRCUIT_REGISTRY_DEPTH: usize = PROOF_ORDER_REGISTRY_DEPTH;

/// Number of leaves in the ACE circuit registry tree.
pub const ACE_CIRCUIT_REGISTRY_LEAF_COUNT: usize = 1 << ACE_CIRCUIT_REGISTRY_DEPTH;
const _: () = assert!(
    PROOF_ORDER_COUNT <= ACE_CIRCUIT_REGISTRY_LEAF_COUNT,
    "ACE_CIRCUIT_REGISTRY_DEPTH must cover every proof-order variant",
);

// NOTE: registry leaves are not checked in. They are recomputed per process from the AIR
// (see `ace_registry_path`) and authenticated against `ACE_CIRCUIT_REGISTRY_ROOT`, which
// is the registry commitment. Checking in the six active Miden VM leaves would be cheap but
// redundant; recomputing them also keeps the Miden VM and PVM on one serving model.

/// Authentication data for one ACE registry slot: its leaf and Merkle path.
///
/// This is the whole registry read surface — the MASM loader consumes exactly one
/// `mtree_get(ACE_REGISTRY_ROOT, ORDER_TAG)`, so one `(leaf, path)` per proof is all a
/// caller ever needs. How the registry is stored behind this call is an implementation
/// detail, which is what lets the precompile VM's 10!-leaf registry swap in a different
/// strategy without touching callers.
///
/// Returns `None` when `tag` does not address a slot of the depth-[`ACE_CIRCUIT_REGISTRY_DEPTH`]
/// tree. Padding slots (tags at or above [`PROOF_ORDER_COUNT`]) resolve like any other slot;
/// the MASM verifier's `assert_valid_order_tag` is what keeps them from being opened.
pub fn ace_registry_path(tag: u32) -> Option<(Word, MerklePath)> {
    #[cfg(feature = "std")]
    let tree = miden_vm_ace_registry();
    #[cfg(not(feature = "std"))]
    let tree = &build_miden_vm_ace_registry();
    registry_path_in(tree, tag)
}

/// Leaf and path for `tag` in a caller-held registry tree.
pub(crate) fn registry_path_in(tree: &MerkleTree, tag: u32) -> Option<(Word, MerklePath)> {
    let index = NodeIndex::new(ACE_CIRCUIT_REGISTRY_DEPTH as u8, u64::from(tag)).ok()?;
    let leaf = tree.get_node(index).ok()?;
    let path = tree.get_path(index).ok()?;
    Some((leaf, path))
}

/// The process-wide Miden VM ACE circuit registry, computed from the AIR on first use.
///
/// Leaves are recomputed (not checked in) and the resulting root is authenticated against
/// the compiled-in [`ACE_CIRCUIT_REGISTRY_ROOT`] before anything can read the tree, so
/// this cache carries no trust: drift between the AIR and the protocol constant fails
/// here, loudly, instead of at every later `mtree_get` with an opaque store miss.
#[cfg(feature = "std")]
fn miden_vm_ace_registry() -> &'static MerkleTree {
    static REGISTRY: std::sync::OnceLock<MerkleTree> = std::sync::OnceLock::new();
    REGISTRY
        .get_or_init(|| build_miden_vm_ace_registry_with(crate::ace::shared_recursive_factory()))
}

/// Rebuilds the Miden VM's eight-leaf ACE registry and authenticates it against the protocol root.
///
/// `std` caches this tree process-wide; `no_std` callers rebuild it on demand so recursive advice
/// generation remains available without global state or synchronization support.
#[cfg(not(feature = "std"))]
fn build_miden_vm_ace_registry() -> MerkleTree {
    let factory = crate::ace::RecursiveAceCircuitFactory::new()
        .expect("recursive-verifier ACE composition must build");
    build_miden_vm_ace_registry_with(&factory)
}

/// Builds the eight-leaf registry from an existing factory and authenticates it against
/// the protocol root, so callers holding a factory pay no second composition build.
pub(crate) fn build_miden_vm_ace_registry_with(
    factory: &crate::ace::RecursiveAceCircuitFactory,
) -> MerkleTree {
    let mut buffer = miden_ace_codegen::ShuffleEncodeBuffer::new();
    let mut leaves = vec![miden_ace_codegen::padding_leaf(); ACE_CIRCUIT_REGISTRY_LEAF_COUNT];
    for order in ProofOrder::variants() {
        let leaf = factory
            .leaf_for_order(&order, &mut buffer)
            .expect("registry leaf must encode for every proof order");
        leaves[order.tag() as usize] = leaf;
    }
    let tree = MerkleTree::new(&leaves).expect("ACE circuit registry has power-of-two leaves");
    assert_eq!(
        tree.root(),
        Word::new(ACE_CIRCUIT_REGISTRY_ROOT),
        "computed ACE registry root does not match ACE_CIRCUIT_REGISTRY_ROOT. If this \
         protocol change is intentional, run `make regenerate-constraints` and record \
         the break; otherwise inspect the unexpected registry drift before updating constants",
    );
    tree
}

/// Observes PCS protocol parameters into the challenger.
///
/// Call on a challenger obtained from `config.challenger()` to complete the
/// domain-separated transcript initialization. The config factories bind the
/// caller-supplied relation digest into the prototype challenger; this function
/// adds the actual PCS parameters used by that config.
pub fn observe_protocol_params(params: &PcsParams, challenger: &mut impl CanObserve<Felt>) {
    // Batch 1: PCS parameters, zero-padded to SPONGE_RATE.
    challenger.observe(Felt::new_unchecked(params.num_queries() as u64));
    challenger.observe(Felt::new_unchecked(params.query_pow_bits() as u64));
    challenger.observe(Felt::new_unchecked(params.deep_pow_bits() as u64));
    challenger.observe(Felt::new_unchecked(params.folding_pow_bits() as u64));
    challenger.observe(Felt::new_unchecked(params.log_blowup() as u64));
    challenger.observe(Felt::new_unchecked(params.log_final_degree() as u64));
    challenger.observe(Felt::new_unchecked(1_u64 << params.log_folding_arity()));
    challenger.observe(Felt::ZERO);
}

// ALGEBRAIC HASHES (RPO, Poseidon2, RPX)
// ================================================================================================

/// Sponge state width in field elements.
const SPONGE_WIDTH: usize = 12;
/// Sponge rate (absorbable elements per permutation).
pub(crate) const SPONGE_RATE: usize = 8;
/// Sponge digest width in field elements.
pub(crate) const DIGEST_WIDTH: usize = 4;
/// Range of capacity slots within the sponge state array.
const CAPACITY_RANGE: core::ops::Range<usize> = SPONGE_RATE..SPONGE_WIDTH;

/// Algebraic LMCS (for RPO, Poseidon2, RPX).
type AlgLmcs<P> = LmcsConfig<
    PackedFelt,
    PackedFelt,
    StatefulSponge<P, SPONGE_WIDTH, SPONGE_RATE, DIGEST_WIDTH>,
    TruncatedPermutation<P, COMPRESSION_INPUTS, DIGEST_WIDTH, SPONGE_WIDTH>,
    SPONGE_WIDTH,
    DIGEST_WIDTH,
>;

/// Algebraic duplex challenger (for RPO, Poseidon2, RPX).
type AlgChallenger<P> = DuplexChallenger<Felt, P, SPONGE_WIDTH, SPONGE_RATE>;

/// Concrete STARK configuration type for RPO.
pub type RpoConfig = MidenStarkConfig<AlgLmcs<RpoPermutation256>, AlgChallenger<RpoPermutation256>>;

/// Concrete STARK configuration type for Poseidon2.
pub type Poseidon2Config =
    MidenStarkConfig<AlgLmcs<Poseidon2Permutation256>, AlgChallenger<Poseidon2Permutation256>>;

/// Concrete STARK configuration type for RPX.
pub type RpxConfig = MidenStarkConfig<AlgLmcs<RpxPermutation256>, AlgChallenger<RpxPermutation256>>;

/// Creates an RPO-based STARK configuration bound to `relation_digest`.
pub fn rpo_config(params: PcsParams, relation_digest: RelationDigest) -> RpoConfig {
    alg_config(params, RpoPermutation256, relation_digest)
}

/// Creates a Poseidon2-based STARK configuration bound to `relation_digest`.
pub fn poseidon2_config(params: PcsParams, relation_digest: RelationDigest) -> Poseidon2Config {
    alg_config(params, Poseidon2Permutation256, relation_digest)
}

/// Creates an RPX-based STARK configuration bound to `relation_digest`.
pub fn rpx_config(params: PcsParams, relation_digest: RelationDigest) -> RpxConfig {
    alg_config(params, RpxPermutation256, relation_digest)
}

/// Internal helper: builds an algebraic STARK configuration from a permutation.
///
/// The prototype challenger has the relation digest pre-loaded in the sponge capacity.
/// When `observe_protocol_params` is called, the first duplexing permutes this
/// capacity together with the PCS parameters written into the rate.
fn alg_config<P>(
    params: PcsParams,
    perm: P,
    relation_digest: RelationDigest,
) -> MidenStarkConfig<AlgLmcs<P>, AlgChallenger<P>>
where
    P: CryptographicPermutation<[Felt; SPONGE_WIDTH]> + Copy,
{
    let lmcs = LmcsConfig::new(StatefulSponge::new(perm), TruncatedPermutation::new(perm));
    let mut state = [Felt::ZERO; SPONGE_WIDTH];
    state[CAPACITY_RANGE].copy_from_slice(&relation_digest);
    let challenger = DuplexChallenger {
        sponge_state: state,
        input_buffer: vec![],
        output_buffer: vec![],
        permutation: perm,
    };
    GenericStarkConfig::new(params, lmcs, Radix2DitParallel::default(), challenger)
}

// BLAKE3
// ================================================================================================

/// Digest size in bytes for Blake3.
const BLAKE_DIGEST_SIZE: usize = 32;

/// Blake3 LMCS.
type BlakeLmcs = LmcsConfig<
    Felt,
    u8,
    ChainingHasher<Blake3Hasher>,
    CompressionFunctionFromHasher<Blake3Hasher, COMPRESSION_INPUTS, BLAKE_DIGEST_SIZE>,
    BLAKE_DIGEST_SIZE,
    BLAKE_DIGEST_SIZE,
>;

/// Blake3 challenger.
type BlakeChallenger =
    SerializingChallenger64<Felt, HashChallenger<u8, Blake3Hasher, BLAKE_DIGEST_SIZE>>;

/// Concrete STARK configuration type for Blake3.
pub type Blake3Config = MidenStarkConfig<BlakeLmcs, BlakeChallenger>;

/// Creates a Blake3_256-based STARK configuration bound to `relation_digest`.
pub fn blake3_256_config(params: PcsParams, relation_digest: RelationDigest) -> Blake3Config {
    let lmcs = LmcsConfig::new(
        ChainingHasher::new(Blake3Hasher),
        CompressionFunctionFromHasher::new(Blake3Hasher),
    );
    let mut challenger = SerializingChallenger64::from_hasher(vec![], Blake3Hasher);
    challenger.observe_slice(&relation_digest);
    GenericStarkConfig::new(params, lmcs, Radix2DitParallel::default(), challenger)
}

// KECCAK
// ================================================================================================

/// Keccak permutation state width (in u64 elements).
const KECCAK_WIDTH: usize = 25;
/// Keccak sponge rate (absorbable u64 elements per permutation).
const KECCAK_RATE: usize = 17;
/// Keccak digest width (in u64 elements).
const KECCAK_DIGEST: usize = 4;
/// Keccak-256 digest size in bytes (for the Fiat-Shamir challenger).
const KECCAK_CHALLENGER_DIGEST_SIZE: usize = 32;

/// Keccak MMCS sponge (padding-free, used for compression).
type KeccakMmcsSponge = PaddingFreeSponge<KeccakF, KECCAK_WIDTH, KECCAK_RATE, KECCAK_DIGEST>;

/// Keccak LMCS using the stateful binary sponge with `[Felt; VECTOR_LEN]` packing.
type KeccakLmcs = LmcsConfig<
    [Felt; VECTOR_LEN],
    [u64; VECTOR_LEN],
    SerializingStatefulSponge<StatefulSponge<KeccakF, KECCAK_WIDTH, KECCAK_RATE, KECCAK_DIGEST>>,
    CompressionFunctionFromHasher<KeccakMmcsSponge, COMPRESSION_INPUTS, KECCAK_DIGEST>,
    KECCAK_WIDTH,
    KECCAK_DIGEST,
>;

/// Keccak challenger.
type KeccakChallenger =
    SerializingChallenger64<Felt, HashChallenger<u8, Keccak256Hash, KECCAK_CHALLENGER_DIGEST_SIZE>>;

/// Concrete STARK configuration type for Keccak.
pub type KeccakConfig = MidenStarkConfig<KeccakLmcs, KeccakChallenger>;

/// Creates a Keccak-based STARK configuration.
///
/// Uses the stateful binary sponge with the Keccak permutation and `[Felt; VECTOR_LEN]` packing
/// for SIMD parallelization.
pub fn keccak_config(params: PcsParams, relation_digest: RelationDigest) -> KeccakConfig {
    let mmcs_sponge = KeccakMmcsSponge::new(KeccakF {});
    let compress = CompressionFunctionFromHasher::new(mmcs_sponge);
    let sponge = SerializingStatefulSponge::new(StatefulSponge::new(KeccakF {}));
    let lmcs = LmcsConfig::new(sponge, compress);
    let mut challenger = SerializingChallenger64::from_hasher(vec![], Keccak256Hash {});
    challenger.observe_slice(&relation_digest);
    GenericStarkConfig::new(params, lmcs, Radix2DitParallel::default(), challenger)
}

#[cfg(test)]
mod tests {
    extern crate alloc;
    use alloc::vec::Vec;

    use miden_ace_codegen::padding_leaf;
    use miden_core::{Felt, Word};
    use miden_crypto::{
        merkle::MerkleTree,
        stark::{challenger::CanObserve, pcs::PcsParams},
    };

    use crate::{ProofOrder, ace};

    const PROTOCOL_ID: u64 = 1;
    const REGEN_HINT: &str = "cargo run -p miden-core-lib --features constraints-tools --bin regenerate-constraints -- --write";

    #[derive(Default)]
    struct RecordingChallenger(Vec<Felt>);

    impl CanObserve<Felt> for RecordingChallenger {
        fn observe(&mut self, value: Felt) {
            self.0.push(value);
        }
    }

    /// Transcript domain separation must bind the parameters actually supplied to the config,
    /// not the Miden VM's current compile-time defaults.
    #[test]
    fn protocol_observation_uses_the_supplied_pcs_params() {
        let params = PcsParams::new(4, 3, 6, 5, 11, 19, 13).expect("valid distinct PCS params");
        let mut challenger = RecordingChallenger::default();
        super::observe_protocol_params(&params, &mut challenger);
        assert_eq!(
            challenger.0,
            [19, 13, 11, 5, 4, 6, 8, 0].map(Felt::new_unchecked),
            "the transcript must encode [queries, query PoW, DEEP PoW, folding PoW, blowup log, \
             final-degree log, folding arity, padding]",
        );
    }

    /// Snapshot test: catches any AIR change that alters the constraint circuit.
    ///
    /// If this test fails, regenerate with:
    /// ```text
    /// cargo run -p miden-core-lib --features constraints-tools --bin regenerate-constraints -- --write
    /// ```
    #[test]
    fn relation_digest_matches_current_air() {
        let mut expected_leaves = vec![padding_leaf(); super::ACE_CIRCUIT_REGISTRY_LEAF_COUNT];
        let mut snapshot_lines = Vec::new();
        let mut expected_metadata = None;

        let factory = ace::RecursiveAceCircuitFactory::new().unwrap();
        let mut buffer = miden_ace_codegen::ShuffleEncodeBuffer::new();
        for order in ProofOrder::variants() {
            let circuit = factory.circuit_for_order(&order).unwrap();

            // Dual-path leaf equality for EVERY order: the encode-only path (which the
            // runtime registry uses) against the assembled-stream path. This catches
            // encoding divergence between the two stream constructions; it is blind to
            // hash-implementation faults, which hit both paths identically (they share
            // the cached sponge states) — those are guarded by the one-shot builder
            // sweep in air/tests/ace_codegen.rs and miden-crypto's packed-vs-scalar
            // differential test.
            assert_eq!(
                factory.leaf_for_order(&order, &mut buffer).unwrap(),
                circuit.commitment,
                "encode-only registry leaf diverges from the assembled circuit for {}",
                order.file_stem(),
            );
            let metadata = (circuit.num_inputs, circuit.num_eval_gates, circuit.stream_len);
            if let Some(expected) = expected_metadata {
                assert_eq!(metadata, expected, "ACE circuit metadata must be uniform");
            } else {
                expected_metadata = Some(metadata);
            }

            let tag = order.tag() as usize;
            assert!(tag < expected_leaves.len(), "proof-order tag does not fit registry tree");
            expected_leaves[tag] = circuit.commitment;

            let commitment: Vec<u64> =
                circuit.commitment.iter().map(Felt::as_canonical_u64).collect();
            snapshot_lines.push(format!(
                "{}:\n  num_inputs: {}\n  num_eval_gates: {}\n  stream_len: {}\n  commitment: {:?}",
                order.file_stem(),
                circuit.num_inputs,
                circuit.num_eval_gates,
                circuit.stream_len,
                commitment,
            ));
        }

        let tree = MerkleTree::new(&expected_leaves).expect("registry tree");
        let registry_root = tree.root();
        assert_eq!(
            Word::new(super::ACE_CIRCUIT_REGISTRY_ROOT),
            registry_root,
            "ACE_CIRCUIT_REGISTRY_ROOT in config.rs is stale. Regenerate with: {REGEN_HINT}"
        );

        // The path-shaped read surface must serve every slot — active and padding —
        // with a leaf and path that verify against the compiled-in root.
        for tag in 0..super::ACE_CIRCUIT_REGISTRY_LEAF_COUNT as u32 {
            let (leaf, path) = super::ace_registry_path(tag).expect("tag addresses a slot");
            assert_eq!(leaf, expected_leaves[tag as usize], "leaf mismatch at tag {tag}");
            let computed = path.compute_root(u64::from(tag), leaf).expect("path root computes");
            assert_eq!(computed, registry_root, "path at tag {tag} does not verify");
        }
        assert!(
            super::ace_registry_path(super::ACE_CIRCUIT_REGISTRY_LEAF_COUNT as u32).is_none(),
            "out-of-range tags must not resolve"
        );

        let digest = super::relation_digest(PROTOCOL_ID, &registry_root);
        let expected: Vec<u64> = digest.iter().map(Felt::as_canonical_u64).collect();

        let snapshot = format!("{}\nrelation_digest: {:?}", snapshot_lines.join("\n"), expected);
        insta::assert_snapshot!(snapshot);

        let actual: Vec<u64> = super::RELATION_DIGEST.iter().map(Felt::as_canonical_u64).collect();
        assert_eq!(
            actual, expected,
            "RELATION_DIGEST in config.rs is stale. Regenerate with: {REGEN_HINT}"
        );
    }
}