neptune-consensus 0.14.0

use neptune_primitives::network::Network;
use tasm_lib::triton_vm;
use tasm_lib::triton_vm::proof::Claim;
use tasm_lib::triton_vm::stark::Stark;
use tokio::task;

use crate::transaction::validity::neptune_proof::Proof;

/// Historical block claims that define the main-net checkpoint: one hex-encoded,
/// bincode-serialized [`Claim`] per line, each prefixed by its block height.
/// Feeding these into [`cache_true_claims`] marks the corresponding historical
/// blocks as valid without re-verifying their proofs.
pub const CHECKPOINT_MAIN: &str = include_str!("../assets/main/checkpoint.dat");

/// Historical block claims that define the testnet-0 checkpoint; see
/// [`CHECKPOINT_MAIN`] for the format.
pub const CHECKPOINT_TESTNET_0: &str = include_str!("../assets/testnet-0/checkpoint.dat");

/// This claims-cache contains claims that are simply defined to be true.
///
/// If the claim is in the cache, then the Triton VM verifier is by-passed
/// without reading the proof.
///
/// Besides tests, it is used for *checkpoints*, where we define historical
/// blocks to be valid.
static CLAIMS_CACHE: std::sync::LazyLock<tokio::sync::Mutex<std::collections::HashSet<Claim>>> =
    std::sync::LazyLock::new(|| tokio::sync::Mutex::new(std::collections::HashSet::new()));

static CLAIMS_CACHE_ENABLED: std::sync::LazyLock<tokio::sync::Mutex<bool>> =
    std::sync::LazyLock::new(|| tokio::sync::Mutex::new(true));

/// Verify a Triton VM (claim, proof) pair for default STARK parameters.
///
/// When the test flag is set, this function checks whether the claim is present
/// in the `CLAIMS_CACHE` and if so returns true early (*i.e.*, without running
/// the verifier). When the test flag is set and the cache does not contain the
/// claim and verification succeeds, the claim is added to the cache. The only
/// other way to populate the cache is through method `cache_true_claim`.
pub async fn verify(claim: Claim, proof: Proof, network: Network) -> bool {
    // security: we do not accept mock proofs unless we ourselves
    // are running a network that accepts mock-proofs, eg regtest.
    if network.use_mock_proof() {
        return proof.is_valid_mock();
    }

    {
        let is_enabled = *CLAIMS_CACHE_ENABLED.lock().await;
        if is_enabled && CLAIMS_CACHE.lock().await.contains(&claim) {
            return true;
        }
    }

    #[cfg(test)]
    let claim_clone = claim.clone();

    let verdict =
        task::spawn_blocking(move || triton_vm::verify(Stark::default(), &claim, &proof.into()))
            .await
            .expect("should be able to verify proof in new tokio task");

    // tbd: we might want to enable a cache for mainnet usage.
    // but we should probably use a cache that has a configurable max
    // size, so we don't blow up RAM.
    #[cfg(test)]
    if verdict {
        cache_true_claims([claim_clone]).await;
    }

    verdict
}

/// Add claims to the `CLAIMS_CACHE`.
pub async fn cache_true_claims<IterClaims: IntoIterator<Item = Claim>>(claims: IterClaims) {
    let mut cache = CLAIMS_CACHE.lock().await;
    for claim in claims {
        cache.insert(claim);
    }
}

/// Disable the true `CLAIMS_CACHE`.
#[cfg(any(test, feature = "test-helpers"))]
pub async fn disable_true_claims_cache() {
    *CLAIMS_CACHE_ENABLED.lock().await = false;
}

/// Enable the true `CLAIMS_CACHE`.
#[cfg(any(test, feature = "test-helpers"))]
pub async fn enable_true_claims_cache() {
    *CLAIMS_CACHE_ENABLED.lock().await = true;
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
pub(crate) mod tests {
    use itertools::Itertools;
    use macro_rules_attr::apply;
    use rand::Rng;
    use tasm_lib::prelude::Tip5;
    use triton_vm::prelude::BFieldCodec;

    use super::*;
    use crate::proof_abstractions::test_runtime::shared_tokio_runtime;

    pub(crate) fn bogus_proof(claim: &Claim) -> Proof {
        Proof::from(Tip5::hash_varlen(&claim.encode()).values().to_vec())
    }

    #[apply(shared_tokio_runtime)]
    async fn test_claims_cache() {
        let network = Network::Main;

        // generate random claim and bogus proof
        let mut rng = rand::rng();
        let some_claim = Claim::new(rng.random())
            .with_input((0..10).map(|_| rng.random()).collect_vec())
            .with_output((0..10).map(|_| rng.random()).collect_vec());
        let some_proof = bogus_proof(&some_claim);

        // verification must fail
        assert!(!verify(some_claim.clone(), some_proof.clone(), network).await);

        // put claim into cache
        cache_true_claims(vec![some_claim.clone()]).await;

        // verification must succeed
        assert!(verify(some_claim, some_proof, network).await);
    }
}