open-standards/wai/wai-quantum/src/quantum_spd_receipt.rs

//! Signed, energy-accounted receipts for sparse Pauli dynamics
//! (`wai.quantum.spd`).
//!
//! A receipt binds three things:
//!
//! - **the job**: the circuit, the observable, the initial basis state and the
//!   truncation settings. Angles and coefficients are hashed as their exact
//!   IEEE-754 bits.
//! - **the result**: the expectation value, the truncation bound, the term
//!   counts and the work done.
//! - **the energy** the sink measured while computing it (`joules_micro`, `0`
//!   if unmetered).
//!
//! The job is run inside [`SpdReceipt::seal`], under a meter the caller
//! supplies, so the sealed result is never a caller's claim. Because the
//! computation is reproducible bit for bit on every platform,
//! [`SpdReceipt::verify_computation`] re-runs the job anywhere and demands the
//! identical result: value, bound, term counts and work. The energy stays what
//! it always is, attested by the signer's key. The receipt is a JWP profile like
//! the crate's others (Merkle root over leaf hashes, Ed25519), and can carry an
//! energy-class label (`quantum_energy::Labelled`).

use crate::merkle::merkle_root;
use crate::quantum_spd::{Op, Pauli, PauliString, RotCircuit, SpdConfig, SpdError, SpdResult, expectation};
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};

const DOMAIN_RECEIPT: &[u8] = b"wai:spd-receipt\x01";
const DOMAIN_RECEIPT_ID: &[u8] = b"wai:spd-receipt-id\x01";
const DOMAIN_CIRCUIT: &[u8] = b"wai:spd-circuit\x01";
const DOMAIN_QUERY: &[u8] = b"wai:spd-query\x01";
const DOMAIN_RESULT: &[u8] = b"wai:spd-result\x01";

/// Everything that determines a result: the circuit, the observable `Σ c·P`,
/// the initial basis state `|ones⟩`, and the truncation settings.
#[derive(Clone, Debug, PartialEq)]
pub struct SpdJob {
    pub circuit: RotCircuit,
    pub observable: Vec<(f64, PauliString)>,
    pub ones: Vec<u32>,
    pub config: SpdConfig,
}

impl SpdJob {
    pub fn run(&self) -> Result<SpdResult, SpdError> {
        expectation(&self.circuit, &self.observable, &self.ones, &self.config)
    }

    /// BLAKE3 over the circuit's canonical encoding.
    pub fn circuit_hash(&self) -> [u8; 32] {
        let mut o = DOMAIN_CIRCUIT.to_vec();
        o.extend_from_slice(&self.circuit.n.to_be_bytes());
        o.extend_from_slice(&(self.circuit.ops.len() as u64).to_be_bytes());
        for op in &self.circuit.ops {
            encode_op(op, &mut o);
        }
        *blake3::hash(&o).as_bytes()
    }

    /// BLAKE3 over the observable, the initial state and the settings.
    pub fn query_hash(&self) -> [u8; 32] {
        let mut o = DOMAIN_QUERY.to_vec();
        o.extend_from_slice(&(self.observable.len() as u64).to_be_bytes());
        for (c, s) in &self.observable {
            o.extend_from_slice(&c.to_bits().to_be_bytes());
            encode_string(s, &mut o);
        }
        o.extend_from_slice(&(self.ones.len() as u64).to_be_bytes());
        for q in &self.ones {
            o.extend_from_slice(&q.to_be_bytes());
        }
        o.extend_from_slice(&self.config.threshold.to_bits().to_be_bytes());
        match self.config.max_weight {
            Some(w) => {
                o.push(1);
                o.extend_from_slice(&w.to_be_bytes());
            }
            None => o.push(0),
        }
        *blake3::hash(&o).as_bytes()
    }
}

fn encode_string(s: &[(u32, Pauli)], o: &mut Vec<u8>) {
    o.extend_from_slice(&(s.len() as u64).to_be_bytes());
    for (q, p) in s {
        o.extend_from_slice(&q.to_be_bytes());
        o.push(match p {
            Pauli::X => b'X',
            Pauli::Y => b'Y',
            Pauli::Z => b'Z',
        });
    }
}

fn encode_op(op: &Op, o: &mut Vec<u8>) {
    let one = |tag: u8, q: u32, o: &mut Vec<u8>| {
        o.push(tag);
        o.extend_from_slice(&q.to_be_bytes());
    };
    let two = |tag: u8, a: u32, b: u32, o: &mut Vec<u8>| {
        o.push(tag);
        o.extend_from_slice(&a.to_be_bytes());
        o.extend_from_slice(&b.to_be_bytes());
    };
    match op {
        Op::H(q) => one(1, *q, o),
        Op::S(q) => one(2, *q, o),
        Op::Sdg(q) => one(3, *q, o),
        Op::SX(q) => one(4, *q, o),
        Op::SXdg(q) => one(5, *q, o),
        Op::X(q) => one(6, *q, o),
        Op::Y(q) => one(7, *q, o),
        Op::Z(q) => one(8, *q, o),
        Op::CX(a, b) => two(9, *a, *b, o),
        Op::CZ(a, b) => two(10, *a, *b, o),
        Op::Swap(a, b) => two(11, *a, *b, o),
        Op::Rot { axis, theta } => {
            o.push(12);
            encode_string(axis, o);
            o.extend_from_slice(&theta.to_bits().to_be_bytes());
        }
        Op::QuarterRot { axis, k } => {
            o.push(13);
            encode_string(axis, o);
            o.push(*k);
        }
    }
}

fn result_hash(r: &SpdResult) -> [u8; 32] {
    let mut o = DOMAIN_RESULT.to_vec();
    o.extend_from_slice(&r.value.to_bits().to_be_bytes());
    o.extend_from_slice(&r.truncation_bound.to_bits().to_be_bytes());
    o.extend_from_slice(&(r.peak_terms as u64).to_be_bytes());
    o.extend_from_slice(&(r.final_terms as u64).to_be_bytes());
    o.extend_from_slice(&r.term_updates.to_be_bytes());
    *blake3::hash(&o).as_bytes()
}

fn hx(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

fn unhex<const N: usize>(s: &str) -> Option<[u8; N]> {
    if s.len() != 2 * N || !s.bytes().all(|b| b.is_ascii_hexdigit()) {
        return None;
    }
    let mut out = [0u8; N];
    for (i, c) in s.as_bytes().chunks(2).enumerate() {
        out[i] = u8::from_str_radix(std::str::from_utf8(c).ok()?, 16).ok()?;
    }
    Some(out)
}

/// A signed record of one sparse-Pauli-dynamics computation and its energy.
#[derive(Clone, Debug, PartialEq)]
pub struct SpdReceipt {
    pub circuit_hash: [u8; 32],
    pub query_hash: [u8; 32],
    pub result: SpdResult,
    pub result_hash: [u8; 32],
    /// Merkle root over (circuit, query, result) hashes.
    pub merkle_root: [u8; 32],
    /// Measured microjoules the sink spent on the run (`0` = unmetered).
    /// Attested by the signature, not verifiable by anyone else.
    pub joules_micro: u64,
    pub parent_receipt_hash: Option<[u8; 32]>,
    pub signer_pubkey: [u8; 32],
    pub signer_id: String,
    pub sig: [u8; 64],
}

impl SpdReceipt {
    fn leaves(&self) -> Vec<[u8; 32]> {
        vec![self.circuit_hash, self.query_hash, self.result_hash]
    }

    fn payload_with(&self, label: Option<&crate::quantum_energy::EnergyClass>) -> Vec<u8> {
        let mut o = crate::quantum_energy::labelled_domain(DOMAIN_RECEIPT, label);
        o.extend_from_slice(&self.circuit_hash);
        o.extend_from_slice(&self.query_hash);
        o.extend_from_slice(&self.result_hash);
        o.extend_from_slice(&self.merkle_root);
        o.extend_from_slice(&self.result.term_updates.to_be_bytes());
        o.extend_from_slice(&self.joules_micro.to_be_bytes());
        crate::quantum_energy::write_label(label, &mut o);
        match self.parent_receipt_hash {
            Some(h) => {
                o.push(1);
                o.extend_from_slice(&h);
            }
            None => o.push(0),
        }
        o.extend_from_slice(&self.signer_pubkey);
        o.extend_from_slice(self.signer_id.as_bytes());
        o
    }

    /// Run `job` under `meter` and seal what it computed. `meter` is handed a
    /// closure that performs the run; it must call it once, and returns the
    /// microjoules it measured (`|run| { run(); 0 }` for an unmetered seal).
    /// The result sealed is the one that closure produced.
    pub fn seal(
        signer: &SigningKey,
        signer_id: impl Into<String>,
        job: &SpdJob,
        parent_receipt_hash: Option<[u8; 32]>,
        meter: impl FnOnce(&mut dyn FnMut()) -> u64,
    ) -> Result<SpdReceipt, SpdError> {
        let mut out = None;
        let joules_micro = meter(&mut || out = Some(job.run()));
        let result = out.ok_or(SpdError::NotRun)??;
        let mut r = SpdReceipt {
            circuit_hash: job.circuit_hash(),
            query_hash: job.query_hash(),
            result,
            result_hash: result_hash(&result),
            merkle_root: [0; 32],
            joules_micro,
            parent_receipt_hash,
            signer_pubkey: signer.verifying_key().to_bytes(),
            signer_id: signer_id.into(),
            sig: [0; 64],
        };
        r.merkle_root = merkle_root(&r.leaves());
        r.sig = signer.sign(&r.payload_with(None)).to_bytes();
        Ok(r)
    }

    fn body_verifies(&self) -> bool {
        result_hash(&self.result) == self.result_hash && merkle_root(&self.leaves()) == self.merkle_root
    }

    /// The receipt is internally consistent and its signature is valid. Does
    /// not re-run the job: see [`Self::verify_computation`].
    pub fn verify(&self) -> bool {
        if !self.body_verifies() {
            return false;
        }
        let Ok(k) = VerifyingKey::from_bytes(&self.signer_pubkey) else {
            return false;
        };
        k.verify(&self.payload_with(None), &Signature::from_bytes(&self.sig)).is_ok()
    }

    /// The re-run half of [`Self::verify_computation`].
    pub(crate) fn computation_matches(&self, job: &SpdJob) -> bool {
        job.circuit_hash() == self.circuit_hash
            && job.query_hash() == self.query_hash
            && job.run().is_ok_and(|r| result_hash(&r) == self.result_hash)
    }

    /// The strong check: `job` is the job this receipt names, re-running it
    /// here reproduces the sealed result bit for bit, and the receipt verifies.
    pub fn verify_computation(&self, job: &SpdJob) -> bool {
        self.computation_matches(job) && self.verify()
    }

    /// Stable id, for chaining a derived receipt via `parent_receipt_hash`.
    pub fn receipt_hash(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(DOMAIN_RECEIPT_ID);
        h.update(&self.payload_with(None));
        h.update(&self.sig);
        *h.finalize().as_bytes()
    }

    /// Canonical JSON, keys sorted. Floating-point figures appear twice: as a
    /// readable number, and as their exact bits (`*_bits`), which are what a
    /// verifier reads back.
    pub fn to_json(&self) -> String {
        let parent = match self.parent_receipt_hash {
            Some(h) => format!("\"{}\"", hx(&h)),
            None => "null".into(),
        };
        let r = &self.result;
        format!(
            "{{\"kind\":\"quantum-spd\",\"circuit_hash\":\"{}\",\"final_terms\":{},\"joules_micro\":{},\
             \"parent_receipt_hash\":{},\"peak_terms\":{},\"query_hash\":\"{}\",\"receipt_hash\":\"{}\",\
             \"result_hash\":\"{}\",\"root_hash\":\"{}\",\"sig\":\"{}\",\"signer_id\":{},\
             \"signer_pubkey\":\"{}\",\"term_updates\":{},\"truncation_bound\":{},\
             \"truncation_bound_bits\":\"{:016x}\",\"value\":{},\"value_bits\":\"{:016x}\"}}",
            hx(&self.circuit_hash),
            r.final_terms,
            self.joules_micro,
            parent,
            r.peak_terms,
            hx(&self.query_hash),
            hx(&self.receipt_hash()),
            hx(&self.result_hash),
            hx(&self.merkle_root),
            hx(&self.sig),
            serde_json::to_string(&self.signer_id).unwrap(),
            hx(&self.signer_pubkey),
            r.term_updates,
            serde_json::to_string(&r.truncation_bound).unwrap(),
            r.truncation_bound.to_bits(),
            serde_json::to_string(&r.value).unwrap(),
            r.value.to_bits(),
        )
    }

    pub fn from_json(s: &str) -> Option<SpdReceipt> {
        let v: serde_json::Value = serde_json::from_str(s).ok()?;
        let o = v.as_object()?;
        if o.get("kind")?.as_str()? != "quantum-spd" {
            return None;
        }
        let u = |k: &str| o.get(k).and_then(|x| x.as_u64());
        let s = |k: &str| o.get(k).and_then(|x| x.as_str());
        let bits = |k: &str| u64::from_str_radix(s(k)?, 16).ok().map(f64::from_bits);
        let parent = match o.get("parent_receipt_hash") {
            Some(serde_json::Value::String(h)) => Some(unhex::<32>(h)?),
            _ => None,
        };
        Some(SpdReceipt {
            circuit_hash: unhex(s("circuit_hash")?)?,
            query_hash: unhex(s("query_hash")?)?,
            result: SpdResult {
                value: bits("value_bits")?,
                truncation_bound: bits("truncation_bound_bits")?,
                peak_terms: usize::try_from(u("peak_terms")?).ok()?,
                final_terms: usize::try_from(u("final_terms")?).ok()?,
                term_updates: u("term_updates")?,
            },
            result_hash: unhex(s("result_hash")?)?,
            merkle_root: unhex(s("root_hash")?)?,
            joules_micro: u("joules_micro")?,
            parent_receipt_hash: parent,
            signer_pubkey: unhex(s("signer_pubkey")?)?,
            signer_id: s("signer_id")?.to_string(),
            sig: unhex(s("sig")?)?,
        })
    }
}

crate::quantum_energy::labellable!(SpdReceipt, DOMAIN_RECEIPT_ID);

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantum_spd::{heavy_hex_127, kicked_ising};

    fn key(b: u8) -> SigningKey {
        SigningKey::from_bytes(&[b; 32])
    }

    fn job() -> SpdJob {
        SpdJob {
            circuit: kicked_ising(127, &heavy_hex_127(), 0.4, 4),
            observable: vec![(1.0, vec![(62, Pauli::Z)])],
            ones: vec![],
            config: SpdConfig { threshold: 1e-4, max_weight: None },
        }
    }

    fn unmetered(run: &mut dyn FnMut()) -> u64 {
        run();
        0
    }

    #[test]
    fn a_sealed_receipt_verifies_and_reproduces() {
        let j = job();
        let r = SpdReceipt::seal(&key(1), "sink-a", &j, None, |run| {
            run();
            4_200
        })
        .unwrap();
        assert!(r.verify());
        assert!(r.verify_computation(&j));
        assert_eq!(r.joules_micro, 4_200);
        assert_eq!(r.result, j.run().unwrap());
    }

    #[test]
    fn the_json_round_trips_to_the_same_bytes() {
        let r = SpdReceipt::seal(&key(2), "sink \"b\"", &job(), Some([7; 32]), unmetered).unwrap();
        let back = SpdReceipt::from_json(&r.to_json()).unwrap();
        assert_eq!(back, r);
        assert_eq!(back.to_json(), r.to_json());
        assert!(back.verify());
    }

    #[test]
    fn tampering_is_caught() {
        let j = job();
        let r = SpdReceipt::seal(&key(3), "sink", &j, None, unmetered).unwrap();
        let mut a = r.clone();
        a.joules_micro = 1;
        assert!(!a.verify(), "the energy is signed");
        let mut b = r.clone();
        b.result.value += 1e-12;
        assert!(!b.verify(), "the result is hashed");
        let mut c = r.clone();
        c.result.term_updates -= 1;
        assert!(!c.verify(), "the work is hashed");
        // A different job — one angle moved by one ulp — is not the job sealed.
        let mut other = j.clone();
        if let Op::Rot { theta, .. } = &mut other.circuit.ops[5] {
            *theta = f64::from_bits(theta.to_bits() + 1);
        }
        assert!(!r.verify_computation(&other));
        let mut looser = j.clone();
        looser.config.threshold = 2e-4;
        assert!(!r.verify_computation(&looser), "the settings are part of the job");
    }

    #[test]
    fn a_meter_that_never_runs_the_job_cannot_seal() {
        assert_eq!(SpdReceipt::seal(&key(4), "sink", &job(), None, |_run| 99), Err(SpdError::NotRun));
    }

    #[test]
    fn an_energy_class_label_seals_over_it() {
        use crate::quantum_energy::{EnergyClass, Labelled};
        let r = SpdReceipt::seal(&key(5), "sink", &job(), None, |run| {
            run();
            1_000
        })
        .unwrap();
        let class = EnergyClass::ModelBased;
        let labelled = Labelled::seal(r, class, &key(5)).unwrap();
        assert!(labelled.verify());
        assert!(!labelled.receipt.verify(), "the inner signature now covers the label");
    }
}