open-standards/wai/wai-quantum/examples/extension_points.rs

//! A host's own components, written outside the engine against its public
//! traits, registered beside the built-ins and run end to end: a simulation
//! method, a compiler pass, a noise channel and a program format.
//!
//! ```sh
//! cargo run --release --example extension_points
//! cargo test --release --example extension_points
//! ```
//!
//! - `basis_table`, a method: a circuit of `x`, `cx`, `ccx` and `swap` gates
//!   keeps one basis state, so its answer is that state, followed bit by bit
//!   and exact. `auto` tries it first for counts; it refuses every other
//!   circuit, and `auto` goes on to the built-ins.
//! - `cancel_pairs`, a pass: adjacent pairs of one self-inverse gate on the
//!   same qubits removed. The engine checks its output against its input by
//!   the equivalence ladder, as it declares.
//! - `bit_flip`, a noise channel: `X` with probability `p` after each gate,
//!   given by its Kraus operators and its Pauli form, which the engine checks
//!   at every job's value; the density matrix takes it.
//! - `gate_list`, a format: `qubits N`, then a gate and its qubits a line,
//!   and `measure`; read to OpenQASM 3 and written from the engine's program.
//!
//! Every result one of them made names them in `components`, each marked
//! `registered` and `certified: false`, and the receipt seals it.

use serde_json::{Value, json};
use wai_quantum::engine::catalog::Plan;
use wai_quantum::engine::formats::{Exported, Imported};
use wai_quantum::engine::registry::{Channel, Check, Component, FrontEnd, Job, Method, Param, Pass, PassInput, PassOutput, Question, Registry};
use wai_quantum::gates::C;
use wai_quantum::qasm::diag::{Diagnostic, FORMAT_UNREAD, WRITE_CANNOT};
use wai_quantum::qasm::{Kind, Op, Program};

/// A method for circuits that permute basis states.
pub struct BasisTable;

fn permutes(p: &Program) -> bool {
    p.ops.iter().all(|o| matches!(o.kind, Kind::X | Kind::Cx | Kind::Ccx | Kind::Swap | Kind::Id))
}

/// The basis state the circuit leaves, qubit by qubit.
fn follow(p: &Program) -> Vec<bool> {
    let mut b = vec![false; p.n_qubits as usize];
    for o in &p.ops {
        let q: Vec<usize> = o.qubits.iter().map(|&q| q as usize).collect();
        match o.kind {
            Kind::X => b[q[0]] = !b[q[0]],
            Kind::Cx if b[q[0]] => b[q[1]] = !b[q[1]],
            Kind::Ccx if b[q[0]] && b[q[1]] => b[q[2]] = !b[q[2]],
            Kind::Swap => b.swap(q[0], q[1]),
            _ => {}
        }
    }
    b
}

impl Method for BasisTable {
    fn name(&self) -> &str {
        "basis_table"
    }
    fn version(&self) -> &str {
        "1.0.0"
    }
    fn classes(&self) -> Vec<String> {
        vec!["exact".into()]
    }
    fn descriptor(&self) -> Value {
        json!({ "circuits": "x, cx, ccx and swap alone: a permutation of basis states", "answers": ["counts", "amplitudes"] })
    }
    fn auto_rank(&self, question: Question) -> Option<u32> {
        // Before every built-in, for counts alone.
        (question == Question::Counts).then_some(5)
    }
    fn plan(&self, job: &Job<'_>) -> Result<Plan, String> {
        let r = job.request()?;
        if !permutes(&r.program) {
            return Err("basis_table takes x, cx, ccx and swap alone".into());
        }
        if !r.observables.is_empty() {
            return Err("basis_table answers counts and amplitudes".into());
        }
        let n = r.program.n_qubits;
        Ok(Plan { method: "basis_table".into(), why: format!("a permutation of basis states on {n} qubits, followed bit by bit"), memory_bytes: 1 << 20, work: r.program.ops.len() as u64, qubits: n, target: None })
    }
    fn run(&self, job: &Job<'_>, _: &Plan, out: &mut Value) -> Result<(), String> {
        let r = job.request()?;
        let b = follow(&r.program);
        out["accuracy"] = json!({ "class": "exact", "note": "the one basis state a permutation leaves, followed bit by bit" });
        if r.shots > 0 {
            let mut counts = serde_json::Map::new();
            counts.insert(r.program.bitstring(|q| b[q as usize]), json!(r.shots));
            out["counts"] = Value::Object(counts);
        }
        if !r.amplitudes.is_empty() {
            let here: String = (0..b.len()).rev().map(|q| if b[q] { '1' } else { '0' }).collect();
            out["amplitudes"] = r.amplitudes.iter().map(|a| json!({ "bits": a, "re": if *a == here { 1.0 } else { 0.0 }, "im": 0.0 })).collect();
        }
        Ok(())
    }
}

/// A pass: adjacent pairs of one self-inverse gate on the same qubits removed.
pub struct CancelPairs;

impl Pass for CancelPairs {
    fn name(&self) -> &str {
        "cancel_pairs"
    }
    fn version(&self) -> &str {
        "1.0.0"
    }
    fn descriptor(&self) -> Value {
        json!({ "does": "adjacent pairs of one self-inverse gate on the same qubits removed" })
    }
    fn check(&self) -> Check {
        Check::Ladder
    }
    fn run(&self, input: &PassInput<'_>) -> Result<PassOutput, String> {
        let inverse = |k: Kind| matches!(k, Kind::X | Kind::Y | Kind::Z | Kind::H | Kind::Cx | Kind::Cz | Kind::Swap | Kind::Ccx);
        let mut out: Vec<Op> = Vec::new();
        let mut dropped = 0;
        for o in input.ops {
            if inverse(o.kind) && out.last().is_some_and(|l| l.kind == o.kind && l.qubits == o.qubits) {
                out.pop();
                dropped += 1;
            } else {
                out.push(o.clone());
            }
        }
        Ok(PassOutput::new(out).with_note(json!({ "how": "adjacent pairs dropped", "pairs": dropped })))
    }
}

/// A channel: the bit flip.
pub struct BitFlip;

fn re(x: f64) -> C {
    C { re: x, im: 0.0 }
}

impl Channel for BitFlip {
    fn name(&self) -> &str {
        "bit_flip"
    }
    fn version(&self) -> &str {
        "1.0.0"
    }
    fn title(&self) -> &str {
        "bit flip"
    }
    fn descriptor(&self) -> Value {
        json!({ "does": "X with probability p after each gate, on each qubit it touched" })
    }
    fn params(&self) -> Vec<Param> {
        vec![Param::new("p", 0.0, 1.0, "a probability")]
    }
    fn arity(&self) -> u32 {
        1
    }
    fn kraus(&self, v: &[f64]) -> Vec<Vec<C>> {
        let (a, b) = ((1.0 - v[0]).sqrt(), v[0].sqrt());
        vec![vec![re(a), re(0.0), re(0.0), re(a)], vec![re(0.0), re(b), re(b), re(0.0)]]
    }
    fn pauli(&self, v: &[f64]) -> Option<Vec<(f64, String)>> {
        Some(vec![(1.0 - v[0], "I".into()), (v[0], "X".into())])
    }
}

/// A format: a list of gates.
pub struct GateList;

impl FrontEnd for GateList {
    fn name(&self) -> &str {
        "gate_list"
    }
    fn version(&self) -> &str {
        "1.0.0"
    }
    fn title(&self) -> &str {
        "a gate list"
    }
    fn descriptor(&self) -> Value {
        json!({ "reads": "qubits N, then a gate and its qubits a line, and measure", "writes": "the same, for gates without parameters" })
    }
    fn extensions(&self) -> Vec<String> {
        vec!["gates".into()]
    }
    fn read(&self, input: &[u8]) -> Result<Imported, Diagnostic> {
        let text = std::str::from_utf8(input).map_err(|_| Diagnostic::new(FORMAT_UNREAD, None, "a gate list is text"))?;
        let mut q = String::from("OPENQASM 3.0;\n");
        let mut declared = false;
        for (i, line) in text.lines().map(str::trim).filter(|l| !l.is_empty()).enumerate() {
            let w: Vec<&str> = line.split_whitespace().collect();
            match (i, w.as_slice()) {
                (0, ["qubits", k]) => {
                    let k: u32 = k.parse().map_err(|_| Diagnostic::new(FORMAT_UNREAD, None, "the first line is qubits N"))?;
                    q.push_str(&format!("qubit[{k}] q;\nbit[{k}] c;\n"));
                    declared = true;
                }
                (0, _) => return Err(Diagnostic::new(FORMAT_UNREAD, None, "a gate list starts with qubits N")),
                (_, ["measure"]) => q.push_str("c = measure q;\n"),
                (_, [g, qs @ ..]) => q.push_str(&format!("{g} {};\n", qs.iter().map(|x| format!("q[{x}]")).collect::<Vec<_>>().join(", "))),
                _ => {}
            }
        }
        if !declared {
            return Err(Diagnostic::new(FORMAT_UNREAD, None, "an empty gate list"));
        }
        Ok(Imported { qasm: q, notes: Vec::new(), summary: json!({ "format": "gate_list" }) })
    }
    fn write(&self, p: &Program, _: Option<&str>) -> Result<Exported, Diagnostic> {
        let mut s = format!("qubits {}\n", p.n_qubits);
        for o in &p.ops {
            if !o.params.is_empty() {
                return Err(Diagnostic::new(WRITE_CANNOT, None, "a gate list holds gates without parameters"));
            }
            let qs: Vec<String> = o.qubits.iter().map(u32::to_string).collect();
            s.push_str(&format!("{} {}\n", format!("{:?}", o.kind).to_lowercase(), qs.join(" ")));
        }
        if !p.measured.is_empty() {
            s.push_str("measure\n");
        }
        Ok(Exported { bytes: s.into_bytes(), notes: Vec::new(), summary: json!({ "format": "gate_list" }) })
    }
}

/// The built-ins, and the four components above after them.
pub fn registry() -> Registry {
    let mut r = Registry::with_builtins();
    for c in [Component::method(BasisTable), Component::pass(CancelPairs), Component::channel(BitFlip), Component::front_end(GateList)] {
        let id = r.register(c).expect("registers");
        assert!(!id.is_builtin());
    }
    r
}

/// A gate list: `x` on each qubit, the pair on qubit 2 a pair the pass
/// cancels, then `cx` and `swap`.
pub const LIST: &str = "qubits 3\nx 0\nx 1\nx 2\nx 2\ncx 0 1\nswap 0 2\nmeasure\n";

/// The four end to end: a gate list read as a job's source, compiled with
/// the pass, run on the method, and the same program under the channel.
pub fn walk_through(r: &Registry) -> Result<Value, String> {
    let list = LIST;
    let job = r.expand(&json!({ "source": { "format": "gate_list", "program": list }, "shots": 100, "compile": { "level": 0, "basis": "standard", "passes": ["cancel_pairs"] } }))?.ok_or("a job with a source expands")?;
    let (plan, result) = r.plan_and_run(&job, 1 << 30)?;
    let noisy = json!({ "qasm": job["qasm"], "observables": ["Z0", "Z2"], "noise": { "bit_flip": 0.05 } });
    let (noisy_plan, noisy_result) = r.plan_and_run(&noisy, 1 << 30)?;
    let written = r.convert_json(&json!({ "from": "qasm", "to": "gate_list", "program": job["qasm"] }))?;
    Ok(json!({ "plan": plan.method, "why": plan.why, "result": result, "noisy_plan": noisy_plan.method, "noisy": noisy_result, "written": written }))
}

fn main() {
    let r = registry();
    let c = r.components();
    println!("{} components, {} registered:", r.identities().len(), c["registered"]);
    for id in r.identities().into_iter().filter(|i| !i.is_builtin()) {
        println!("  {:<10} {:<14} {:<7} {}", id.kind.as_str(), id.name, id.version, id.digest);
    }
    match walk_through(&r) {
        Ok(w) => {
            println!("\nthe job planned on {}: {}", w["plan"], w["why"]);
            println!("counts {}", w["result"]["counts"]);
            println!("compiled by {}", w["result"]["compiled"]["source"]);
            println!("under bit_flip on {}: {}", w["noisy_plan"], w["noisy"]["expectations"]);
            println!("written back as a gate list:\n{}", w["written"]["program"].as_str().unwrap_or(""));
            println!("made by: {}", serde_json::to_string_pretty(&w["result"]["components"]).unwrap_or_default());
        }
        Err(e) => {
            eprintln!("{e}");
            std::process::exit(1);
        }
    }
}

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

    #[test]
    fn the_four_run_end_to_end_and_are_named_in_what_they_make() {
        let r = registry();
        let w = walk_through(&r).unwrap();
        // x on every qubit, the pair on q[2] cancelled, then cx and swap: q[2] alone reads 1.
        assert_eq!(w["plan"], "basis_table", "{}", w["why"]);
        assert_eq!(w["result"]["counts"], json!({ "100": 100 }), "{w}");
        let names: Vec<(&str, &str)> = w["result"]["components"].as_array().unwrap().iter().map(|c| (c["name"].as_str().unwrap(), c["origin"].as_str().unwrap())).collect();
        assert_eq!(names, [("gate_list", "registered"), ("lower", "built-in"), ("cancel_pairs", "registered"), ("basis_table", "registered")]);
        assert!(w["result"]["components"].as_array().unwrap().iter().filter(|c| c["origin"] == "registered").all(|c| c["certified"] == false));
        // The same counts as the bit-exact method, a built-in, gives.
        let mut exact = r.expand(&json!({ "source": { "format": "gate_list", "program": LIST }, "shots": 100 })).unwrap().unwrap();
        exact["method"] = json!("exact");
        assert_eq!(r.plan_and_run(&exact, 1 << 30).unwrap().1["counts"], w["result"]["counts"]);
        // Under the bit flip, on the density matrix: ⟨Z⟩ of a flipped qubit
        // after the gates that touch it, each flip with probability 0.05.
        assert_eq!(w["noisy_plan"], "density_matrix");
        let z: Vec<f64> = w["noisy"]["expectations"].as_array().unwrap().iter().map(|e| e["value"].as_f64().unwrap()).collect();
        assert!(z.iter().all(|v| v.abs() < 1.0 && v.abs() > 0.5), "{z:?}");
        // Written back as a gate list, read again and held to the program.
        assert_eq!(w["written"]["check"]["verdict"]["verdict"], "equivalent", "{}", w["written"]);
    }

    #[test]
    fn the_receipt_seals_them() {
        let r = registry();
        let key = ed25519_dalek_key();
        let job = json!({ "qasm": "OPENQASM 3.0;\nqubit[2] q;\nbit[2] c;\nx q[0];\ncx q[0], q[1];\nc = measure q;\n", "shots": 10 });
        let body = r.job(&job, 1 << 30, &key, json!({ "unmetered": "an example" })).unwrap();
        assert!(wai_quantum::engine::job::verify(&body).unwrap().0);
        let mut forged = body.clone();
        forged["result"]["components"][1]["origin"] = json!("built-in");
        assert!(!wai_quantum::engine::job::verify(&forged).unwrap().0, "a registered method cannot pass for a built-in");
    }

    fn ed25519_dalek_key() -> ed25519_dalek::SigningKey {
        ed25519_dalek::SigningKey::from_bytes(&[9; 32])
    }

    #[test]
    fn what_they_do_not_take_is_refused_with_its_code() {
        let r = registry();
        let codes = |e: String| wai_quantum::engine::refusal::codes_in(&e).iter().map(|c| c.to_string()).collect::<Vec<_>>();
        let bell = "OPENQASM 3.0;\nqubit[2] q;\nh q[0];\ncx q[0], q[1];\n";
        assert_eq!(codes(r.plan(&json!({ "qasm": bell, "shots": 1, "method": "basis_table" }), 1 << 30).unwrap_err()), ["J011"]);
        assert_eq!(codes(r.plan(&json!({ "qasm": bell, "shots": 1, "method": "stabilizer", "noise": { "bit_flip": 0.1 } }), 1 << 30).unwrap_err()), ["J061"]);
        assert_eq!(codes(r.plan(&json!({ "qasm": bell, "shots": 1, "compile": { "passes": ["no_such_pass"] } }), 1 << 30).unwrap_err()), ["J060"]);
        let mut again = registry();
        assert!(again.register(Component::method(BasisTable)).unwrap_err().ends_with("[J063]"));
        // The built-in registry holds none of them.
        assert!(Registry::builtins().plan(&json!({ "qasm": bell, "shots": 1, "method": "basis_table" }), 1 << 30).unwrap_err().ends_with("[J010]"));
    }
}