mod call_metrics;
pub use call_metrics::*;
mod inclusion;
pub use inclusion::*;
use algorithms::snark::varuna::VarunaVersion;
use circuit::Assignment;
use console::{
network::prelude::*,
program::{InputID, Locator},
};
use ledger_block::{Execution, Fee, Transition};
use ledger_query::QueryTrait;
use synthesizer_snark::{Proof, ProvingKey, VerifyingKey};
use once_cell::sync::OnceCell;
use std::collections::HashMap;
#[derive(Clone, Debug, Default)]
pub struct Trace<N: Network> {
transitions: Vec<Transition<N>>,
transition_tasks: HashMap<Locator<N>, (ProvingKey<N>, Vec<Assignment<N::Field>>)>,
inclusion_tasks: Inclusion<N>,
call_metrics: Vec<CallMetrics<N>>,
inclusion_assignments: OnceCell<Vec<InclusionAssignmentWrapper<N>>>,
global_state_root: OnceCell<N::StateRoot>,
}
impl<N: Network> Trace<N> {
pub fn new() -> Self {
Self {
transitions: Vec::new(),
transition_tasks: HashMap::new(),
inclusion_tasks: Inclusion::new(),
inclusion_assignments: OnceCell::new(),
global_state_root: OnceCell::new(),
call_metrics: Vec::new(),
}
}
pub fn transitions(&self) -> &[Transition<N>] {
&self.transitions
}
pub fn call_metrics(&self) -> &[CallMetrics<N>] {
&self.call_metrics
}
}
impl<N: Network> Trace<N> {
pub fn insert_transition(
&mut self,
input_ids: &[InputID<N>],
transition: &Transition<N>,
(proving_key, assignment): (ProvingKey<N>, Assignment<N::Field>),
metrics: CallMetrics<N>,
) -> Result<()> {
ensure!(self.inclusion_assignments.get().is_none());
ensure!(self.global_state_root.get().is_none());
self.inclusion_tasks.insert_transition(input_ids, transition)?;
let locator = Locator::new(*transition.program_id(), *transition.function_name());
self.transition_tasks.entry(locator).or_insert((proving_key, vec![])).1.push(assignment);
self.transitions.push(transition.clone());
self.call_metrics.push(metrics);
Ok(())
}
}
impl<N: Network> Trace<N> {
pub fn is_fee(&self) -> bool {
self.is_fee_private() || self.is_fee_public()
}
pub fn is_fee_private(&self) -> bool {
self.transitions.len() == 1 && self.transitions[0].is_fee_private()
}
pub fn is_fee_public(&self) -> bool {
self.transitions.len() == 1 && self.transitions[0].is_fee_public()
}
pub fn is_upgrade(&self) -> bool {
self.transitions.len() == 1 && self.transitions[0].is_upgrade()
}
}
impl<N: Network> Trace<N> {
pub fn prepare(&mut self, query: &dyn QueryTrait<N>) -> Result<()> {
let (inclusion_assignments, global_state_root) = self.inclusion_tasks.prepare(&self.transitions, query)?;
self.inclusion_assignments
.set(inclusion_assignments)
.map_err(|_| anyhow!("Failed to set inclusion assignments"))?;
self.global_state_root.set(global_state_root).map_err(|_| anyhow!("Failed to set global state root"))?;
Ok(())
}
#[cfg(feature = "async")]
pub async fn prepare_async(&mut self, query: &dyn QueryTrait<N>) -> Result<()> {
let (inclusion_assignments, global_state_root) =
self.inclusion_tasks.prepare_async(&self.transitions, query).await?;
self.inclusion_assignments
.set(inclusion_assignments)
.map_err(|_| anyhow!("Failed to set inclusion assignments"))?;
self.global_state_root.set(global_state_root).map_err(|_| anyhow!("Failed to set global state root"))?;
Ok(())
}
pub fn prove_execution<A: circuit::Aleo<Network = N>, R: Rng + CryptoRng>(
&self,
locator: &str,
varuna_version: VarunaVersion,
rng: &mut R,
) -> Result<Execution<N>> {
ensure!(!self.is_fee(), "The trace cannot call 'prove_execution' for a fee type");
ensure!(
self.transitions.iter().all(|transition| !(transition.is_fee_private() || transition.is_fee_public())),
"The trace cannot prove execution for a fee, call 'prove_fee' instead"
);
let inclusion_assignments =
self.inclusion_assignments.get().ok_or_else(|| anyhow!("Inclusion assignments have not been set"))?;
let global_state_root =
self.global_state_root.get().ok_or_else(|| anyhow!("Global state root has not been set"))?;
let proving_tasks = self.transition_tasks.values().cloned().collect();
let (global_state_root, proof) = Self::prove_batch::<A, R>(
locator,
varuna_version,
proving_tasks,
inclusion_assignments,
*global_state_root,
rng,
)?;
Execution::from(self.transitions.iter().cloned(), global_state_root, Some(proof))
}
pub fn prove_fee<A: circuit::Aleo<Network = N>, R: Rng + CryptoRng>(
&self,
varuna_version: VarunaVersion,
rng: &mut R,
) -> Result<Fee<N>> {
let is_fee_public = self.is_fee_public();
let is_fee_private = self.is_fee_private();
ensure!(is_fee_public || is_fee_private, "The trace cannot call 'prove_fee' for an execution type");
let inclusion_assignments =
self.inclusion_assignments.get().ok_or_else(|| anyhow!("Inclusion assignments have not been set"))?;
match is_fee_public {
true => ensure!(inclusion_assignments.is_empty(), "Expected 0 inclusion assignments for proving the fee"),
false => ensure!(inclusion_assignments.len() == 1, "Expected 1 inclusion assignment for proving the fee"),
}
let global_state_root =
self.global_state_root.get().ok_or_else(|| anyhow!("Global state root has not been set"))?;
let fee_transition = &self.transitions[0];
let proving_tasks = self.transition_tasks.values().cloned().collect();
let (global_state_root, proof) = Self::prove_batch::<A, R>(
"credits.aleo/fee (private or public)",
varuna_version,
proving_tasks,
inclusion_assignments,
*global_state_root,
rng,
)?;
Ok(Fee::from_unchecked(fee_transition.clone(), global_state_root, Some(proof)))
}
pub fn verify_execution_proof(
locator: &str,
varuna_version: VarunaVersion,
inclusion_version: InclusionVersion,
verifier_inputs: Vec<(VerifyingKey<N>, Vec<Vec<N::Field>>)>,
execution: &Execution<N>,
) -> Result<()> {
let global_state_root = execution.global_state_root();
if global_state_root == N::StateRoot::default() {
bail!("Inclusion expected the global state root in the execution to *not* be zero")
}
let Some(proof) = execution.proof() else { bail!("Expected the execution to contain a proof") };
match Self::verify_batch(
locator,
varuna_version,
inclusion_version,
verifier_inputs,
global_state_root,
execution.transitions(),
proof,
) {
Ok(()) => Ok(()),
Err(e) => bail!("Execution is invalid - {e}"),
}
}
pub fn verify_fee_proof(
varuna_version: VarunaVersion,
inclusion_version: InclusionVersion,
verifier_inputs: (VerifyingKey<N>, Vec<Vec<N::Field>>),
fee: &Fee<N>,
) -> Result<()> {
let global_state_root = fee.global_state_root();
if global_state_root == N::StateRoot::default() {
bail!("Inclusion expected the global state root in the fee to *not* be zero")
}
let Some(proof) = fee.proof() else { bail!("Expected the fee to contain a proof") };
match Self::verify_batch(
"credits.aleo/fee (private or public)",
varuna_version,
inclusion_version,
vec![verifier_inputs],
global_state_root,
[fee.transition()].into_iter(),
proof,
) {
Ok(()) => Ok(()),
Err(e) => bail!("Fee is invalid - {e}"),
}
}
}
impl<N: Network> Trace<N> {
fn prove_batch<A: circuit::Aleo<Network = N>, R: Rng + CryptoRng>(
locator: &str,
varuna_version: VarunaVersion,
mut proving_tasks: Vec<(ProvingKey<N>, Vec<Assignment<N::Field>>)>,
inclusion_assignments: &[InclusionAssignmentWrapper<N>],
global_state_root: N::StateRoot,
rng: &mut R,
) -> Result<(N::StateRoot, Proof<N>)> {
if global_state_root == N::StateRoot::default() {
bail!("Inclusion expected the global state root in the execution to *not* be zero")
}
let mut batch_inclusions = Vec::with_capacity(inclusion_assignments.len());
let mut inclusion_version = None;
for assignment in inclusion_assignments.iter() {
match &mut inclusion_version {
None => inclusion_version = Some(assignment),
Some(expected) if std::mem::discriminant(expected) == std::mem::discriminant(&assignment) => {}
Some(_) => bail!("Inclusion version expected to be the same across iterations."),
}
let assignment = match assignment {
InclusionAssignmentWrapper::V0(assignment_v0) => {
if global_state_root != assignment_v0.state_path.global_state_root() {
bail!("Inclusion expected the global state root to be the same across iterations")
}
assignment_v0.to_circuit_assignment::<A>()?
}
InclusionAssignmentWrapper::V1(assignment_v1) => {
if global_state_root != assignment_v1.state_path.global_state_root() {
bail!("Inclusion expected the global state root to be the same across iterations")
}
assignment_v1.to_circuit_assignment::<A>()?
}
};
batch_inclusions.push(assignment);
}
if !batch_inclusions.is_empty() {
let proving_key = match inclusion_version {
Some(InclusionAssignmentWrapper::V0(..)) => ProvingKey::<N>::new(N::inclusion_v0_proving_key().clone()),
Some(InclusionAssignmentWrapper::V1(..)) => ProvingKey::<N>::new(N::inclusion_proving_key().clone()),
None => bail!("Invalid or missing inclusion version"),
};
proving_tasks.push((proving_key, batch_inclusions));
}
let proof = ProvingKey::prove_batch(locator, varuna_version, &proving_tasks, rng)?;
Ok((global_state_root, proof))
}
fn verify_batch<'a>(
locator: &str,
varuna_version: VarunaVersion,
inclusion_version: InclusionVersion,
mut verifier_inputs: Vec<(VerifyingKey<N>, Vec<Vec<N::Field>>)>,
global_state_root: N::StateRoot,
transitions: impl ExactSizeIterator<Item = &'a Transition<N>>,
proof: &Proof<N>,
) -> Result<()> {
let batch_inclusion_inputs =
Inclusion::prepare_verifier_inputs(global_state_root, inclusion_version, transitions)?;
if !batch_inclusion_inputs.is_empty() {
let verifying_key = match inclusion_version {
InclusionVersion::V0 => N::inclusion_v0_verifying_key().clone(),
InclusionVersion::V1 => N::inclusion_verifying_key().clone(),
};
let num_variables = verifying_key.circuit_info.num_public_and_private_variables as u64;
verifier_inputs.push((VerifyingKey::<N>::new(verifying_key, num_variables), batch_inclusion_inputs));
}
VerifyingKey::verify_batch(locator, varuna_version, verifier_inputs, proof)
.map_err(|e| anyhow!("Failed to verify proof - {e}"))
}
}