wai-quantum 0.4.0

//! Signed, re-checkable resource estimates (`wai.quantum.resource`).
//!
//! A resource estimate is the number a claim stands on: so many physical
//! qubits, so much time. The receipt binds that number to everything it came
//! from:
//!
//! - **the job**: the program's counts, the physical qubit, the code (or the
//!   rule "cheapest code this qubit supports"), the error budget, the slowdown
//!   and the search limits. Rates and times are hashed as their exact IEEE-754
//!   bits.
//! - **the estimate**: every figure in it, down to each factory round's copy
//!   count and output error.
//! - **the energy** the sink spent producing it (`joules_micro`, `0` if
//!   unmetered).
//!
//! The estimate is a pure function of the job, computed with `+ − × ÷` and
//! in-crate logarithms. [`ResourceReceipt::verify_computation`] therefore
//! re-runs the job anywhere and demands the identical estimate, bit for bit.
//! Change one input (an error rate by one ulp, the budget, the code) and the
//! check fails.
//!
//! The receipt is a JWP profile like the crate's others (Merkle root over leaf
//! hashes, Ed25519). It can carry an energy-class label
//! (`quantum_energy::Labelled`).

use crate::merkle::merkle_root;
use crate::quantum_resource::{
    estimate_best, estimate_best_isa, estimate, estimate_isa, Budget, Estimate, EstimateConfig, EstimateError,
    Factory, InstructionSet, IsaCounts, LogicalCounts, PhysicalQubit, QecModel, Round, SearchLimits, UnitKind,
};
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
use serde_json::{json, Value};

const DOMAIN_RECEIPT: &[u8] = b"wai:resource-receipt\x01";
const DOMAIN_RECEIPT_ID: &[u8] = b"wai:resource-receipt-id\x01";
const DOMAIN_JOB: &[u8] = b"wai:resource-job\x01";
const DOMAIN_RESULT: &[u8] = b"wai:resource-estimate\x01";

/// The program being estimated, at the level it is known.
#[derive(Clone, Debug, PartialEq)]
pub enum Program {
    /// Algorithm-level counts, laid out by the estimator.
    Logical(LogicalCounts),
    /// Layout-level counts, given directly.
    Isa(IsaCounts),
}

/// Everything that determines an estimate.
#[derive(Clone, Debug, PartialEq)]
pub struct ResourceJob {
    pub program: Program,
    pub qubit: PhysicalQubit,
    /// The code to use, or `None` for the cheapest the qubit supports.
    pub code: Option<QecModel>,
    pub config: EstimateConfig,
}

impl ResourceJob {
    pub fn run(&self) -> Result<Estimate, EstimateError> {
        match (&self.program, &self.code) {
            (Program::Logical(c), Some(q)) => estimate(c, &self.qubit, q, &self.config),
            (Program::Logical(c), None) => estimate_best(c, &self.qubit, &self.config),
            (Program::Isa(i), Some(q)) => estimate_isa(i, &self.qubit, q, &self.config),
            (Program::Isa(i), None) => estimate_best_isa(i, &self.qubit, &self.config),
        }
    }

    /// BLAKE3 over the job's canonical encoding.
    pub fn job_hash(&self) -> [u8; 32] {
        let mut o = DOMAIN_JOB.to_vec();
        match &self.program {
            Program::Logical(c) => {
                o.push(1);
                for x in [c.qubits, c.measurements, c.rotations, c.rotation_depth, c.t_gates, c.ccz] {
                    u(x, &mut o);
                }
            }
            Program::Isa(i) => {
                o.push(2);
                isa(i, &mut o);
            }
        }
        qubit(&self.qubit, &mut o);
        match &self.code {
            Some(q) => {
                o.push(1);
                code(q, &mut o);
            }
            None => o.push(0),
        }
        f(self.config.error_budget, &mut o);
        u(self.config.slowdown, &mut o);
        u(self.config.limits.max_distance as u64, &mut o);
        u(self.config.limits.max_rounds as u64, &mut o);
        *blake3::hash(&o).as_bytes()
    }
}

fn u(x: u64, o: &mut Vec<u8>) {
    o.extend_from_slice(&x.to_be_bytes());
}

fn f(x: f64, o: &mut Vec<u8>) {
    o.extend_from_slice(&x.to_bits().to_be_bytes());
}

fn text(s: &str, o: &mut Vec<u8>) {
    u(s.len() as u64, o);
    o.extend_from_slice(s.as_bytes());
}

fn set(s: InstructionSet) -> u8 {
    match s {
        InstructionSet::GateBased => 1,
        InstructionSet::Measurement => 2,
    }
}

fn isa(i: &IsaCounts, o: &mut Vec<u8>) {
    for x in [i.tiles, i.min_steps, i.t_states, i.t_per_rotation] {
        u(x, o);
    }
    o.push(i.synthesizes as u8);
}

fn qubit(q: &PhysicalQubit, o: &mut Vec<u8>) {
    text(&q.name, o);
    o.push(set(q.instruction_set));
    for x in [q.gate_ns, q.meas_ns, q.clifford_error, q.t_error] {
        f(x, o);
    }
}

fn code(q: &QecModel, o: &mut Vec<u8>) {
    text(&q.name, o);
    o.push(set(q.instruction_set));
    f(q.prefactor, o);
    f(q.threshold, o);
    u(q.qubits.0, o);
    u(q.qubits.1, o);
    u(q.qubits.2 as u64, o);
    f(q.cycle.0, o);
    f(q.cycle.1, o);
}

fn kind_tag(k: UnitKind) -> u8 {
    match k {
        UnitKind::SpaceEfficient => 1,
        UnitKind::ReedMuller => 2,
    }
}

fn result_hash(e: &Estimate) -> [u8; 32] {
    let mut o = DOMAIN_RESULT.to_vec();
    text(&e.qubit, &mut o);
    text(&e.code, &mut o);
    isa(&e.isa, &mut o);
    for x in [e.budget.logical, e.budget.distillation, e.budget.synthesis] {
        f(x, &mut o);
    }
    u(e.distance as u64, &mut o);
    u(e.tile_qubits, &mut o);
    f(e.step_ns, &mut o);
    u(e.steps, &mut o);
    f(e.runtime_ns, &mut o);
    f(e.logical_error, &mut o);
    f(e.t_target, &mut o);
    match &e.factory {
        Some(fa) => {
            o.push(1);
            u(fa.rounds.len() as u64, &mut o);
            for r in &fa.rounds {
                o.push(kind_tag(r.kind));
                u(r.distance as u64, &mut o);
                u(r.copies, &mut o);
                u(r.unit_qubits, &mut o);
                f(r.duration_ns, &mut o);
                f(r.failure, &mut o);
                f(r.output_error, &mut o);
            }
            u(fa.qubits, &mut o);
            f(fa.duration_ns, &mut o);
            f(fa.output_error, &mut o);
            u(fa.outputs, &mut o);
        }
        None => o.push(0),
    }
    for x in [e.factories, e.algorithm_qubits, e.factory_qubits, e.physical_qubits] {
        u(x, &mut o);
    }
    *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)
}

fn bits(x: f64) -> String {
    format!("{:016x}", x.to_bits())
}

fn unbits(v: &Value) -> Option<f64> {
    u64::from_str_radix(v.as_str()?, 16).ok().map(f64::from_bits)
}

/// A signed record of one resource estimate and the energy spent on it.
#[derive(Clone, Debug, PartialEq)]
pub struct ResourceReceipt {
    pub job_hash: [u8; 32],
    pub estimate: Estimate,
    pub result_hash: [u8; 32],
    /// Merkle root over (job, estimate) hashes.
    pub merkle_root: [u8; 32],
    /// Measured microjoules the sink spent (`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 ResourceReceipt {
    fn leaves(&self) -> Vec<[u8; 32]> {
        vec![self.job_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.job_hash);
        o.extend_from_slice(&self.result_hash);
        o.extend_from_slice(&self.merkle_root);
        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).
    pub fn seal(
        signer: &SigningKey,
        signer_id: impl Into<String>,
        job: &ResourceJob,
        parent_receipt_hash: Option<[u8; 32]>,
        meter: impl FnOnce(&mut dyn FnMut()) -> u64,
    ) -> Result<ResourceReceipt, EstimateError> {
        let mut out = None;
        let joules_micro = meter(&mut || out = Some(job.run()));
        let estimate = out.ok_or(EstimateError::NotRun)??;
        let mut r = ResourceReceipt {
            job_hash: job.job_hash(),
            result_hash: result_hash(&estimate),
            estimate,
            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.estimate) == 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: &ResourceJob) -> bool {
        job.job_hash() == self.job_hash && job.run().is_ok_and(|e| result_hash(&e) == self.result_hash)
    }

    /// The strong check: `job` is the job this receipt names, re-running it
    /// here reproduces the sealed estimate bit for bit, and the receipt
    /// verifies.
    pub fn verify_computation(&self, job: &ResourceJob) -> 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, `kind` first and then keys sorted. Every floating-point
    /// figure appears as its exact bits (`*_bits`), which are what a verifier
    /// reads back; the headline runtime also appears as a readable number.
    pub fn to_json(&self) -> String {
        let e = &self.estimate;
        let factory = e.factory.as_ref().map(|fa| {
            json!({
                "duration_ns_bits": bits(fa.duration_ns),
                "output_error_bits": bits(fa.output_error),
                "outputs": fa.outputs,
                "qubits": fa.qubits,
                "rounds": fa.rounds.iter().map(|r| json!({
                    "copies": r.copies,
                    "distance": r.distance,
                    "duration_ns_bits": bits(r.duration_ns),
                    "failure_bits": bits(r.failure),
                    "kind": match r.kind { UnitKind::SpaceEfficient => "space-efficient", UnitKind::ReedMuller => "reed-muller" },
                    "output_error_bits": bits(r.output_error),
                    "unit_qubits": r.unit_qubits,
                })).collect::<Vec<_>>(),
            })
        });
        let v = json!({
            "kind": "quantum-resource",
            "estimate": {
                "algorithm_qubits": e.algorithm_qubits,
                "budget_bits": [bits(e.budget.logical), bits(e.budget.distillation), bits(e.budget.synthesis)],
                "code": e.code,
                "distance": e.distance,
                "factories": e.factories,
                "factory": factory,
                "factory_qubits": e.factory_qubits,
                "isa": {
                    "min_steps": e.isa.min_steps,
                    "synthesizes": e.isa.synthesizes,
                    "t_per_rotation": e.isa.t_per_rotation,
                    "t_states": e.isa.t_states,
                    "tiles": e.isa.tiles,
                },
                "logical_error_bits": bits(e.logical_error),
                "physical_qubits": e.physical_qubits,
                "qubit": e.qubit,
                "runtime_ns_bits": bits(e.runtime_ns),
                "runtime_s": e.runtime_ns / 1e9,
                "step_ns_bits": bits(e.step_ns),
                "steps": e.steps,
                "t_target_bits": bits(e.t_target),
                "tile_qubits": e.tile_qubits,
            },
            "job_hash": hx(&self.job_hash),
            "joules_micro": self.joules_micro,
            "parent_receipt_hash": self.parent_receipt_hash.map(|h| hx(&h)),
            "receipt_hash": hx(&self.receipt_hash()),
            "result_hash": hx(&self.result_hash),
            "root_hash": hx(&self.merkle_root),
            "sig": hx(&self.sig),
            "signer_id": self.signer_id,
            "signer_pubkey": hx(&self.signer_pubkey),
        });
        // The crate's serde_json keeps insertion order, and keys are written in
        // sorted order after `kind`, so this is canonical.
        v.to_string()
    }

    pub fn from_json(s: &str) -> Option<ResourceReceipt> {
        let v: Value = serde_json::from_str(s).ok()?;
        if v.get("kind")?.as_str()? != "quantum-resource" {
            return None;
        }
        let e = v.get("estimate")?;
        let n = |o: &Value, k: &str| o.get(k).and_then(Value::as_u64);
        let b = |o: &Value, k: &str| o.get(k).and_then(unbits);
        let i = e.get("isa")?;
        let budget = e.get("budget_bits")?.as_array()?;
        let factory = match e.get("factory")? {
            Value::Null => None,
            fa => {
                let mut rounds = Vec::new();
                for r in fa.get("rounds")?.as_array()? {
                    rounds.push(Round {
                        kind: match r.get("kind")?.as_str()? {
                            "space-efficient" => UnitKind::SpaceEfficient,
                            "reed-muller" => UnitKind::ReedMuller,
                            _ => return None,
                        },
                        distance: u32::try_from(n(r, "distance")?).ok()?,
                        copies: n(r, "copies")?,
                        unit_qubits: n(r, "unit_qubits")?,
                        duration_ns: b(r, "duration_ns_bits")?,
                        failure: b(r, "failure_bits")?,
                        output_error: b(r, "output_error_bits")?,
                    });
                }
                Some(Factory {
                    rounds,
                    qubits: n(fa, "qubits")?,
                    duration_ns: b(fa, "duration_ns_bits")?,
                    output_error: b(fa, "output_error_bits")?,
                    outputs: n(fa, "outputs")?,
                })
            }
        };
        let estimate = Estimate {
            qubit: e.get("qubit")?.as_str()?.into(),
            code: e.get("code")?.as_str()?.into(),
            isa: IsaCounts {
                tiles: n(i, "tiles")?,
                min_steps: n(i, "min_steps")?,
                t_states: n(i, "t_states")?,
                t_per_rotation: n(i, "t_per_rotation")?,
                synthesizes: i.get("synthesizes")?.as_bool()?,
            },
            budget: Budget {
                logical: unbits(budget.first()?)?,
                distillation: unbits(budget.get(1)?)?,
                synthesis: unbits(budget.get(2)?)?,
            },
            distance: u32::try_from(n(e, "distance")?).ok()?,
            tile_qubits: n(e, "tile_qubits")?,
            step_ns: b(e, "step_ns_bits")?,
            steps: n(e, "steps")?,
            runtime_ns: b(e, "runtime_ns_bits")?,
            logical_error: b(e, "logical_error_bits")?,
            t_target: b(e, "t_target_bits")?,
            factory,
            factories: n(e, "factories")?,
            algorithm_qubits: n(e, "algorithm_qubits")?,
            factory_qubits: n(e, "factory_qubits")?,
            physical_qubits: n(e, "physical_qubits")?,
        };
        let s = |k: &str| v.get(k).and_then(Value::as_str);
        let parent = match v.get("parent_receipt_hash") {
            Some(Value::String(h)) => Some(unhex::<32>(h)?),
            _ => None,
        };
        Some(ResourceReceipt {
            job_hash: unhex(s("job_hash")?)?,
            estimate,
            result_hash: unhex(s("result_hash")?)?,
            merkle_root: unhex(s("root_hash")?)?,
            joules_micro: v.get("joules_micro")?.as_u64()?,
            parent_receipt_hash: parent,
            signer_pubkey: unhex(s("signer_pubkey")?)?,
            signer_id: s("signer_id")?.to_string(),
            sig: unhex(s("sig")?)?,
        })
    }
}

impl Default for ResourceJob {
    /// 2048-bit factoring on the `(ns, 10⁻⁴)` gate-based qubit, `ε = 1/3`.
    fn default() -> Self {
        ResourceJob {
            program: Program::Logical(LogicalCounts {
                qubits: 12_581,
                measurements: 1_080_000_000,
                rotations: 12,
                rotation_depth: 12,
                t_gates: 12,
                ccz: 3_730_000_000,
            }),
            qubit: PhysicalQubit::ns_e4(),
            code: None,
            config: EstimateConfig { error_budget: 1.0 / 3.0, slowdown: 1, limits: SearchLimits::default() },
        }
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

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

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

    #[test]
    fn a_sealed_estimate_verifies_and_reproduces() {
        let j = ResourceJob::default();
        let r = ResourceReceipt::seal(&key(1), "sink-a", &j, None, |run| {
            run();
            1_700
        })
        .unwrap();
        assert!(r.verify());
        assert!(r.verify_computation(&j));
        assert_eq!(r.estimate, j.run().unwrap());
        assert_eq!((r.estimate.distance, r.estimate.factories), (13, 18));
        assert_eq!(r.joules_micro, 1_700);
    }

    #[test]
    fn the_json_round_trips_to_the_same_bytes() {
        // With a factory, and a Clifford-only program without one.
        let clifford = ResourceJob {
            program: Program::Logical(LogicalCounts { qubits: 50, measurements: 1_000, ..Default::default() }),
            ..ResourceJob::default()
        };
        for job in [ResourceJob::default(), clifford] {
            let r = ResourceReceipt::seal(&key(2), "sink \"b\"", &job, Some([7; 32]), unmetered).unwrap();
            let back = ResourceReceipt::from_json(&r.to_json()).unwrap();
            assert_eq!(back, r);
            assert_eq!(back.to_json(), r.to_json());
            assert!(back.verify());
            assert!(back.verify_computation(&job));
        }
    }

    #[test]
    fn tampering_is_caught() {
        let j = ResourceJob::default();
        let r = ResourceReceipt::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.estimate.physical_qubits -= 1;
        assert!(!b.verify(), "the headline is hashed");
        let mut c = r.clone();
        c.estimate.factory.as_mut().unwrap().rounds[0].copies = 15;
        assert!(!c.verify(), "every factory round is hashed");
        // One ulp on the physical error rate is a different job.
        let mut other = j.clone();
        other.qubit.clifford_error = f64::from_bits(other.qubit.clifford_error.to_bits() + 1);
        assert!(!r.verify_computation(&other));
        let mut fixed = j.clone();
        fixed.code = Some(QecModel::surface_gate());
        assert!(!r.verify_computation(&fixed), "choosing the code is part of the job");
        let mut slower = j.clone();
        slower.config.slowdown = 2;
        assert!(!r.verify_computation(&slower));
    }

    #[test]
    fn a_meter_that_never_runs_the_job_cannot_seal() {
        assert_eq!(
            ResourceReceipt::seal(&key(4), "sink", &ResourceJob::default(), None, |_run| 99),
            Err(EstimateError::NotRun)
        );
    }

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