open-standards/wai/wai-quantum-writer/tests/write.rs

//! The loop, with the scripted model ([`wai_quantum_writer::scripted`]): its replies set in
//! advance, what it is told recorded. No checkpoint is needed, so these run anywhere the
//! engine builds.

use serde_json::{json, Value};
use wai_quantum_writer::scripted::Scripted;
use wai_quantum_writer::{block_closed, check_task, program_of, prompt, write, write_with, Model, Options, NO_ANSWER, SCHEMA};

const BELL: &str = "```qasm\nOPENQASM 3.0;\ninclude \"stdgates.inc\";\nqubit[2] q;\nbit[2] c;\nh q[0];\ncx q[0], q[1];\nc = measure q;\n```";
const NO_CNOT: &str = "```qasm\nOPENQASM 3.0;\ninclude \"stdgates.inc\";\nqubit[2] q;\nbit[2] c;\nh q[0];\nc = measure q;\n```";

fn bell_task() -> Value {
    json!({
        "schema": "wai.quantum.exercise/1",
        "title": "A Bell pair",
        "task": "Entangle two qubits so that their measured bits always agree.",
        "starter": ["OPENQASM 3.0;", "include \"stdgates.inc\";", "qubit[2] q;", "bit[2] c;", "c = measure q;"],
        "intent": { "distribution": { "00": 0.5, "11": 0.5 } }
    })
}

#[test]
fn a_program_that_passes_at_once_ends_the_loop() {
    let mut m = Scripted::new(&[Some(BELL)]);
    let w = write(&mut m, &bell_task(), &Options::default()).unwrap();
    let t = &w.transcript;
    assert_eq!(t["schema"], SCHEMA);
    assert_eq!((t["passed"].as_bool(), t["passed_round"].as_u64(), t["level"].as_str()), (Some(true), Some(0), Some("proof")), "{t:#}");
    assert_eq!(t["attempts"].as_array().unwrap().len(), 1);
    assert!(t["program"].as_str().unwrap().contains("cx q[0], q[1];"));
    // The task went to the model in words, with its starter program.
    assert!(m.told[1].contains("A Bell pair") && m.told[1].contains("00 with probability 0.5") && m.told[1].contains("qubit[2] q;"), "{}", m.told[1]);
}

#[test]
fn what_went_wrong_goes_back_in_the_engines_words_and_the_repair_passes() {
    let mut m = Scripted::new(&[Some(NO_CNOT), Some(BELL)]);
    let w = write(&mut m, &bell_task(), &Options::default()).unwrap();
    let t = &w.transcript;
    assert_eq!(t["passed_round"].as_u64(), Some(1), "{t:#}");
    assert_eq!(t["attempts"][0]["passed"], false);
    assert_eq!(t["attempts"][0]["check"]["what_went_wrong"][0]["kind"], "distribution");
    let told = &m.told[2];
    assert!(told.starts_with("The engine checked that program"), "{told}");
    assert!(told.contains("the engine:") && told.contains("Where to look:"), "{told}");
    assert_eq!(w.timing.len(), 2);
}

#[test]
fn a_reply_that_never_leaves_its_thinking_is_no_answer_and_is_asked_again() {
    let mut m = Scripted::new(&[None, Some(BELL)]);
    let w = write(&mut m, &bell_task(), &Options::default()).unwrap();
    let t = &w.transcript;
    assert_eq!((t["attempts"][0]["answered"].as_bool(), t["attempts"][0]["check"].is_null()), (Some(false), true), "{t:#}");
    assert_eq!(m.told[2], NO_ANSWER);
    assert_eq!(t["passed_round"].as_u64(), Some(1));
}

#[test]
fn the_rounds_run_out_without_a_pass() {
    let mut m = Scripted::new(&[Some(NO_CNOT); 3]);
    let w = write(&mut m, &bell_task(), &Options { rounds: 2, ..Options::default() }).unwrap();
    let t = &w.transcript;
    assert_eq!((t["passed"].as_bool(), t["passed_round"].is_null(), t["level"].is_null()), (Some(false), true, true), "{t:#}");
    assert_eq!(t["attempts"].as_array().unwrap().len(), 3);
}

#[test]
fn the_same_model_task_and_options_give_the_same_transcript_bytes() {
    let run = || {
        let mut m = Scripted::new(&[Some(NO_CNOT), Some(BELL)]);
        serde_json::to_vec(&write(&mut m, &bell_task(), &Options::default()).unwrap().transcript).unwrap()
    };
    assert_eq!(run(), run());
}

/// A task that asks nothing is refused before the model is asked: the model here would
/// answer with a program that passes anything, so only the task's check can refuse it.
#[test]
fn a_task_needs_something_the_program_must_do() {
    let mut m = Scripted::new(&[Some(BELL)]);
    let r = write(&mut m, &json!({ "task": "anything" }), &Options::default());
    assert!(r.is_err(), "a task that asks nothing was written for: {:#}", r.unwrap().transcript);
    assert!(m.told.is_empty(), "the model was asked");
}

#[test]
fn the_program_is_the_answers_code_block() {
    assert_eq!(program_of("Here:\n```qasm\nOPENQASM 3.0;\nqubit q;\n```\nDone."), "OPENQASM 3.0;\nqubit q;\n");
    assert_eq!(program_of("OPENQASM 3.0;\nqubit q;"), "OPENQASM 3.0;\nqubit q;\n");
    assert!(block_closed("```qasm\nx\n```") && !block_closed("```qasm\nx\n"));
    assert!(prompt(&bell_task()).starts_with("Task: A Bell pair"));
}

#[test]
fn each_attempt_is_handed_on_as_it_lands() {
    let mut m = Scripted::new(&[Some(NO_CNOT), Some(BELL)]);
    let mut seen: Vec<(u64, bool)> = Vec::new();
    let w = write_with(&mut m, &bell_task(), &Options::default(), &mut |a, (reply, check)| {
        assert!(reply >= 0.0 && check >= 0.0);
        seen.push((a["round"].as_u64().unwrap(), a["passed"].as_bool().unwrap()));
    })
    .unwrap();
    assert_eq!(seen, [(0, false), (1, true)]);
    // What was handed on is what the transcript holds, and write gives the same bytes.
    let again = write(&mut Scripted::new(&[Some(NO_CNOT), Some(BELL)]), &bell_task(), &Options::default()).unwrap();
    assert_eq!(serde_json::to_vec(&w.transcript).unwrap(), serde_json::to_vec(&again.transcript).unwrap());
}

/// A task must give the engine something to check: an intent, or cases that each have
/// one (a case without one takes the task's).
#[test]
fn a_task_reads_only_with_something_to_check() {
    let case = |intent: Option<Value>| {
        let mut c = json!({ "name": "n = 2", "inputs": { "n": 2 } });
        if let Some(i) = intent {
            c["intent"] = i;
        }
        c
    };
    let dist = json!({ "distribution": { "0": 1.0 } });
    assert!(check_task(&bell_task()).is_ok());
    assert!(check_task(&json!({ "cases": [case(Some(dist.clone()))] })).is_ok());
    assert!(check_task(&json!({ "intent": dist, "cases": [case(None)] })).is_ok());
    assert!(check_task(&json!({ "intent": [dist, dist] })).is_ok());
    for (task, why) in [
        (json!("a Bell pair"), "not a JSON object"),
        (json!({ "task": "anything" }), "neither"),
        (json!({ "cases": [case(Some(dist.clone())), case(None)] }), "case 2 has no intent"),
        (json!({ "cases": [] }), "not a list of one or more"),
        (json!({ "intent": "00 and 11" }), "neither an object nor a list"),
        (json!({ "intent": [] }), "neither an object nor a list"),
        (json!({ "schema": "wai.quantum.job/1", "intent": dist }), "its schema"),
        (json!({ "cases": [{ "inputs": 3, "intent": dist }] }), "inputs are not an object"),
        (json!({ "intent": dist, "cases": [{ "name": "a", "inputs": { "err": 0 } }, { "name": "b" }] }), "no program can pass every case"),
        (json!({ "intent": dist, "cases": [{ "inputs": { "err": 0 } }, { "inputs": { "n": 0 } }] }), "no program can pass every case"),
    ] {
        let e = check_task(&task).unwrap_err();
        assert!(e.contains(why), "{task}: {e}");
    }
    // write refuses what check_task refuses, before the model is asked anything.
    let mut m = Scripted::new(&[Some(BELL)]);
    assert!(write(&mut m, &json!({ "cases": [case(None)] }), &Options::default()).is_err());
    assert!(m.told.is_empty());
}

/// A script reads from JSON, names itself by its digest, plays from its first reply in
/// each conversation, and says so when it runs out.
#[test]
fn a_script_plays_from_its_start_in_each_conversation() {
    let m = Scripted::from_json(&json!([null, BELL])).unwrap();
    assert_eq!(m.script(), json!([null, BELL]));
    assert_eq!(m.identity()["digest"], wai_quantum_writer::digest(&json!([null, BELL])));
    assert_ne!(m.identity()["digest"], Scripted::new(&[Some(BELL)]).identity()["digest"]);
    assert!(Scripted::from_json(&json!([1])).unwrap_err().contains("reply 1"));
    assert!(Scripted::from_json(&json!({ "replies": [] })).is_err());
    let mut m = m;
    let first = write(&mut m, &bell_task(), &Options::default()).unwrap();
    let second = write(&mut m, &bell_task(), &Options::default()).unwrap();
    assert_eq!(first.transcript, second.transcript);
    assert_eq!(first.transcript["passed_round"], 1);
    let mut short = Scripted::new(&[Some(NO_CNOT)]);
    let e = write(&mut short, &bell_task(), &Options::default()).unwrap_err();
    assert!(e.contains("the script has 1 replies"), "{e}");
}

/// A task is read in its canonical form, the form its digest names: the same task with its
/// keys in another order, or a whole number written as a float, is the same prompt, the same
/// check and the same transcript, and a bound written `0.0` is the bound 0.
#[test]
fn one_task_digest_names_one_prompt_and_one_check() {
    let run = |task: Value| {
        let mut m = Scripted::new(&[Some(BELL)]);
        let w = write(&mut m, &task, &Options { rounds: 0, ..Options::default() }).unwrap();
        (serde_json::to_vec(&w.transcript).unwrap(), m.told[1].clone(), w.transcript)
    };
    let task = |d: Value, max: Value| json!({ "title": "A Bell pair", "intent": { "distribution": d }, "constraints": { "counts": { "cx": { "max": max } } } });
    let (a, prompt_a, t) = run(task(json!({ "00": 0.5, "11": 0.5 }), json!(0)));
    let (b, prompt_b, _) = run(task(json!({ "11": 0.5, "00": 0.5 }), json!(0.0)));
    assert_eq!(prompt_a, prompt_b, "one task, two prompts");
    assert_eq!(a, b, "one task, two transcripts");
    // The bound holds: the Bell pair has a cx, and the task allows none.
    assert_eq!((t["passed"].as_bool(), t["attempts"][0]["check"]["constraints"][0]["holds"].as_bool()), (Some(false), Some(false)), "{t:#}");
}

/// A task is held to the course's form before the model is asked anything: a bound that is
/// not whole numbers, a constraint the course does not name, or a job that asks for more than
/// shots, seed and noise (a repair the transcript would not show) is refused.
#[test]
fn a_malformed_task_is_refused_before_the_model_is_asked() {
    for extra in [
        json!({ "constraints": { "t_count": 0 } }),
        json!({ "constraints": { "depth": { "max": 3 } } }),
        json!({ "constraints": { "counts": { "cx": { "max": -1 } } } }),
        json!({ "job": { "repair": true } }),
        json!({ "job": { "dyadic_bits": 3 } }),
        json!({ "job": 1 }),
        json!({ "cases": vec![json!({ "inputs": {} }); 100] }),
    ] {
        let mut task = bell_task();
        for (k, v) in extra.as_object().unwrap() {
            task[k] = v.clone();
        }
        let mut m = Scripted::new(&[Some(BELL)]);
        let r = write(&mut m, &task, &Options::default());
        assert!(r.is_err(), "{task} was written for: {:#}", r.unwrap().transcript);
        assert!(m.told.is_empty(), "{task}: the model was asked");
    }
}

/// The transcript records what decides its checks and its replies beside the model: the
/// memory each check's run was planned in, the thinking cap, and each reply's thinking closed
/// or not.
#[test]
fn the_transcript_records_the_check_s_memory_and_the_thinking_cap() {
    let mut m = Scripted::new(&[Some(BELL)]);
    let opts = Options { job_bytes: 512 << 20, max_new: 900, ..Options::default() };
    let t = write(&mut m, &bell_task(), &opts).unwrap().transcript;
    assert_eq!(t["check_memory_bytes"].as_u64(), Some(512 << 20), "{t:#}");
    assert_eq!(t["decoding"], json!({ "method": "greedy", "max_new": 900, "think_max": 600, "rounds": 4 }), "{t:#}");
    assert_eq!(t["attempts"][0]["thinking_closed"], false, "{t:#}");
    assert_eq!(Options { think_max: Some(5000), ..Options::default() }.thinking_cap(), 3000);
}

/// An answer with nothing in it is no program: the model is asked again, and the engine is not
/// asked to check an empty one.
#[test]
fn an_empty_answer_is_no_program() {
    let mut m = Scripted::new(&[Some("  \n "), Some(BELL)]);
    let t = write(&mut m, &bell_task(), &Options::default()).unwrap().transcript;
    assert_eq!((t["attempts"][0]["answered"].as_bool(), t["attempts"][0]["check"].is_null()), (Some(false), true), "{t:#}");
    assert_eq!(m.told[2], NO_ANSWER);
    assert_eq!(t["passed_round"].as_u64(), Some(1));
}