use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
use lens::WhirParams;
use nebu::{Fp3, Goldilocks};
use super::circuit::air::CircuitAir;
use super::circuit::builder::Builder;
use super::circuit::trace::{self, Pre};
use super::decide::{self, Decider, KeyWords};
use super::finalv::{self, Publics};
use super::ops::{Native, Ops};
use super::params::Params;
use super::perm::tag;
use super::program;
use super::prove::{self, AccData, StepInput, pbar_nox, pbar_v};
use super::relation::{G_POINT, Relation, WORD};
use super::sponge::Sponge;
use super::state::{self, CtxParts, State, ZeroWord};
use super::step::{self, StepProof};
use super::whir;
use super::word::{Digest, Word};
use crate::air::Public;
use crate::air::num::{Graph, record};
use crate::execution::ExecutionNoun;
use crate::machine::air::{Constants, Machine};
use crate::machine::layout::BLOCK;
use crate::machine::{self, MachineStatement, Run, statement};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct IvcProof {
pub log_rows: u32,
pub start: u64,
pub segments: u64,
pub chain: Digest,
pub state: State,
pub step: StepProof,
pub decider: Decider,
}
pub struct Key {
pub params: Params,
pub pre: Pre,
pub sparse: Vec<Vec<(u32, Fp3)>>,
pub kw_root: Digest,
kw: OnceLock<KeyWords>,
pub key_ext: bool,
pub dcfg: whir::Config,
pub g: Graph,
}
fn layout(p: &Params, air: &CircuitAir) -> Result<(Pre, usize), String> {
let mut b = Builder::new(false);
let live = b.live_var();
let parts = CtxParts {
statement: [Goldilocks::ZERO; 4],
chain: [Goldilocks::ZERO; 4],
g0: Fp3::ZERO,
pn0: Fp3::ZERO,
pv0: Fp3::ZERO,
};
program::run(
&mut b,
p,
live,
&program::dummy_state(p),
&program::dummy_proof(p),
&parts,
|b, c| b.set_output(c),
);
let (_, pre, out) = trace::generate(&b, air, p.n)?;
Ok((pre, out))
}
pub const SHIPPED_POW: u8 = 16;
pub const SHIPPED_RATE: u8 = 3;
pub fn plan(whir: &WhirParams, n: usize) -> Result<(Params, usize, [usize; 3]), String> {
let p = Params::new(whir, n)?;
let mut b = Builder::new(false);
let live = b.live_var();
let parts = CtxParts {
statement: [Goldilocks::ZERO; 4],
chain: [Goldilocks::ZERO; 4],
g0: Fp3::ZERO,
pn0: Fp3::ZERO,
pv0: Fp3::ZERO,
};
program::run(
&mut b,
&p,
live,
&program::dummy_state(&p),
&program::dummy_proof(&p),
&parts,
|b, c| b.set_output(c),
);
let census = [
b.gates.len(),
b.bits.len(),
b.chains.iter().map(|c| c.blocks.len()).sum(),
];
Ok((p, b.rows(), census))
}
type Cache = Mutex<HashMap<([u8; 8], usize), Arc<Key>>>;
fn key_cache() -> &'static Cache {
static CACHE: OnceLock<Cache> = OnceLock::new();
CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
pub fn cached(whir: &WhirParams, n: usize) -> bool {
key_cache()
.lock()
.expect("key cache")
.contains_key(&(whir.header(), n))
}
pub fn key(whir: &WhirParams, n: usize) -> Result<Arc<Key>, String> {
let cache = key_cache();
if let Some(k) = cache.lock().expect("key cache").get(&(whir.header(), n)) {
return Ok(k.clone());
}
let mut p = Params::new(whir, n)?;
let air = CircuitAir::default();
let (_, out) = layout(&p, &air)?;
p.out_row = out;
let (pre, out2) = layout(&p, &air)?;
if out2 != out {
return Err("recursion: the circuit layout moved".into());
}
let (kw, _) = KeyWords::commit(p.cfg.layout, p.n, &pre, p.arity);
let k = Arc::new(Key::assemble(p, pre, kw.root, Some(kw))?);
cache
.lock()
.expect("key cache")
.insert((whir.header(), n), k.clone());
Ok(k)
}
pub fn install(k: Key) -> Arc<Key> {
let cache = key_cache();
let mut c = cache.lock().expect("key cache");
c.entry((k.params.whir.header(), k.params.n))
.or_insert_with(|| Arc::new(k))
.clone()
}
impl Key {
fn assemble(
params: Params,
pre: Pre,
kw_root: Digest,
kw: Option<KeyWords>,
) -> Result<Self, String> {
let sparse = super::vkey::sparse(&pre);
let key_ext = super::vkey::ext(&pre);
let dcfg = whir::Config::derive(
¶ms.whir,
params.vars,
&[1, 2],
params.cfg.acc_claims() + 2,
)?
.with_arity(params.arity);
let g = g_graph(¶ms);
let cell = OnceLock::new();
if let Some(kw) = kw {
let _ = cell.set(kw);
}
Ok(Key {
params,
pre,
sparse,
kw_root,
kw: cell,
key_ext,
dcfg,
g,
})
}
pub fn words(&self) -> &KeyWords {
self.kw.get_or_init(|| {
let (kw, _) = KeyWords::commit(
self.params.cfg.layout,
self.params.n,
&self.pre,
self.params.arity,
);
assert_eq!(
kw.root, self.kw_root,
"recursion: the key's words do not match its root"
);
kw
})
}
pub fn layout_bytes(&self) -> Vec<u8> {
let mut w = crate::envelope::codec::Writer::default();
w.raw(LAYOUT_TAG);
w.raw(&self.params.whir.header());
w.varint(self.params.n as u64);
w.varint(self.params.out_row as u64);
super::vkey::write_digest(&mut w, &self.kw_root);
super::vkey::write_pre(&mut w, &self.pre);
w.bytes
}
pub fn from_layout(bytes: &[u8]) -> Result<Self, String> {
use crate::execution::succinct::SuccinctPcs;
let e = |e: crate::envelope::EnvelopeError| format!("ivc key layout: {e}");
let mut r = crate::envelope::codec::Reader::new(bytes);
if r.raw(LAYOUT_TAG.len()).map_err(e)? != LAYOUT_TAG {
return Err("ivc key layout: tag".into());
}
let whir = lens::Whir::params_from_header(r.raw(8).map_err(e)?)
.map_err(|_| "ivc key layout: header")?;
let n = r.varint().map_err(e)? as usize;
if !(10..=20).contains(&n) {
return Err("ivc key layout: step size".into());
}
let out_row = r.varint().map_err(e)? as usize;
if out_row >= 1 << n {
return Err("ivc key layout: output row".into());
}
let kw_root = super::vkey::read_digest(&mut r)?;
let pre = super::vkey::read_pre(&mut r, 1 << n)?;
r.finish().map_err(e)?;
let mut p = Params::new(&whir, n)?;
p.out_row = out_row;
Self::assemble(p, pre, kw_root, None)
}
}
const LAYOUT_TAG: &[u8; 8] = b"ZHKIVC02";
fn g_graph(p: &Params) -> Graph {
use crate::machine::air::{KCONSTS, KConst};
let constants = Constants {
fml0: 0,
obj0: 0,
p: 0,
output: [Goldilocks::ZERO; 4],
cycles: 0,
root: [Goldilocks::ZERO; 4],
};
let rel = Relation::new(
Machine::new(constants, &[], 1 << p.n, 0, 1 << p.n),
[Fp3::ZERO; 2],
);
record(G_POINT + 2 + KCONSTS, |i| {
let (pt, rest) = i.split_at(G_POINT);
vec![rel.g_with(pt, &[rest[0], rest[1]], &KConst::from_slice(&rest[2..]))]
})
}
pub fn statement_digest(st: &MachineStatement, n: usize, start: u64, segments: u64) -> Digest {
let mut o = Native::new();
let mut sp = Sponge::new(&mut o, tag::CTX + 16);
let bytes = st.bytes();
let base = |v: u64| Fp3::from_base(Goldilocks::new(v));
sp.absorb(&mut o, base(bytes.len() as u64));
for c in bytes.chunks(4) {
let mut b = [0u8; 4];
b[..c.len()].copy_from_slice(c);
sp.absorb(&mut o, base(u64::from(u32::from_le_bytes(b))));
}
for v in [n as u64, start, segments] {
sp.absorb(&mut o, base(v));
}
core::array::from_fn(|_| sp.squeeze(&mut o).c0)
}
pub fn run_challenges(sd: Digest, chain: Digest) -> [Fp3; 2] {
let mut o = Native::new();
let mut sp = Sponge::new(&mut o, tag::CTX + 17);
for v in sd.iter().chain(&chain) {
sp.absorb(&mut o, Fp3::from_base(*v));
}
[sp.squeeze_ext(&mut o), sp.squeeze_ext(&mut o)]
}
fn digest_vals(d: &[Fp3; 4]) -> Digest {
[d[0].c0, d[1].c0, d[2].c0, d[3].c0]
}
fn parts(
p: &Params,
key: &[Vec<(u32, Fp3)>],
rel: &Relation,
global: &[Public],
sd: Digest,
chain: Digest,
) -> CtxParts {
CtxParts {
statement: sd,
chain,
g0: rel.g(&vec![Fp3::ZERO; G_POINT]),
pn0: pbar_nox(global, &vec![Fp3::ZERO; p.dims.pn], p.n),
pv0: pbar_v(key, &vec![Fp3::ZERO; p.dims.pv], p.n),
}
}
pub fn prove(
program: &ExecutionNoun,
input: &[u64],
budget: u64,
whir: &WhirParams,
n: u32,
) -> Result<(MachineStatement, IvcProof), String> {
let run =
machine::execute_exact(program, input, budget, n).map_err(|e| format!("machine: {e:?}"))?;
let proof = prove_run(&run, whir)?;
Ok((run.statement, proof))
}
fn timer() -> impl Fn(&str) {
timer_pub(" ")
}
pub fn timer_pub(indent: &'static str) -> impl Fn(&str) {
let on = std::env::var_os("ZHENG_TIMING").is_some();
let clock = std::time::Instant::now();
let cpu0 = if on { thread_cpu_ms() } else { 0.0 };
let all0 = if on { process_cpu_ms() } else { 0.0 };
move |what: &str| {
if on {
let (a, b) = super::perm::count();
eprintln!(
"{indent}{what}: {:.2} ms (this thread's CPU {:.2} ms, all threads' {:.0} ms) ยท permutations {a} single, {b} batched",
clock.elapsed().as_secs_f64() * 1e3,
thread_cpu_ms() - cpu0,
process_cpu_ms() - all0
);
}
}
}
pub fn thread_cpu_ms() -> f64 {
cpu_clock(if cfg!(target_os = "macos") { 16 } else { 3 })
}
pub fn process_cpu_ms() -> f64 {
cpu_clock(if cfg!(target_os = "macos") { 12 } else { 2 })
}
fn cpu_clock(id: u32) -> f64 {
#[repr(C)]
struct Timespec {
sec: i64,
nsec: i64,
}
unsafe extern "C" {
fn clock_gettime(clock: u32, tp: *mut Timespec) -> i32;
}
let mut t = Timespec { sec: 0, nsec: 0 };
if unsafe { clock_gettime(id, &mut t) } != 0 {
return 0.0;
}
t.sec as f64 * 1e3 + t.nsec as f64 / 1e6
}
fn nox_row0(seg: &crate::air::Trace, n2: &crate::air::Trace) -> Vec<Goldilocks> {
seg.row(0).iter().chain(n2.row(0)).copied().collect()
}
fn keep_words_bytes() -> usize {
std::env::var("ZHENG_KEEP_WORDS_BYTES")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(2 << 30)
}
pub fn prove_run(run: &Run, whir: &WhirParams) -> Result<IvcProof, String> {
let lap = timer();
let n = run.seg_log as usize;
let k = key(whir, n)?;
let p = &k.params;
let air = CircuitAir::default();
let segs = run.segments();
let layout = p.cfg.layout;
let mut chain = [Goldilocks::ZERO; 4];
let mut ood_a = Vec::with_capacity(segs);
let mut kept_a: Vec<Option<Word>> = Vec::with_capacity(segs);
let mut kept_bytes = 0usize;
for i in 0..segs {
let w = Word::commit_base_a(layout, &run.segment(i).column_major(WORD), p.arity);
let mut o = Native::new();
let root = w.root().map(Fp3::from_base);
let z = step::pre_points(&mut o, root, p.fresh);
let ys: Vec<Fp3> = z.iter().map(|&z| w.univariate(z)).collect();
chain = digest_vals(&step::chain_next(
&mut o,
chain.map(Fp3::from_base),
root,
&ys,
));
ood_a.push(ys);
kept_bytes += w.bytes();
kept_a.push((kept_bytes <= keep_words_bytes()).then_some(w));
}
lap("pre-commit");
let sd = statement_digest(&run.statement, n, run.start as u64, segs as u64);
let ch = run_challenges(sd, chain);
let global = Machine::new(run.constants.clone(), &run.init, run.start, 0, segs << n).publics;
let rels: Vec<Relation> = (0..segs)
.map(|i| Relation::new(run.machine(i), ch))
.collect();
let ctx_parts = parts(p, &k.sparse, &rels[0], &global, sd, chain);
let mut carry = Fp3::ZERO;
let mut n2s = Vec::with_capacity(segs);
for (i, rel) in rels.iter().enumerate() {
let (w2, out) = machine::phase2_build(&rel.machine, &run.segment(i), &ch, carry);
carry = out;
n2s.push(w2);
}
let mut acc = AccData::Zero(ZeroWord::new(&layout, p.arity));
let mut prev: Option<(State, StepProof)> = None;
let mut last_state = None;
for i in 0..segs {
let (st_prev, pf_prev) = match &prev {
Some((s, f)) => (s.clone(), f.clone()),
None => (program::dummy_state(p), program::dummy_proof(p)),
};
let mut b = Builder::new(i > 0);
let live = b.live_var();
let (out, h) = program::run(&mut b, p, live, &st_prev, &pf_prev, &ctx_parts, |b, c| {
b.set_output(c)
});
b.finish()?;
let st: State = state::StateV::from_items(
&p.dims,
&out.items()
.iter()
.map(|&(v, _)| b.value(v))
.collect::<Vec<_>>(),
);
let x = digest_vals(&h.map(|v| b.value(v)));
lap(&format!("circuit {i}"));
let (v1, out_row) = trace::generate_trace(&b, &air, n)?;
if out_row != p.out_row {
return Err("recursion: the circuit's layout differs from its key".into());
}
lap(&format!("trace {i}"));
let seg = run.segment(i);
let word_a = match kept_a[i].take() {
Some(w) => w,
None => Word::commit_base_a(layout, &seg.column_major(WORD), p.arity),
};
let j = (i + 1) % segs;
let b_out = nox_row0(&run.segment(j), &n2s[j]);
let input = StepInput {
rel: &rels[i],
global: &global,
n1: &seg,
n2: &n2s[i],
word_a: &word_a,
ood_a: &ood_a[i],
b_out: &b_out,
v1: &v1,
key: &k.pre,
sparse: &k.sparse,
};
let (pf, data, _) = prove::prove(p, &input, &st, x, &acc)?;
let mut o = Native::batched();
let next = step::verify(&mut o, p, &st, x.map(Fp3::from_base), &pf);
o.finish().map_err(|e| format!("step {i}: {e}"))?;
lap(&format!("self-check {i}"));
acc = data;
last_state = Some(next);
prev = Some((st, pf));
lap(&format!("step {i}"));
}
let (state, step) = prev.expect("a segment");
let fin = last_state.expect("a segment");
let AccData::Word(word) = acc else {
return Err("recursion: the last accumulator".into());
};
let decider = decide::prove(&k.dcfg, &word, &fin.acc, &fin.pv, k.words(), n)?;
lap("decide");
Ok(IvcProof {
log_rows: n as u32,
start: run.start as u64,
segments: segs as u64,
chain,
state,
step,
decider,
})
}
pub struct Prepared {
pub key: Arc<Key>,
pub header: (u32, u64, u64, Digest),
pub publics: Publics<Fp3>,
pub global: Vec<Public>,
}
pub fn prepare(
st: &MachineStatement,
whir: &WhirParams,
log_rows: u32,
start: u64,
segments: u64,
chain: Digest,
) -> Result<Prepared, String> {
let n = log_rows as usize;
if !(10..=20).contains(&n) {
return Err("recursion: step size".into());
}
let key = key(whir, n)?;
let derived = st.derive()?;
let (segs, st_row) = (segments as usize, start as usize);
if segs == 0
|| segs > 1 << 32
|| !st_row.is_multiple_of(BLOCK)
|| st_row <= derived.entries.len()
|| st_row + BLOCK > segs << n
{
return Err("recursion: geometry".into());
}
let constants = Constants {
fml0: derived.fml0,
obj0: derived.obj0,
p: derived.entries.len() as u64,
output: derived.output,
cycles: st.cycles,
root: derived.root,
};
let init = statement::init_columns(&derived.entries);
let sd = statement_digest(st, n, start, segments);
let ch = run_challenges(sd, chain);
let global = Machine::new(constants.clone(), &init, st_row, 0, segs << n).publics;
let rel = Relation::new(
Machine::new(constants.clone(), &init, st_row, 0, 1 << n),
ch,
);
let ctx = state::ctx_native(&parts(&key.params, &key.sparse, &rel, &global, sd, chain));
let base = |d: Digest| d.map(Fp3::from_base);
let publics = Publics {
ctx: base(ctx),
chain: base(chain),
segments: Fp3::from_base(Goldilocks::new(segments)),
ch,
k: constants.lift(),
};
Ok(Prepared {
key,
header: (log_rows, start, segments, chain),
publics,
global,
})
}
pub fn verify(st: &MachineStatement, proof: &IvcProof, whir: &WhirParams) -> Result<(), String> {
let prep = prepare(
st,
whir,
proof.log_rows,
proof.start,
proof.segments,
proof.chain,
)?;
verify_prepared(&prep, proof)
}
pub fn check_shapes(key: &Key, proof: &IvcProof) -> Result<(), String> {
let p = &key.params;
step::check_shape(p, &proof.step)?;
decide::check_shape(&key.dcfg, &proof.decider)?;
let s = &proof.state;
let shape_ok = s.b_first.len() == p.dims.boundary
&& s.b_last.len() == p.dims.boundary
&& s.acc.rho.len() == p.dims.vars
&& s.acc.ood.len() == p.dims.ood
&& s.acc.spot.len() == p.dims.spot
&& s.g.point.len() == p.dims.g
&& s.pn.point.len() == p.dims.pn
&& s.pv.point.len() == p.dims.pv
&& state::is_canonical(s);
if shape_ok {
Ok(())
} else {
Err("recursion: state shape".into())
}
}
pub fn verify_prepared(prep: &Prepared, proof: &IvcProof) -> Result<(), String> {
verify_claim(prep, proof).map(|_| ())
}
pub fn verify_claim(prep: &Prepared, proof: &IvcProof) -> Result<state::ClaimV<Fp3>, String> {
let lap = timer();
if prep.header != (proof.log_rows, proof.start, proof.segments, proof.chain) {
return Err("recursion: header".into());
}
let k = &prep.key;
check_shapes(k, proof)?;
lap("setup");
let mut o = Native::batched();
let pubs = finalv::constants(&mut o, &prep.publics);
let pn = finalv::run(&mut o, k, &pubs, &proof.state, &proof.step, &proof.decider);
o.finish()?;
lap("final verifier");
if pbar_nox(&prep.global, &pn.point, k.params.n) != pn.value {
return Err("recursion: deferred nox publics".into());
}
lap("nox publics");
Ok(pn)
}