use super::ControlBody;
use crate::circuit::classical_expr::ClassicalExpr;
use crate::circuit::{CircuitError, ClassicalType, ClassicalValue, ClassicalVar, Qubit};
use std::collections::BTreeSet;
#[derive(Debug, Clone, PartialEq)]
pub struct IfOp {
condition: ClassicalExpr,
then_body: ControlBody,
else_body: Option<ControlBody>,
}
impl IfOp {
pub fn new(
condition: ClassicalExpr,
then_body: ControlBody,
else_body: Option<ControlBody>,
) -> Result<Self, CircuitError> {
if condition.ty() != ClassicalType::Bool {
return Err(CircuitError::InvalidOperation(format!(
"if condition must be Bool, got {:?}",
condition.ty()
)));
}
Ok(Self {
condition,
then_body,
else_body,
})
}
pub fn condition(&self) -> &ClassicalExpr {
&self.condition
}
pub fn then_body(&self) -> &ControlBody {
&self.then_body
}
pub fn else_body(&self) -> Option<&ControlBody> {
self.else_body.as_ref()
}
pub fn classical_var_reads(&self) -> BTreeSet<ClassicalVar> {
self.condition.vars()
}
pub fn classical_value_reads(&self) -> BTreeSet<ClassicalValue> {
self.condition.values()
}
pub fn used_qubits(&self) -> BTreeSet<Qubit> {
let mut qubits = self.then_body.used_qubits();
if let Some(else_body) = &self.else_body {
qubits.extend(else_body.used_qubits());
}
qubits
}
}
#[cfg(test)]
mod tests {
use super::IfOp;
use crate::circuit::{
CircuitId, ClassicalExpr, ClassicalType, ClassicalValue, ClassicalVar, ControlBody,
};
#[test]
fn if_requires_bool_condition() {
let condition = ClassicalExpr::bool_literal(true);
let op = IfOp::new(condition, ControlBody::new(vec![]), None).unwrap();
assert!(op.else_body().is_none());
assert!(
IfOp::new(
ClassicalExpr::bit_literal(true),
ControlBody::new(vec![]),
None
)
.is_err()
);
}
#[test]
fn if_reports_condition_var_reads() {
let circuit_id = CircuitId::new();
let bit = ClassicalExpr::var(ClassicalVar::new(circuit_id, 3, ClassicalType::Bit));
let condition = ClassicalExpr::bit_to_bool(bit).unwrap();
let op = IfOp::new(condition, ControlBody::new(vec![]), None).unwrap();
assert!(op.classical_var_reads().contains(&ClassicalVar::new(
circuit_id,
3,
ClassicalType::Bit
)));
}
#[test]
fn if_reports_condition_value_reads() {
let value = ClassicalValue::new(CircuitId::new(), 4, ClassicalType::Bit);
let condition = ClassicalExpr::bit_to_bool(value.expr()).unwrap();
let op = IfOp::new(condition, ControlBody::new(vec![]), None).unwrap();
assert!(op.classical_var_reads().is_empty());
assert!(op.classical_value_reads().contains(&value));
}
}