soft3/zheng/rs/src/execution/statement.rs

//! Canonical, bounded public statements; no trace or prover-defined matrices.
use super::relation::{ExecutionNoun, ExecutionRelation, SubjectShape, compile_relation};
use super::{DirectProof, proof};
use nebu::Goldilocks as F;
use std::collections::BTreeMap;

pub const MAX_PROGRAM_NODES: usize = 4096;
pub const MAX_DEPTH: usize = 128;
pub const MAX_INPUTS: usize = 64;
pub const MAX_OUTPUTS: usize = 4096;

/// Flat prefix encoding avoids recursive untrusted deserialization.
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum NounToken {
    Atom(u64),
    Pair,
}

#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ExecutionStatement {
    #[cfg_attr(
        feature = "serde",
        serde(deserialize_with = "super::statement_wire::program")
    )]
    pub program: Vec<NounToken>,
    #[cfg_attr(
        feature = "serde",
        serde(deserialize_with = "super::statement_wire::inputs")
    )]
    pub public_input: Vec<u64>,
    #[cfg_attr(
        feature = "serde",
        serde(deserialize_with = "super::statement_wire::outputs")
    )]
    pub public_output: Vec<u64>,
    pub cycles: u64,
    pub budget: u64,
}

fn canonical(values: &[u64]) -> bool {
    values.iter().all(|&v| v < nebu::field::P)
}

impl ExecutionStatement {
    pub fn program_noun(&self) -> Result<ExecutionNoun, String> {
        if self.program.is_empty() || self.program.len() > MAX_PROGRAM_NODES {
            return Err("program size limit".into());
        }
        let mut stack: Vec<(ExecutionNoun, usize)> = Vec::new();
        for token in self.program.iter().rev() {
            match token {
                NounToken::Atom(v) => {
                    if *v >= nebu::field::P {
                        return Err("noncanonical program atom".into());
                    }
                    stack.push((ExecutionNoun::Atom(*v), 0));
                }
                NounToken::Pair => {
                    let (a, da) = stack.pop().ok_or("malformed program prefix")?;
                    let (b, db) = stack.pop().ok_or("malformed program prefix")?;
                    let depth = 1 + da.max(db);
                    if depth > MAX_DEPTH {
                        return Err("program depth limit".into());
                    }
                    stack.push((ExecutionNoun::Pair(Box::new(a), Box::new(b)), depth));
                }
            }
        }
        if stack.len() != 1 {
            return Err("trailing program tokens".into());
        }
        stack
            .pop()
            .map(|(n, _)| n)
            .ok_or_else(|| "empty program".into())
    }

    pub(crate) fn validate_bounds(&self) -> Result<(), String> {
        if self.public_input.len() > MAX_INPUTS
            || self.public_output.len() > MAX_OUTPUTS
            || !canonical(&self.public_input)
            || !canonical(&self.public_output)
            || self.budget >= nebu::field::P
            || self.cycles > self.budget
        {
            return Err("invalid public statement bounds or field values".into());
        }
        Ok(())
    }

    pub(crate) fn relation(&self) -> Result<ExecutionRelation, String> {
        self.validate_bounds()?;
        let mut shape = SubjectShape::Atom;
        for _ in &self.public_input {
            shape = SubjectShape::Pair(Box::new(SubjectShape::Atom), Box::new(shape));
        }
        let relation = compile_relation(&self.program_noun()?, &shape)
            .map_err(|e| format!("unsupported execution relation: {e:?}"))?;
        if !relation.lookups.is_empty() {
            return Err("state lookup requires authenticated state execution protocol".into());
        }
        Ok(relation)
    }

    pub(crate) fn inputs(&self) -> Vec<F> {
        self.public_input
            .iter()
            .rev()
            .map(|&v| F::new(v))
            .chain([F::ZERO])
            .collect()
    }

    pub(crate) fn bindings(&self, relation: &ExecutionRelation) -> Result<Vec<(usize, F)>, String> {
        self.bindings_with_inputs(relation, &self.inputs())
    }
    pub(crate) fn bindings_with_inputs(
        &self,
        relation: &ExecutionRelation,
        inputs: &[F],
    ) -> Result<Vec<(usize, F)>, String> {
        if inputs.len() != relation.input_indices.len() {
            return Err("input shape mismatch".into());
        }
        if self.public_output.len() != relation.output_indices.len() {
            return Err("output shape does not match program".into());
        }
        let mut result = BTreeMap::new();
        for (i, v) in relation
            .input_indices
            .iter()
            .copied()
            .zip(inputs.iter().copied())
            .chain(
                relation
                    .output_indices
                    .iter()
                    .copied()
                    .zip(self.public_output.iter().map(|&v| F::new(v))),
            )
            .chain([(relation.cost_index, F::new(self.cycles))])
        {
            if result.insert(i, v).is_some_and(|old| old != v) {
                return Err("conflicting public coordinates".into());
            }
        }
        Ok(result.into_iter().collect())
    }

    /// What the relation depends on, hashed: see [`super::vk`].
    pub fn program_key(&self) -> [u8; 32] {
        super::vk::program_key(
            &self.program,
            self.public_input.len(),
            super::vk::StatementKind::Execution,
        )
    }

    /// The relation from `vk` when it was derived for this statement's
    /// program key, compiled otherwise; bounds are always checked.
    pub(crate) fn keyed<'a>(
        &self,
        vk: Option<&'a super::VerifyingKey>,
    ) -> Result<std::borrow::Cow<'a, super::VerifyingKey>, String> {
        match vk {
            Some(vk) => {
                self.validate_bounds()?;
                Ok(std::borrow::Cow::Borrowed(vk.check(self.program_key())?))
            }
            None => Ok(std::borrow::Cow::Owned(super::VerifyingKey::for_execution(
                self,
            )?)),
        }
    }

    /// Stable encoding independent of serde, usize width and JSON formatting.
    pub fn transcript_bytes(&self) -> Vec<u8> {
        let mut bytes = b"zheng-nox-public-execution-v2".to_vec();
        bytes.extend_from_slice(&(self.program.len() as u64).to_le_bytes());
        for token in &self.program {
            match token {
                NounToken::Atom(v) => {
                    bytes.push(0);
                    bytes.extend_from_slice(&v.to_le_bytes());
                }
                NounToken::Pair => bytes.push(1),
            }
        }
        for values in [&self.public_input, &self.public_output] {
            bytes.extend_from_slice(&(values.len() as u64).to_le_bytes());
            for v in values {
                bytes.extend_from_slice(&v.to_le_bytes());
            }
        }
        bytes.extend_from_slice(&self.cycles.to_le_bytes());
        bytes.extend_from_slice(&self.budget.to_le_bytes());
        bytes
    }

    pub fn encode_program(noun: &ExecutionNoun) -> Result<Vec<NounToken>, String> {
        let mut tokens = Vec::new();
        let mut pending = vec![(noun, 0)];
        while let Some((node, depth)) = pending.pop() {
            if depth > MAX_DEPTH || tokens.len() >= MAX_PROGRAM_NODES {
                return Err("program size/depth limit".into());
            }
            match node {
                ExecutionNoun::Atom(v) => {
                    if *v >= nebu::field::P {
                        return Err("noncanonical program atom".into());
                    }
                    tokens.push(NounToken::Atom(*v));
                }
                ExecutionNoun::Pair(a, b) => {
                    tokens.push(NounToken::Pair);
                    pending.push((b, depth + 1));
                    pending.push((a, depth + 1));
                }
            }
        }
        Ok(tokens)
    }
}

/// The statement, relation, witness and pinned coordinates of one honest
/// public execution; shared by every public-statement prover (v2, v3 and
/// the succinct profile).
#[allow(clippy::type_complexity)]
pub(crate) fn prepare(
    program: &ExecutionNoun,
    input: &[u64],
    budget: u64,
) -> Result<
    (
        ExecutionStatement,
        ExecutionRelation,
        crate::types::CCSWitness,
        Vec<(usize, F)>,
    ),
    String,
> {
    let mut statement = ExecutionStatement {
        program: ExecutionStatement::encode_program(program)?,
        public_input: input.to_vec(),
        public_output: vec![],
        cycles: 0,
        budget,
    };
    let relation = statement.relation()?;
    let witness = relation
        .witness(&statement.inputs())
        .map_err(|e| format!("execution witness: {e:?}"))?;
    statement.public_output = relation
        .output_indices
        .iter()
        .map(|&i| witness.z[i].as_u64())
        .collect();
    statement.cycles = witness.z[relation.cost_index].as_u64();
    if statement.cycles > budget {
        return Err("execution cost exceeds budget".into());
    }
    let public = statement.bindings(&relation)?;
    Ok((statement, relation, witness, public))
}

/// Input values and the witness are public in this certificate. Secret-input
/// proving uses [`super::private::prepare_execution`] and [`super::zk::prove`];
/// see `specs/native-private-ccs.md` for its disclosure and security contract.
pub fn prove_execution(
    program: &ExecutionNoun,
    input: &[u64],
    budget: u64,
) -> Result<(ExecutionStatement, DirectProof), String> {
    let (statement, relation, witness, public) = prepare(program, input, budget)?;
    let proof = proof::prove(
        &relation.instance,
        &witness,
        &statement.transcript_bytes(),
        &public,
    )
    .map_err(|e| e.to_string())?;
    Ok((statement, proof))
}

/// Derives the full relation and verifies the claimed public values. This does
/// not run nox, construct a witness, or accept a supplied execution trace.
pub fn verify_execution(statement: &ExecutionStatement, proof: &DirectProof) -> Result<(), String> {
    let relation = statement.relation()?;
    proof::verify(
        &relation.instance,
        proof,
        &statement.transcript_bytes(),
        &statement.bindings(&relation)?,
    )
    .map_err(|e| e.to_string())
}

/// Profile v3: prove by producing the free witness positions. The statement is
/// the same as profile v2; the certificate replaces the Spartan proof.
pub fn certify_execution(
    program: &ExecutionNoun,
    input: &[u64],
    budget: u64,
) -> Result<(ExecutionStatement, super::certificate::Certificate), String> {
    let (statement, relation, witness, public) = prepare(program, input, budget)?;
    let certificate = super::certificate::certify(&relation.instance, &witness, &public)
        .map_err(|e| e.to_string())?;
    Ok((statement, certificate))
}

/// Profile v3: derive the relation from the statement and check the
/// certificate exactly. Runs no nox, accepts no relation from the prover.
pub fn verify_certificate(
    statement: &ExecutionStatement,
    certificate: &super::certificate::Certificate,
) -> Result<(), String> {
    let relation = statement.relation()?;
    super::certificate::verify(
        &relation.instance,
        certificate,
        &statement.bindings(&relation)?,
    )
    .map_err(|e| e.to_string())
}

/// [`verify_certificate`] with a cached verifying key: `vk` must have been
/// derived for this statement's program key, or the statement is rejected.
pub fn verify_certificate_with(
    statement: &ExecutionStatement,
    certificate: &super::certificate::Certificate,
    vk: &super::VerifyingKey,
) -> Result<(), String> {
    statement.validate_bounds()?;
    let relation = vk.check(statement.program_key())?.relation();
    super::certificate::verify(
        &relation.instance,
        certificate,
        &statement.bindings(relation)?,
    )
    .map_err(|e| e.to_string())
}