open-standards/wai/wai-quantum/tests/qasm3_conformance.rs

//! The OpenQASM 3 conformance corpus (`wai/quantum-conformance/qasm3/`):
//! every program read, run and checked against what its manifest says the
//! reader must do with it.
//!
//! The cases come from `bundle.json`, which holds each program's text inline,
//! so the same test runs where no file can be read (under wasm32-wasip2, by
//! wasmtime); natively, a second test holds the bundle equal to the files.
//! Run with `WAI_QASM3_BUNDLE_WRITE=1` to write the bundle from the files.

#![cfg(feature = "quantum_engine")]

use serde_json::{Value, json};
use wai_quantum::engine::job;
use wai_quantum::qasm::{self, Inputs, Program};
use wai_quantum::quantum_sv::StateVector;

const BUNDLE: &str = include_str!("../../quantum-conformance/qasm3/bundle.json");

fn inputs_of(v: &Value) -> Inputs {
    v.as_object().map(|o| o.iter().map(|(k, v)| (k.clone(), v.as_i64().expect("an integer input"))).collect()).unwrap_or_default()
}

/// The program's amplitudes, from |0…0⟩.
fn amplitudes(p: &Program, col: usize) -> Result<(Vec<f64>, Vec<f64>), String> {
    let mut sv = StateVector::zero(p.n_qubits).map_err(|e| format!("{e:?}"))?;
    sv.re[0] = 0.0;
    sv.re[col] = 1.0;
    sv.run(&p.matrices(), 1).map_err(|e| format!("{e:?}"))?;
    Ok((sv.re.clone(), sv.im.clone()))
}

fn run(qasm: &str, extra: Value) -> Result<Value, String> {
    let mut req = json!({ "qasm": qasm, "seed": 7 });
    if let (Some(o), Some(e)) = (req.as_object_mut(), extra.as_object()) {
        o.extend(e.clone());
    }
    job::plan_and_run(&req, 1 << 30).map(|(_, result)| result)
}

fn check(case: &Value) -> Result<(), String> {
    let src = case["qasm"].as_str().ok_or("no program")?;
    let e = &case["expect"];
    let ins = inputs_of(&e["inputs"]);
    let reading = qasm::read_with(src, &ins);
    if let Some(code) = e["refused"].as_str() {
        return match reading {
            Err(d) if d.code == code => Ok(()),
            Err(d) => Err(format!("refused with {}, not {code}: {d}", d.code)),
            Ok(_) => Err(format!("read, where it must be refused with {code}")),
        };
    }
    let r = reading.map_err(|d| format!("does not read: {d}"))?;
    let mut got: Vec<&str> = r.diagnostics.iter().map(|d| d.code).collect();
    got.sort_unstable();
    got.dedup();
    let want: Vec<&str> = e["warns"].as_array().map(|a| a.iter().filter_map(Value::as_str).collect()).unwrap_or_default();
    if got != want {
        return Err(format!("warns {got:?}, not {want:?}: {:?}", r.diagnostics.iter().map(|d| d.to_string()).collect::<Vec<_>>()));
    }
    // What the engine runs: a generator written out at its inputs.
    let text = if ins.is_empty() { src.to_string() } else { qasm::instantiate(src, &ins)? };
    let p = &r.program;
    if let Some(bits) = e["deterministic"].as_str() {
        let res = run(&text, json!({ "shots": 64 }))?;
        if res["counts"] != json!({ bits: 64 }) {
            return Err(format!("counts {}, not every shot {bits}", res["counts"]));
        }
    }
    if let Some(want) = e["distribution"].as_object() {
        if p.dynamic.is_some() {
            return Err("a distribution is checked on a program whose measurements all come at the end".into());
        }
        let (re, im) = amplitudes(p, 0)?;
        let n = p.n_qubits;
        let mut probs = std::collections::BTreeMap::<String, f64>::new();
        for (i, (a, b)) in re.iter().zip(&im).enumerate() {
            let key = p.bitstring(|q| i >> (n - 1 - q) & 1 == 1);
            *probs.entry(key).or_insert(0.0) += a * a + b * b;
        }
        for (k, &v) in &probs {
            let w = want.get(k).and_then(Value::as_f64).unwrap_or(0.0);
            if (v - w).abs() > 1e-9 {
                return Err(format!("P({k}) = {v}, not {w}"));
            }
        }
        if let Some(k) = want.keys().find(|k| !probs.contains_key(*k)) {
            return Err(format!("no outcome {k}"));
        }
    }
    if let Some(want) = e["expectations"].as_object() {
        let obs: Vec<&String> = want.keys().collect();
        let res = run(&text, json!({ "observables": obs }))?;
        for (o, got) in obs.iter().zip(res["expectations"].as_array().ok_or("no expectations")?) {
            let (g, w) = (got["value"].as_f64().unwrap_or(f64::NAN), want[*o].as_f64().unwrap_or(f64::NAN));
            if (g - w).abs() > 1e-9 {
                return Err(format!("⟨{o}⟩ = {g}, not {w}"));
            }
        }
    }
    if let Some(r_src) = case["reference_qasm"].as_str() {
        let q = qasm::read(r_src).map_err(|d| format!("the reference does not read: {d}"))?.program;
        if p.dynamic.is_some() || q.dynamic.is_some() || p.n_qubits != q.n_qubits || p.n_qubits > 10 {
            return Err("an equivalence is checked on two programs of the same width, measured at the end".into());
        }
        // |tr(U†V)| = 2^n exactly when V = e^{iφ} U.
        let d = 1usize << p.n_qubits;
        let (mut tr_re, mut tr_im) = (0.0, 0.0);
        for col in 0..d {
            let (ur, ui) = amplitudes(p, col)?;
            let (vr, vi) = amplitudes(&q, col)?;
            for row in 0..d {
                tr_re += ur[row] * vr[row] + ui[row] * vi[row];
                tr_im += ur[row] * vi[row] - ui[row] * vr[row];
            }
        }
        let overlap = tr_re.hypot(tr_im) / d as f64;
        if (overlap - 1.0).abs() > 1e-9 {
            return Err(format!("not equivalent to the reference: |tr(U†V)|/2^n = {overlap}"));
        }
    }
    if let Some(want) = e["exact"].as_str() {
        match (qasm::exact_circuit(&text), want) {
            (Ok(_), "ok") => {}
            (Err(d), code) if d.code == code => {}
            (Ok(_), code) => return Err(format!("the bit-exact circuit takes it, where it must refuse it with {code}")),
            (Err(d), _) => return Err(format!("the bit-exact circuit refuses it: {d}")),
        }
    }
    if e.get("outputs").is_some() || e.get("timing").is_some() {
        let res = run(&text, json!({ "shots": 16 }))?;
        if let Some(want) = e.get("outputs")
            && &res["outputs"] != want
        {
            return Err(format!("outputs {}, not {want}", res["outputs"]));
        }
        if e.get("timing").is_some() && res["timing"]["applied"] != false {
            return Err(format!("no label for durations not applied: {}", res["timing"]));
        }
    }
    if let Some(n) = e["calibrations"].as_u64()
        && p.extra.calibrations.len() as u64 != n
    {
        return Err(format!("{} calibration statements kept, not {n}", p.extra.calibrations.len()));
    }
    Ok(())
}

#[test]
fn every_program_reads_as_its_manifest_says() {
    let bundle: Value = serde_json::from_str(BUNDLE).expect("bundle.json reads");
    let cases = bundle["cases"].as_array().expect("cases");
    assert!(cases.len() >= 50, "{} cases", cases.len());
    let failures: Vec<String> = cases.iter().filter_map(|c| check(c).err().map(|e| format!("{}: {e}", c["file"].as_str().unwrap_or("?")))).collect();
    assert!(failures.is_empty(), "{} of {} cases:\n{}", failures.len(), cases.len(), failures.join("\n"));
}

/// The bundle is the files: each case's program and reference as the
/// manifest names them, and the manifest's expectations.
#[cfg(not(target_os = "wasi"))]
#[test]
fn the_bundle_is_the_corpus() {
    let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../quantum-conformance/qasm3");
    let manifest: Value = serde_json::from_str(&std::fs::read_to_string(dir.join("manifest.json")).unwrap()).unwrap();
    let cases: Vec<Value> = manifest["cases"]
        .as_array()
        .unwrap()
        .iter()
        .map(|c| {
            let mut d = json!({ "file": c["file"], "qasm": std::fs::read_to_string(dir.join(c["file"].as_str().unwrap())).unwrap(), "expect": c["expect"] });
            if let Some(r) = c["expect"]["equivalent"].as_str() {
                d["reference_qasm"] = json!(std::fs::read_to_string(dir.join(r)).unwrap());
            }
            d
        })
        .collect();
    if std::env::var_os("WAI_QASM3_BUNDLE_WRITE").is_some() {
        let mut bundle: Value = serde_json::from_str(BUNDLE).unwrap_or(json!({}));
        bundle["cases"] = Value::Array(cases);
        std::fs::write(dir.join("bundle.json"), serde_json::to_string_pretty(&bundle).unwrap() + "\n").unwrap();
        return;
    }
    let bundle: Value = serde_json::from_str(BUNDLE).unwrap();
    assert_eq!(bundle["cases"], Value::Array(cases), "bundle.json differs from the files; write it with WAI_QASM3_BUNDLE_WRITE=1");
}