open-standards/wai/wai-rs/src/quantum_envelope.rs

//! Wrapping a quantum circuit in a WAI envelope, and opening one.
//!
//! The computation itself lives in the dependency-light [`wai_quantum`] crate;
//! this is the container half — the part that belongs to WAI, not to quantum.
//!
//! The envelope's capability is the circuit's: `wai.quantum.circuit2` for a
//! [`Gateset::V2`] circuit, whose payload is a `WQC2` container, and
//! `wai.quantum.circuit` for a [`Gateset::V1`] one, whose payload is `WQC1`.

use wai_quantum::quantum::{from_wqc, to_wqc, Circuit, Gateset, Measure, QuantumError};

/// Wrap a circuit in a v1.0 WAI envelope under its gateset's capability
/// ([`Gateset::capability`]). The payload is the circuit's `WQC` container;
/// `intent=replicate` is sound because the reconstruction is byte-exact
/// (statevector-equivalence).
///
/// Fails as [`to_wqc`] does, for a circuit its gateset does not allow, and with
/// [`QuantumError::TooManyQubits`] for a circuit of 64 qubits or more, whose
/// dimension the manifest's `target.dim` cannot hold.
pub fn wrap_envelope(c: &Circuit, measure: Option<Measure>) -> Result<crate::container::Wai, QuantumError> {
    use crate::container::{Conditioning, Manifest, ModelRequirement};
    let payload = to_wqc(c, measure)?;
    let dim = 1u64.checked_shl(u32::from(c.n_qubits)).ok_or(QuantumError::TooManyQubits(c.n_qubits))?;
    let manifest = Manifest::new("quantum", "replicate", ModelRequirement::new(c.gateset.capability()), Conditioning::new("wqc"))
        .with_target(serde_json::json!({ "n_qubits": c.n_qubits, "dim": dim }));
    Ok(crate::container::Wai::new(manifest, payload))
}

/// Open a quantum envelope: decode its `WQC` payload, then require the
/// envelope's capability to be the one the container's magic names.
///
/// A payload that does not decode fails as [`from_wqc`] does. A decoded circuit
/// whose gateset is not the one the capability names — a `WQC2` payload under
/// `wai.quantum.circuit`, a `WQC1` payload under `wai.quantum.circuit2`, or any
/// payload under another capability — is [`QuantumError::Contract`].
pub fn open_envelope(w: &crate::container::Wai) -> Result<(Circuit, Option<Measure>), QuantumError> {
    let (c, m) = from_wqc(&w.payload)?;
    let cap = w.capability();
    if Gateset::from_capability(cap) != Some(c.gateset) {
        // The capability is the file's; quote at most 64 characters of it.
        let quoted: String = cap.chars().take(64).collect();
        return Err(QuantumError::Contract(format!(
            "the envelope's capability {quoted:?} does not name the payload's container, {} ({})",
            String::from_utf8_lossy(&c.gateset.magic()),
            c.gateset.capability()
        )));
    }
    Ok((c, m))
}

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

    fn bell(gs: Gateset) -> Circuit {
        let mut c = Circuit::with_gateset(2, gs);
        c.h(0).cx(0, 1);
        c
    }

    /// The envelope's bytes, read back.
    fn reread(w: &Wai) -> Wai {
        Wai::from_bytes(&w.to_bytes().unwrap()).unwrap()
    }

    /// The capability follows the circuit's gateset, and the circuit — gateset
    /// included — and its measurement survive the envelope.
    #[test]
    fn envelope_wraps_under_the_gatesets_capability() {
        for (gs, cap, magic) in [
            (Gateset::V2, crate::codecs::CAP_QUANTUM_CIRCUIT2, b"WQC2"),
            (Gateset::V1, crate::codecs::CAP_QUANTUM_CIRCUIT, b"WQC1"),
        ] {
            // The registry's strings are copies of the crate's.
            assert_eq!(gs.capability(), cap);
            let c = bell(gs);
            let m = Some(Measure { seed: 7, shots: 1024 });
            let wai = wrap_envelope(&c, m).unwrap();
            assert_eq!(wai.capability(), cap);
            assert_eq!(&wai.payload[0..4], magic);
            assert_eq!(wai.manifest.target, serde_json::json!({ "n_qubits": 2, "dim": 4 }));
            let back = reread(&wai);
            assert_eq!(open_envelope(&back).unwrap(), (c.clone(), m));
            let (c2, _) = from_wqc(&back.payload).unwrap();
            assert_eq!(c2, c);
        }
        assert_eq!(wrap_envelope(&Circuit::new(1), None).unwrap().capability(), "wai.quantum.circuit2");
    }

    /// An envelope whose capability is not the one its payload's magic names is
    /// refused, whichever way round, and so is a quantum payload under any
    /// other capability. A payload that is not a container fails as the
    /// container reader does, and a circuit its gateset forbids is never wrapped.
    #[test]
    fn a_capability_that_does_not_name_the_payload_is_refused() {
        for (gs, other) in [(Gateset::V2, "wai.quantum.circuit"), (Gateset::V1, "wai.quantum.circuit2"), (Gateset::V2, "wai.image.png")] {
            let mut wai = wrap_envelope(&bell(gs), None).unwrap();
            wai.manifest.model_requirement.capability = other.to_string();
            let e = open_envelope(&reread(&wai)).unwrap_err();
            assert_eq!(e.kind(), "contract", "{gs:?} under {other}: {e}");
        }
        let mut wai = wrap_envelope(&bell(Gateset::V2), None).unwrap();
        wai.payload = b"not a circuit".to_vec();
        assert_eq!(open_envelope(&wai).unwrap_err().kind(), "malformed");
        // A repeated control: allowed in circuit/1, refused by circuit2.
        let mut c = Circuit::with_gateset(3, Gateset::V2);
        c.ccx(0, 0, 2);
        assert_eq!(wrap_envelope(&c, None).err().expect("refused").kind(), "invalid");
        c.gateset = Gateset::V1;
        assert_eq!(open_envelope(&wrap_envelope(&c, None).unwrap()).unwrap().0, c);
        // A circuit whose dimension the manifest cannot hold.
        assert_eq!(wrap_envelope(&Circuit::new(64), None).err(), Some(QuantumError::TooManyQubits(64)));
    }
}