open-standards/wai/wai-quantum-cli/tests/dap.rs

//! `wai-quantum dap` driven as an editor drives it: the binary started, a
//! scripted session of debug adapter protocol messages written to its
//! standard input, and every response and event read back from its standard
//! output and checked.

use serde_json::{Value, json};
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, ChildStdin, Command, Stdio};
use std::sync::mpsc::{Receiver, channel};
use std::time::Duration;

struct Editor {
    child: Child,
    stdin: ChildStdin,
    rx: Receiver<Value>,
    seq: i64,
    events: Vec<Value>,
}

fn read_message(r: &mut impl BufRead) -> Option<Value> {
    let mut len = 0usize;
    loop {
        let mut line = String::new();
        if r.read_line(&mut line).ok()? == 0 {
            return None;
        }
        let line = line.trim_end();
        if line.is_empty() {
            break;
        }
        if let Some(v) = line.strip_prefix("Content-Length:") {
            len = v.trim().parse().ok()?;
        }
    }
    let mut buf = vec![0u8; len];
    r.read_exact(&mut buf).ok()?;
    serde_json::from_slice(&buf).ok()
}

impl Editor {
    fn start() -> Editor {
        let mut child = Command::new(env!("CARGO_BIN_EXE_wai-quantum")).arg("dap").stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::null()).spawn().expect("the binary starts");
        let stdin = child.stdin.take().unwrap();
        let mut stdout = BufReader::new(child.stdout.take().unwrap());
        let (tx, rx) = channel();
        std::thread::spawn(move || {
            while let Some(m) = read_message(&mut stdout) {
                if tx.send(m).is_err() {
                    break;
                }
            }
        });
        Editor { child, stdin, rx, seq: 1, events: Vec::new() }
    }

    fn recv(&mut self) -> Value {
        self.rx.recv_timeout(Duration::from_secs(120)).expect("the adapter answers within two minutes")
    }

    /// A request's response; the events that come first are kept.
    fn request(&mut self, command: &str, arguments: Value) -> Value {
        let seq = self.seq;
        self.seq += 1;
        let body = json!({ "seq": seq, "type": "request", "command": command, "arguments": arguments }).to_string();
        write!(self.stdin, "Content-Length: {}\r\n\r\n{body}", body.len()).unwrap();
        self.stdin.flush().unwrap();
        loop {
            let m = self.recv();
            if m["type"] == "response" && m["request_seq"] == seq {
                assert_eq!(m["command"], command, "{m}");
                return m;
            }
            self.events.push(m);
        }
    }

    fn ok(&mut self, command: &str, arguments: Value) -> Value {
        let r = self.request(command, arguments);
        assert_eq!(r["success"], true, "{command}: {r}");
        r["body"].clone()
    }

    /// The next event of this name (kept ones first).
    fn event(&mut self, name: &str) -> Value {
        if let Some(i) = self.events.iter().position(|m| m["event"] == name) {
            return self.events.remove(i);
        }
        loop {
            let m = self.recv();
            if m["type"] == "event" && m["event"] == name {
                return m;
            }
            self.events.push(m);
        }
    }

    fn line(&mut self) -> u64 {
        let st = self.ok("stackTrace", json!({ "threadId": 1 }));
        st["stackFrames"][0]["line"].as_u64().unwrap()
    }

    fn vars(&mut self, reference: i64) -> Vec<Value> {
        self.ok("variables", json!({ "variablesReference": reference }))["variables"].as_array().unwrap().clone()
    }
}

/// Teleportation of |1⟩: two measurements in the middle (lines 9 and 10),
/// corrections conditioned on them, q2 measured at the end (line 13).
const TELEPORT: &str = "OPENQASM 3.0;\nqubit[3] q;\nbit[3] c;\nx q[0];\nh q[1];\ncx q[1], q[2];\ncx q[0], q[1];\nh q[0];\nc[0] = measure q[0];\nc[1] = measure q[1];\nif (c[1]) x q[2];\nif (c[0]) z q[2];\nc[2] = measure q[2];\n";

fn program(name: &str, text: &str) -> std::path::PathBuf {
    let dir = std::env::temp_dir().join(format!("wai-quantum-dap-{name}-{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    let p = dir.join("teleport.qasm");
    std::fs::write(&p, text).unwrap();
    p
}

#[test]
fn an_editor_s_session_from_initialize_to_disconnect() {
    let path = program("session", TELEPORT);
    let mut e = Editor::start();
    let caps = e.ok("initialize", json!({ "clientID": "test", "adapterID": "wai-quantum", "linesStartAt1": true, "columnsStartAt1": true, "pathFormat": "path" }));
    assert_eq!((caps["supportsStepBack"].as_bool(), caps["supportsConfigurationDoneRequest"].as_bool()), (Some(true), Some(true)), "{caps}");
    e.ok("launch", json!({ "program": path, "stopOnEntry": false }));
    e.event("initialized");
    // A breakpoint on a measurement is verified; one on a declaration is not.
    let bps = e.ok("setBreakpoints", json!({ "source": { "path": path }, "breakpoints": [{ "line": 9 }, { "line": 2 }] }));
    assert_eq!(bps["breakpoints"][0], json!({ "id": 1, "verified": true, "line": 9 }));
    assert_eq!(bps["breakpoints"][1]["verified"], false);
    let locs = e.ok("breakpointLocations", json!({ "source": { "path": path }, "line": 1, "endLine": 20 }));
    assert_eq!(locs["breakpoints"].as_array().unwrap().len(), 10, "{locs}");
    e.ok("configurationDone", json!({}));
    let st = e.event("stopped");
    assert_eq!((st["body"]["reason"].as_str(), st["body"]["hitBreakpointIds"].clone()), (Some("breakpoint"), json!([1])), "{st}");
    assert_eq!(e.ok("threads", json!({}))["threads"][0]["id"], 1);
    let frames = e.ok("stackTrace", json!({ "threadId": 1 }));
    let f = &frames["stackFrames"][0];
    assert_eq!((f["line"].as_u64(), f["source"]["path"].as_str()), (Some(9), path.to_str()), "{frames}");
    assert!(f["name"].as_str().unwrap().starts_with("measure q0 into bit 0 (position 5)"), "{frames}");

    // The scopes: the engine's probes.
    let scopes = e.ok("scopes", json!({ "frameId": 1 }));
    let names: Vec<&str> = scopes["scopes"].as_array().unwrap().iter().map(|s| s["name"].as_str().unwrap()).collect();
    assert_eq!(names, ["Qubits", "State", "Classical bits", "Path", "Branches"]);
    let qubits = e.vars(1);
    assert_eq!(qubits.len(), 3);
    // q0 is in |−⟩ on its own; q1 shares a Bell pair with q2.
    assert!(qubits[0]["value"].as_str().unwrap().ends_with("purity 1.000000"), "{qubits:?}");
    assert!(qubits[1]["value"].as_str().unwrap().ends_with("purity 0.500000"), "{qubits:?}");
    let q0 = e.vars(qubits[0]["variablesReference"].as_i64().unwrap());
    assert_eq!(q0.iter().map(|v| v["name"].as_str().unwrap()).collect::<Vec<_>>(), ["x", "y", "z", "length", "purity", "p(1)"]);
    let state = e.vars(2);
    assert!(state.iter().any(|v| v["name"] == "method" && v["value"].as_str().unwrap().starts_with("exact (bit-exact)")), "{state:?}");
    assert_eq!(state.iter().filter(|v| v["name"].as_str().unwrap().starts_with('|')).count(), 4, "four amplitudes: {state:?}");
    assert_eq!(e.vars(3)[0]["value"], "000 (the highest first)");

    // A step over the measurement: the console says what it read.
    e.ok("next", json!({ "threadId": 1 }));
    assert_eq!(e.event("stopped")["body"]["reason"], "step");
    let said = e.event("output");
    assert!(said["body"]["output"].as_str().unwrap().starts_with("measure q0 (line 9) read 0 with probability 0.5"), "{said}");
    assert_eq!(e.line(), 10);
    let path_vars = e.vars(4);
    assert!(path_vars.iter().any(|v| v["name"] == "measure q0 (line 9)" && v["value"].as_str().unwrap().starts_with("0 with p 0.5")), "{path_vars:?}");
    // And back, to the same state.
    e.ok("stepBack", json!({ "threadId": 1 }));
    e.event("stopped");
    assert_eq!(e.line(), 9);
    assert_eq!(e.vars(4)[0]["value"], "5");

    // The branch tree: four leaves, a quarter each, exactly.
    let br = e.vars(5);
    assert!(br[0]["value"].as_str().unwrap().starts_with("exact"), "{br:?}");
    let first = e.vars(br[1]["variablesReference"].as_i64().unwrap());
    assert_eq!(first.len(), 2, "{first:?}");
    assert!(first[0]["name"].as_str().unwrap().starts_with("reads 0, p 1/2, weight 1/2 (0.500000): then measure q1 into bit 1 (line 10)"), "{first:?}");
    let second = e.vars(first[1]["variablesReference"].as_i64().unwrap());
    assert!(second.iter().all(|v| v["value"].as_str().unwrap().starts_with("weight 1/4 (0.250000)")), "{second:?}");

    // Another path, by evaluation: both read 1, so both corrections run, and q2 reads 1.
    let r = e.ok("evaluate", json!({ "expression": "outcomes 1 1", "context": "repl" }));
    assert!(r["result"].as_str().unwrap().contains("outcomes 11"), "{r}");
    e.event("stopped");
    let cut = e.ok("evaluate", json!({ "expression": "cut 2", "context": "repl" }));
    assert!(cut["result"].as_str().unwrap().contains("bits across [2]"), "{cut}");
    e.ok("continue", json!({ "threadId": 1 }));
    let end = e.event("stopped");
    assert_eq!(end["body"]["reason"], "end", "{end}");
    assert_eq!(e.vars(3)[0]["value"], "111 (the highest first)");
    assert!(e.ok("stackTrace", json!({ "threadId": 1 }))["stackFrames"][0]["name"].as_str().unwrap().starts_with("the end of the program"));
    // Back to the breakpoint, then on to the end, and once more ends the session.
    e.ok("reverseContinue", json!({ "threadId": 1 }));
    assert_eq!(e.event("stopped")["body"]["reason"], "breakpoint");
    assert_eq!(e.line(), 9);
    e.ok("continue", json!({ "threadId": 1 }));
    assert_eq!(e.event("stopped")["body"]["reason"], "end");
    e.ok("continue", json!({ "threadId": 1 }));
    e.event("terminated");
    assert_eq!(e.event("exited")["body"]["exitCode"], 0);
    e.ok("disconnect", json!({}));
    assert_eq!(e.child.wait().unwrap().code(), Some(0));
    let _ = std::fs::remove_dir_all(path.parent().unwrap());
}

#[test]
fn a_program_too_wide_is_refused_at_launch_with_the_reason_and_stop_on_entry_stops_first() {
    let wide = program("wide", "OPENQASM 3.0;\nqubit[17] q;\nh q[0];\n");
    let mut e = Editor::start();
    e.ok("initialize", json!({}));
    let r = e.request("launch", json!({ "program": wide }));
    assert_eq!(r["success"], false);
    assert!(r["message"].as_str().unwrap().contains("at most 16"), "{r}");
    assert_eq!(e.request("stackTrace", json!({ "threadId": 1 }))["success"], false);
    // A noisy program, from its first step, on the density matrix.
    let small = program("entry", TELEPORT);
    e.ok("launch", json!({ "program": small, "stopOnEntry": true, "noise": { "depolarizing_2q": 0.02 } }));
    e.event("initialized");
    e.ok("configurationDone", json!({}));
    assert_eq!(e.event("stopped")["body"]["reason"], "entry");
    assert_eq!(e.line(), 4);
    let state = e.vars(2);
    assert!(state.iter().any(|v| v["name"] == "method" && v["value"].as_str().unwrap().starts_with("density_matrix")), "{state:?}");
    assert!(e.request("evaluate", json!({ "expression": "nonsense" }))["message"].as_str().unwrap().contains("evaluate takes"));
    e.ok("disconnect", json!({}));
    assert_eq!(e.child.wait().unwrap().code(), Some(0));
    let _ = std::fs::remove_dir_all(wide.parent().unwrap());
    let _ = std::fs::remove_dir_all(small.parent().unwrap());
}