guppy_opt.rs

//! Tests optimizing Guppy-generated programs.

use rayon::iter::ParallelIterator;
use smol_str::SmolStr;
use std::collections::HashMap;
use std::fs;
use std::io::BufReader;
use std::path::Path;
use tket::extension::{TKET1_EXTENSION_ID, TKET_EXTENSION_ID};

use hugr::algorithms::ComposablePass;
use hugr::{Hugr, HugrView};
use rstest::{fixture, rstest};
use tket::passes::NormalizeGuppy;
use tket::serialize::pytket::{EncodeOptions, EncodedCircuit};
use tket::Circuit;

use tket1_passes::Tket1Circuit;
use tket_qsystem::QSystemPass;

const GUPPY_EXAMPLES_DIR: &str = "../test_files/guppy_optimization";

enum HugrFileType {
    Original,
    Flat,
    Optimized,
}

fn load_guppy_circuit(name: &str, file_type: HugrFileType) -> std::io::Result<Hugr> {
    let suffix = match file_type {
        HugrFileType::Original => "",
        HugrFileType::Flat => ".flat",
        HugrFileType::Optimized => ".opt",
    };
    let file = Path::new(GUPPY_EXAMPLES_DIR).join(format!("{name}/{name}{suffix}.hugr"));
    let reader = fs::File::open(file)?;
    let reader = BufReader::new(reader);
    Ok(Hugr::load(reader, None).unwrap())
}

#[fixture]
fn guppy_angles() -> Hugr {
    load_guppy_circuit("angles", HugrFileType::Original).unwrap()
}

#[fixture]
fn guppy_false_branch() -> Hugr {
    load_guppy_circuit("false_branch", HugrFileType::Original).unwrap()
}

#[fixture]
fn guppy_nested() -> Hugr {
    load_guppy_circuit("nested", HugrFileType::Original).unwrap()
}

#[fixture]
fn guppy_ranges() -> Hugr {
    load_guppy_circuit("ranges", HugrFileType::Original).unwrap()
}

#[fixture]
fn guppy_simple_cx() -> Hugr {
    load_guppy_circuit("simple_cx", HugrFileType::Original).unwrap()
}

fn run_pytket(h: &mut Hugr) {
    let circ = Circuit::new(h);
    let mut encoded =
        EncodedCircuit::new(&circ, EncodeOptions::new().with_subcircuits(true)).unwrap();

    encoded
        .par_iter_mut()
        .for_each(|(_region, serial_circuit)| {
            let mut circuit_ptr = Tket1Circuit::from_serial_circuit(serial_circuit).unwrap();
            circuit_ptr
                .clifford_simp(tket_json_rs::OpType::CX, true)
                .unwrap();
            *serial_circuit = circuit_ptr.to_serial_circuit().unwrap();
        });

    encoded.reassemble_inplace(circ.into_hugr(), None).unwrap();
}

fn count_gates(h: &impl HugrView) -> HashMap<SmolStr, usize> {
    let mut counts = HashMap::new();
    for n in h.nodes() {
        if let Some(eop) = h.get_optype(n).as_extension_op() {
            if [TKET_EXTENSION_ID, TKET1_EXTENSION_ID].contains(eop.extension_id()) {
                *counts.entry(eop.qualified_id()).or_default() += 1;
            }
        }
    }
    counts
}

/// Run some simple optimization passes on the guppy-generated HUGRs and validate the result.
///
/// This test is intended to check the current status of the Guppy optimization passes.
///

#[rstest]
#[should_panic = "xfail"]
#[case::angles("angles", Some(vec![
    ("tket.quantum.Rz", 2), ("tket.quantum.MeasureFree", 1), ("tket.quantum.H", 2), ("tket.quantum.QAlloc", 1)
]))]
#[should_panic = "xfail"]
#[case::simple_cx("simple_cx", Some(vec![
    ("tket.quantum.QAlloc", 2), ("tket.quantum.MeasureFree", 2),
]))]
#[should_panic = "xfail"]
#[case::nested("nested", Some(vec![
    ("tket.quantum.CZ", 6), ("tket.quantum.QAlloc", 3), ("tket.quantum.MeasureFree", 3), ("tket.quantum.H", 6)
]))]
#[should_panic = "xfail"]
#[case::ranges("ranges", Some(vec![
    ("tket.quantum.H", 8), ("tket.quantum.MeasureFree", 4), ("tket.quantum.QAlloc", 4), ("tket.quantum.CX", 6)
]))]
#[should_panic = "xfail"]
#[case::false_branch("false_branch", Some(vec![
    ("TKET1.tk1op", 2), ("tket.quantum.QAlloc", 1), ("tket.quantum.MeasureFree", 1)
]))]
#[cfg_attr(miri, ignore)] // Opening files is not supported in (isolated) miri
fn optimise_guppy_pytket(#[case] name: &str, #[case] xfail: Option<Vec<(&str, usize)>>) {
    let mut hugr = load_guppy_circuit(name, HugrFileType::Flat)
        .unwrap_or_else(|_| load_guppy_circuit(name, HugrFileType::Original).unwrap());
    run_pytket(&mut hugr);
    let should_xfail = xfail.is_some();
    let expected_counts = match xfail {
        Some(counts) => counts.into_iter().map(|(k, v)| (k.into(), v)).collect(),
        None => count_gates(&load_guppy_circuit(name, HugrFileType::Optimized).unwrap()),
    };
    assert_eq!(count_gates(&hugr), expected_counts);
    if should_xfail {
        panic!("xfail");
    }
}

#[rstest]
#[case::angles("angles")]
#[should_panic]
#[case::nested("nested")]
#[should_panic]
#[case::ranges("ranges")]
#[cfg_attr(miri, ignore)] // Opening files is not supported in (isolated) miri
fn flatten_guppy(#[case] name: &str) {
    let mut hugr = load_guppy_circuit(name, HugrFileType::Original).unwrap();
    NormalizeGuppy::default().run(&mut hugr).unwrap();
    let target = load_guppy_circuit(name, HugrFileType::Flat).unwrap();
    assert_eq!(count_gates(&hugr), count_gates(&target));
}

/// Check that each example optimizes to the full extent given by the .opt (and .flat) .hugr files.
#[rstest]
#[case::angles("angles")]
#[case::false_branch("false_branch")]
#[case::simple_cx("simple_cx")]
#[case::nested("nested")]
#[case::ranges("ranges")]
#[should_panic] // This does not yet pass for any case!
fn optimise_guppy(#[case] name: &str) {
    let mut hugr = load_guppy_circuit(name, HugrFileType::Original).unwrap();
    let flat = count_gates(
        load_guppy_circuit(name, HugrFileType::Flat)
            .ok()
            .as_ref()
            .unwrap_or(&hugr),
    );
    let opt = count_gates(&load_guppy_circuit(name, HugrFileType::Optimized).unwrap());

    NormalizeGuppy::default().run(&mut hugr).unwrap();
    assert_eq!(count_gates(&hugr), flat);

    run_pytket(&mut hugr);

    assert_eq!(count_gates(&hugr), opt);

    // Lower to QSystem. This may blow up the HUGR size.
    QSystemPass::default().run(&mut hugr).unwrap();

    hugr.validate().unwrap_or_else(|e| panic!("{e}"));
}