soft3/zheng/rs/tests/common/mod.rs

//! Fixtures shared by the succinct-profile tests and the bake-off example:
//! joy's compiled `hash.tri` and `add.tri`, a chain of hemera hashes built
//! for the relation compiler, and synthetic CCS relations of a given size.
#![allow(dead_code)]

use nebu::Goldilocks as F;
use zheng::execution::ExecutionNoun;
use zheng::types::{CCSInstance, SparseMatrix};

/// joy 0.5 `hash.tri`, as compiled.
pub const HASH: &str = "[2 [[3 [[4 [[9 [[0 3] [1 0]]] [[1 0] [8 [1 0]]]]] [0 1]]] [1 [2 [[0 3] [1 [2 [[3 [[5 [[0 2] [1 0]]] [1 0]]] [1 [15 [3 [[0 2] [3 [[1 0] [3 [[1 0] [3 [[1 0] [3 [[1 0] [3 [[1 0] [3 [[1 0] [3 [[1 0] [1 0]]]]]]]]]]]]]]]]]]]]]]]]]]]";
/// joy 0.5 `add.tri`, as compiled.
pub const ADD: &str = "[2 [[3 [[4 [[9 [[0 7] [1 0]]] [[1 0] [8 [1 0]]]]] [0 1]]] [1 [2 [[0 3] [1 [2 [[3 [[5 [[0 2] [1 0]]] [3 [[5 [[0 6] [1 0]]] [1 0]]]]] [1 [2 [[3 [[5 [[0 2] [1 0]]] [0 1]]] [1 [7 [[0 2] [0 14]]]]]]]]]]]]]]]";

pub const P: u64 = 0xFFFF_FFFF_0000_0001;

pub fn parse(text: &str) -> ExecutionNoun {
    fn go(b: &[u8], p: &mut usize) -> ExecutionNoun {
        while b[*p].is_ascii_whitespace() {
            *p += 1;
        }
        if b[*p] == b'[' {
            *p += 1;
            let a = go(b, p);
            let c = go(b, p);
            while b[*p].is_ascii_whitespace() {
                *p += 1;
            }
            *p += 1;
            ExecutionNoun::Pair(Box::new(a), Box::new(c))
        } else {
            let s = *p;
            while b[*p].is_ascii_digit() {
                *p += 1;
            }
            ExecutionNoun::Atom(std::str::from_utf8(&b[s..*p]).unwrap().parse().unwrap())
        }
    }
    go(text.as_bytes(), &mut 0)
}

fn pair(a: ExecutionNoun, b: ExecutionNoun) -> ExecutionNoun {
    ExecutionNoun::Pair(Box::new(a), Box::new(b))
}
fn op(tag: u64, arg: ExecutionNoun) -> ExecutionNoun {
    pair(ExecutionNoun::Atom(tag), arg)
}
fn axis(a: u64) -> ExecutionNoun {
    op(0, ExecutionNoun::Atom(a))
}
fn quote(v: u64) -> ExecutionNoun {
    op(1, ExecutionNoun::Atom(v))
}
/// `[3 [a [3 [b … [3 [g h]]]]]]`: the 8-element hash input.
fn cons8(items: Vec<ExecutionNoun>) -> ExecutionNoun {
    let mut it = items.into_iter().rev();
    let mut acc = it.next().unwrap();
    for x in it {
        acc = op(3, pair(x, acc));
    }
    acc
}
/// `[2 [x [1 y]]]`: run `x` on the subject, then `y` on its result.
fn compose(x: ExecutionNoun, y: ExecutionNoun) -> ExecutionNoun {
    op(2, pair(x, op(1, y)))
}

/// Axis of input `k` in the subject `[i0 [i1 … [i_{n−1} 0]]]`.
fn input_axis(k: usize) -> ExecutionNoun {
    axis((1u64 << (k + 2)) - 2)
}

/// One hemera hash of `n` public inputs, `hash([i0 … i_{n−1}])` — the shape
/// of a device-side commitment (a mining ticket hashes beacon, cluster,
/// nonce, miner and commitment; a tip hashes height and two digests).
pub fn hash_inputs(n: usize) -> ExecutionNoun {
    op(15, cons8((0..n).map(input_axis).collect()))
}

/// `groups` independent hashes of `per` inputs each, consed: a settlement
/// batch of `groups` tickets. Inputs `groups · per` (≤ 64).
pub fn hash_batch(groups: usize, per: usize) -> ExecutionNoun {
    cons8(
        (0..groups)
            .map(|g| {
                op(
                    15,
                    cons8((0..per).map(|k| input_axis(g * per + k)).collect()),
                )
            })
            .collect(),
    )
}

/// `n` chained hemera hashes: `d_1 = H(a, 0⁷)`, `d_{i+1} = H(d_i, 0⁴)`, over
/// the subject `[a 0]`; a digest is the noun `[[d0 d1] [d2 d3]]` (axes 4–7). Composition is balanced so the formula depth stays
/// `O(log n)` within the statement's depth limit. Each hash of a digest
/// costs ~2,670 wires; the compiler's cap of 32,768 ops/rows admits
/// `n ≤ 11` (27,520 free wires).
pub fn hash_chain(n: usize) -> ExecutionNoun {
    assert!(n >= 1);
    let first = op(
        15,
        cons8(
            std::iter::once(axis(2))
                .chain((0..7).map(|_| quote(0)))
                .collect(),
        ),
    );
    let step = || {
        op(
            15,
            cons8(
                [4, 5, 6, 7]
                    .into_iter()
                    .map(axis)
                    .chain((0..4).map(|_| quote(0)))
                    .collect(),
            ),
        )
    };
    fn tree(lo: usize, hi: usize, leaf: &dyn Fn(usize) -> ExecutionNoun) -> ExecutionNoun {
        if hi - lo == 1 {
            return leaf(lo);
        }
        let mid = lo + (hi - lo) / 2;
        compose(tree(lo, mid, leaf), tree(mid, hi, leaf))
    }
    let leaf = |i: usize| if i == 0 { first.clone() } else { step() };
    tree(0, n, &leaf)
}

/// A synthetic relation with `2^k` rows and `2^k` columns in the shape the
/// compiler emits (five matrices: `a·b = c` and `x^7 = w`):
/// `z = (1, x, w_2, …, w_{2^k−1})`, row `r` defines `w_{r+2}` as
/// `(w_{r+1} + (r+1))^7` for even `r` and `w_{r+1} · w_r` for odd `r`.
/// Pinned: `z[0] = 1`, the input `z[1]` and the output `z[2^k − 1]`.
pub struct Synthetic {
    pub instance: CCSInstance,
    pub z: Vec<F>,
    pub pins: Vec<(usize, F)>,
    pub statement: Vec<u8>,
}

pub fn synthetic(k: u32, x: u64) -> Synthetic {
    let n = 1usize << k;
    let mut m = vec![SparseMatrix::new(n, n); 5];
    let mut z = vec![F::ZERO; n];
    z[0] = F::ONE;
    z[1] = F::new(x);
    for r in 0..n - 2 {
        if r % 2 == 0 {
            let c = F::new(r as u64 + 1);
            m[3].set(r, r + 1, F::ONE);
            m[3].set(r, 0, c);
            m[4].set(r, r + 2, F::ONE);
            z[r + 2] = (z[r + 1] + c).exp(7);
        } else {
            m[0].set(r, r + 1, F::ONE);
            m[1].set(r, r, F::ONE);
            m[2].set(r, r + 2, F::ONE);
            z[r + 2] = z[r + 1] * z[r];
        }
    }
    let instance = CCSInstance {
        matrices: m,
        multisets: vec![vec![0, 1], vec![2], vec![3; 7], vec![4]],
        coeffs: vec![F::ONE, -F::ONE, F::ONE, -F::ONE],
        num_rows: n,
        num_cols: n,
    };
    let pins = vec![(0, F::ONE), (1, z[1]), (n - 1, z[n - 1])];
    let mut statement = b"zheng-synthetic-v1".to_vec();
    statement.extend((k as u64).to_le_bytes());
    for &(_, v) in &pins[1..] {
        statement.extend(v.as_u64().to_le_bytes());
    }
    Synthetic {
        instance,
        z,
        pins,
        statement,
    }
}

/// Run `program` on the subject `[input… 0]` with native nox; returns the
/// output atoms (left to right) and the reductions spent.
pub fn native(program: &ExecutionNoun, input: &[u64], budget: u64) -> (Vec<u64>, u64) {
    fn arena(r: &mut nox::Reduction<8192>, n: &ExecutionNoun) -> nox::Order {
        match n {
            ExecutionNoun::Atom(v) => r.atom(F::new(*v)).unwrap(),
            ExecutionNoun::Pair(a, b) => {
                let a = arena(r, a);
                let b = arena(r, b);
                r.pair(a, b).unwrap()
            }
        }
    }
    fn leaves(r: &nox::Reduction<8192>, n: nox::Order, out: &mut Vec<u64>) {
        match r.atom_value(n) {
            Some(v) => out.push(v.as_u64()),
            None => {
                leaves(r, r.head(n).unwrap(), out);
                leaves(r, r.tail(n).unwrap(), out);
            }
        }
    }
    let mut subject = ExecutionNoun::Atom(0);
    for &v in input.iter().rev() {
        subject = pair(ExecutionNoun::Atom(v), subject);
    }
    let mut r = Box::new(nox::Reduction::<8192>::new());
    let s = arena(&mut r, &subject);
    let f = arena(&mut r, program);
    match nox::reduce(
        &mut r,
        s,
        f,
        budget,
        &nox::call::NullCalls,
        &mut nox::trace::NoTrace,
    ) {
        nox::Outcome::Ok(out, remaining) => {
            let mut atoms = Vec::new();
            leaves(&r, out, &mut atoms);
            (atoms, budget - remaining)
        }
        other => panic!("native nox: {other:?}"),
    }
}

// ── machine fixtures ────────────────────────────────────────────────

fn m_pr(a: ExecutionNoun, b: ExecutionNoun) -> ExecutionNoun {
    ExecutionNoun::Pair(Box::new(a), Box::new(b))
}
fn m_at(v: u64) -> ExecutionNoun {
    ExecutionNoun::Atom(v)
}
fn m_axis(n: u64) -> ExecutionNoun {
    m_pr(m_at(0), m_at(n))
}
fn m_quote(n: ExecutionNoun) -> ExecutionNoun {
    m_pr(m_at(1), n)
}
fn m_op(t: u64, a: ExecutionNoun, b: ExecutionNoun) -> ExecutionNoun {
    m_pr(m_at(t), m_pr(a, b))
}

/// One Merkle level on the state `[cur [[sib dir] rest]]`:
/// `[hash(dir = 0 ? [cur sib] : [sib cur]) rest]`.
fn merkle_level() -> ExecutionNoun {
    let left = m_op(3, m_axis(2), m_axis(12));
    let right = m_op(3, m_axis(12), m_axis(2));
    let pick = m_pr(m_at(4), m_pr(m_axis(13), m_pr(left, right)));
    m_op(3, m_pr(m_at(15), pick), m_axis(7))
}

/// Apply `f` `k` times by nested compose.
fn repeat(f: &ExecutionNoun, k: usize) -> ExecutionNoun {
    (1..k).fold(f.clone(), |acc, _| m_op(2, acc, m_quote(f.clone())))
}

/// The digest noun `[[a b] [c d]]`.
pub fn digest_noun(d: [u64; 4]) -> ExecutionNoun {
    m_pr(m_pr(m_at(d[0]), m_at(d[1])), m_pr(m_at(d[2]), m_at(d[3])))
}

/// A `levels`-hash Merkle path check on input `[x 0]`: the leaf is
/// `hash(x)`, the siblings and directions are quoted in the program, the
/// output is `eq(computed root, root)` (0 when the path verifies). With
/// `root = None` the program returns the computed root instead.
pub fn merkle_program(levels: usize, root: Option<&ExecutionNoun>) -> ExecutionNoun {
    let mut path = m_at(0);
    for i in (0..levels).rev() {
        let sib = digest_noun([1000 + i as u64, 7 * i as u64, 3, i as u64 * i as u64]);
        path = m_pr(m_pr(sib, m_at((i % 2) as u64)), path);
    }
    let start = m_op(3, m_pr(m_at(15), m_axis(2)), m_quote(path));
    let folded = m_op(2, start, m_quote(repeat(&merkle_level(), levels)));
    let computed = m_op(2, folded, m_quote(m_axis(2)));
    match root {
        None => computed,
        Some(r) => m_op(9, computed, m_quote(r.clone())),
    }
}

/// A `levels`-level hemera Merkle path (`hemera::merkle`: `hash_node`,
/// the root flag on the last level) on input `[x 0]`, each level one call
/// of the node jet (`zheng::execution::node`). The leaf is `hash(x)`
/// (pattern 15); the siblings and directions of [`merkle_program`] are
/// quoted in the program (direction 0: the sibling is the left child).
/// The output is `eq(computed root, root)`, 0 when the path verifies; with
/// `root = None` the computed root itself.
///
/// The subject while folding is `[cur [path [NODE [LEVEL ROOT]]]]`: the
/// node formula and both level formulas ride in it once, and every level
/// is the 9-token step `[2 [[0 1] [0 30￿31]]]`, so a depth-32 path stays
/// under the 4,096-node program cap.
pub fn node_merkle_program(levels: usize, root: Option<&ExecutionNoun>) -> ExecutionNoun {
    assert!(levels >= 1);
    use zheng::execution::node;
    let mut path = m_at(0);
    for i in (0..levels).rev() {
        let sib = digest_noun([1000 + i as u64, 7 * i as u64, 3, i as u64 * i as u64]);
        path = m_pr(m_pr(sib, m_at((i % 2) as u64)), path);
    }
    // subject: cur 2, path 6, sibling 24, direction 25, rest 13, [NODE …] 7, NODE 14
    let level = |flags: u64| {
        let children = m_pr(
            m_at(4),
            m_pr(
                m_axis(25),
                m_pr(
                    m_op(3, m_axis(24), m_axis(2)),
                    m_op(3, m_axis(2), m_axis(24)),
                ),
            ),
        );
        let next = m_op(2, m_op(3, children, m_quote(m_at(flags))), m_axis(14));
        m_op(3, next, m_op(3, m_axis(13), m_axis(7)))
    };
    let lib = m_pr(
        node::formula().clone(),
        m_pr(level(node::FLAGS_NODE), level(node::FLAGS_ROOT)),
    );
    let start = m_op(3, m_pr(m_at(15), m_axis(2)), m_quote(m_pr(path, lib)));
    let step = |axis: u64| m_op(2, m_axis(1), m_axis(axis));
    let mut steps = step(31);
    if levels > 1 {
        steps = m_op(2, repeat(&step(30), levels - 1), m_quote(steps));
    }
    let folded = m_op(2, start, m_quote(steps));
    let computed = m_op(2, folded, m_quote(m_axis(2)));
    match root {
        None => computed,
        Some(r) => m_op(9, computed, m_quote(r.clone())),
    }
}

/// The root [`node_merkle_program`] computes, natively through hemera.
pub fn node_merkle_root(levels: usize, x: u64) -> [u64; 4] {
    use zheng::execution::node;
    let leaf = nox::data::hash::hash_atom(F::new(x));
    let rate: [hemera::field::Goldilocks; 8] = core::array::from_fn(|i| {
        hemera::field::Goldilocks::new(if i < 4 { leaf[i].as_u64() } else { 0 })
    });
    let state = hemera::StepSponge::absorb(&rate).last().unwrap();
    let mut cur: [u64; 4] = core::array::from_fn(|i| state[i].as_canonical_u64());
    for i in 0..levels {
        let sib = [1000 + i as u64, 7 * i as u64, 3, i as u64 * i as u64];
        let flags = if i + 1 == levels {
            node::FLAGS_ROOT
        } else {
            node::FLAGS_NODE
        };
        cur = if i % 2 == 0 {
            node::native(sib, cur, flags)
        } else {
            node::native(cur, sib, flags)
        };
    }
    cur
}

/// A balanced tree of `2^k` additions on input `[x 0]`: `f_0 = x + 1`,
/// `f_k = f_{k−1} + f_{k−1}` — nox depth `k + 1`, about `3·2^k` cycles.
pub fn tree_program(k: usize) -> ExecutionNoun {
    (0..k).fold(m_op(5, m_axis(2), m_quote(m_at(1))), |f, _| {
        m_op(5, f.clone(), f)
    })
}

/// Doubling recursion over a subject that carries its own formula:
/// `f(n) = n = 0 ? 1 : f(n−1) + f(n−1)` on input `[n 0]`, so `f(n) = 2^n`
/// after about `2^{n+1}` calls. A small program (tens of tokens) whose
/// run grows exponentially, with nox depth about `4n`.
pub fn rec_program() -> ExecutionNoun {
    // inside F the subject is [n F]
    let n = m_axis(2);
    let me = m_axis(3);
    let call = m_op(
        2,
        m_op(3, m_op(6, n.clone(), m_quote(m_at(1))), me.clone()),
        me,
    );
    let body = m_op(5, call.clone(), call);
    let test = m_op(9, n, m_quote(m_at(0)));
    let f = m_pr(m_at(4), m_pr(test, m_pr(m_quote(m_at(1)), body)));
    // start: subject [x 0] → [x F], formula F
    m_op(2, m_op(3, m_axis(2), m_quote(f.clone())), m_quote(f))
}

/// Install the prover backend named by `ZHENG_BACKEND` (`cpu`, `metal`,
/// `webgpu`, `auto`; unset = cpu) and report it on stderr. Proof bytes do
/// not depend on it.
pub fn backend() -> &'static str {
    let choice = std::env::var("ZHENG_BACKEND").unwrap_or_else(|_| "cpu".into());
    let (name, device) = choice
        .parse::<cyber_lens_gpu::Choice>()
        .and_then(cyber_lens_gpu::install)
        .unwrap_or_else(|e| panic!("ZHENG_BACKEND: {e}"));
    eprintln!("backend: {name} ({device})");
    name
}

/// With a `prof-*` backend: the device operations' wall times since the
/// last call, labelled.
#[allow(dead_code)]
pub fn prof_report(label: &str) {
    if std::env::var("ZHENG_BACKEND").is_ok_and(|b| b.starts_with("prof-")) {
        eprint!(
            "device ops, {label}:\n{}",
            cyber_lens_gpu::profile::report()
        );
    }
}