use hugr_core::Hugr;
use hugr_core::builder::{
BuildError, BuildHandle, HugrBuilder, ModuleBuilder, dataflow::FunctionBuilder,
};
use hugr_core::ops::handle::FuncID;
#[expect(unused)]
pub(crate) mod test_quantum_extension {
use hugr_core::extension::Version;
use std::sync::{Arc, LazyLock};
use hugr_core::ops::{OpName, OpNameRef};
use hugr_core::std_extensions::arithmetic::float_ops;
use hugr_core::std_extensions::logic;
use hugr_core::types::FuncValueType;
use hugr_core::{
Extension,
extension::{
ExtensionId, ExtensionRegistry, PRELUDE,
prelude::{bool_t, qb_t},
},
ops::ExtensionOp,
std_extensions::arithmetic::float_types,
type_row,
types::{PolyFuncTypeRV, Signature},
};
fn one_qb_func() -> PolyFuncTypeRV {
FuncValueType::new_endo([qb_t()]).into()
}
fn two_qb_func() -> PolyFuncTypeRV {
FuncValueType::new_endo(vec![qb_t(), qb_t()]).into()
}
pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("test.quantum");
fn extension() -> Arc<Extension> {
Extension::new_arc(
EXTENSION_ID,
Version::new(0, 0, 0),
|extension, extension_ref| {
extension
.add_op(
OpName::new_inline("H"),
"Hadamard".into(),
one_qb_func(),
extension_ref,
)
.unwrap();
extension
.add_op(
OpName::new_inline("RzF64"),
"Rotation specified by float".into(),
Signature::new(vec![qb_t(), float_types::float64_type()], vec![qb_t()]),
extension_ref,
)
.unwrap();
extension
.add_op(
OpName::new_inline("CX"),
"CX".into(),
two_qb_func(),
extension_ref,
)
.unwrap();
extension
.add_op(
OpName::new_inline("Measure"),
"Measure a qubit, returning the qubit and the measurement result.".into(),
Signature::new(vec![qb_t()], vec![qb_t(), bool_t()]),
extension_ref,
)
.unwrap();
extension
.add_op(
OpName::new_inline("QAlloc"),
"Allocate a new qubit.".into(),
Signature::new(type_row![], vec![qb_t()]),
extension_ref,
)
.unwrap();
extension
.add_op(
OpName::new_inline("QDiscard"),
"Discard a qubit.".into(),
Signature::new(vec![qb_t()], type_row![]),
extension_ref,
)
.unwrap();
},
)
}
pub static EXTENSION: LazyLock<Arc<Extension>> = LazyLock::new(extension);
pub static REG: LazyLock<ExtensionRegistry> = LazyLock::new(|| {
ExtensionRegistry::new([
EXTENSION.clone(),
PRELUDE.clone(),
float_types::EXTENSION.clone(),
float_ops::EXTENSION.clone(),
logic::EXTENSION.clone(),
])
});
fn get_gate(gate_name: &OpNameRef) -> ExtensionOp {
EXTENSION.instantiate_extension_op(gate_name, []).unwrap()
}
pub(crate) fn h_gate() -> ExtensionOp {
get_gate("H")
}
pub(crate) fn cx_gate() -> ExtensionOp {
get_gate("CX")
}
pub(crate) fn measure() -> ExtensionOp {
get_gate("Measure")
}
pub(crate) fn rz_f64() -> ExtensionOp {
get_gate("RzF64")
}
pub(crate) fn q_alloc() -> ExtensionOp {
get_gate("QAlloc")
}
pub(crate) fn q_discard() -> ExtensionOp {
get_gate("QDiscard")
}
}
pub(crate) fn build_simple_hugr(
num_qubits: usize,
f: impl FnOnce(FunctionBuilder<&mut Hugr>) -> Result<BuildHandle<FuncID<true>>, BuildError>,
) -> Result<Hugr, BuildError> {
use hugr_core::{extension::prelude::qb_t, types::Signature};
let qb_row = vec![qb_t(); num_qubits];
let signature = Signature::new(qb_row.clone(), qb_row);
let mut module_builder = ModuleBuilder::new();
let f_builder = module_builder.define_function("main", signature)?;
f(f_builder)?;
Ok(module_builder.finish_hugr()?)
}