1use super::{ControlBody, ForOp, IfOp, SwitchCase, SwitchOp, WhileOp};
14use crate::circuit::{ClassicalValue, ClassicalVar, Qubit};
15use std::collections::BTreeSet;
16use std::fmt;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
20pub enum ClassicalControlKind {
21 If,
22 While,
23 For,
24 Switch,
25 Break,
26 Continue,
27}
28
29impl ClassicalControlKind {
30 pub const fn name(self) -> &'static str {
32 match self {
33 Self::If => "if",
34 Self::While => "while",
35 Self::For => "for",
36 Self::Switch => "switch",
37 Self::Break => "break",
38 Self::Continue => "continue",
39 }
40 }
41}
42
43#[derive(Debug, Clone, PartialEq)]
45pub enum ClassicalControlOp {
46 If(IfOp),
48 While(WhileOp),
50 For(ForOp),
52 Switch(SwitchOp),
54 Break,
56 Continue,
58}
59
60impl ClassicalControlOp {
61 pub(crate) fn bodies(&self) -> impl Iterator<Item = &ControlBody> + '_ {
68 let (first, cases, last): (Option<&ControlBody>, &[SwitchCase], Option<&ControlBody>) =
69 match self {
70 Self::If(op) => (Some(op.then_body()), &[], op.else_body()),
71 Self::While(op) => (Some(op.body()), &[], None),
72 Self::For(op) => (Some(op.body()), &[], None),
73 Self::Switch(op) => (None, op.cases(), op.default()),
74 Self::Break | Self::Continue => (None, &[], None),
75 };
76 first
77 .into_iter()
78 .chain(cases.iter().map(SwitchCase::body))
79 .chain(last)
80 }
81
82 pub const fn kind(&self) -> ClassicalControlKind {
84 match self {
85 Self::If(_) => ClassicalControlKind::If,
86 Self::While(_) => ClassicalControlKind::While,
87 Self::For(_) => ClassicalControlKind::For,
88 Self::Switch(_) => ClassicalControlKind::Switch,
89 Self::Break => ClassicalControlKind::Break,
90 Self::Continue => ClassicalControlKind::Continue,
91 }
92 }
93
94 pub const fn name(&self) -> &'static str {
96 self.kind().name()
97 }
98
99 pub fn classical_var_reads(&self) -> BTreeSet<ClassicalVar> {
101 match self {
102 Self::If(op) => op.classical_var_reads(),
103 Self::While(op) => op.classical_var_reads(),
104 Self::For(op) => op.classical_var_reads(),
105 Self::Switch(op) => op.classical_var_reads(),
106 Self::Break | Self::Continue => BTreeSet::new(),
107 }
108 }
109
110 pub fn classical_value_reads(&self) -> BTreeSet<ClassicalValue> {
112 match self {
113 Self::If(op) => op.classical_value_reads(),
114 Self::While(op) => op.classical_value_reads(),
115 Self::For(op) => op.classical_value_reads(),
116 Self::Switch(op) => op.classical_value_reads(),
117 Self::Break | Self::Continue => BTreeSet::new(),
118 }
119 }
120
121 pub fn classical_writes(&self) -> BTreeSet<ClassicalVar> {
123 match self {
124 Self::For(op) => op.classical_writes(),
125 Self::If(_) | Self::While(_) | Self::Switch(_) | Self::Break | Self::Continue => {
126 BTreeSet::new()
127 }
128 }
129 }
130
131 pub fn used_qubits(&self) -> BTreeSet<Qubit> {
133 match self {
134 Self::If(op) => op.used_qubits(),
135 Self::While(op) => op.used_qubits(),
136 Self::For(op) => op.used_qubits(),
137 Self::Switch(op) => op.used_qubits(),
138 Self::Break | Self::Continue => BTreeSet::new(),
139 }
140 }
141
142 pub fn has_measurement(&self) -> bool {
144 self.bodies().any(ControlBody::has_measurement)
145 }
146
147 pub fn reads_value(&self, value: ClassicalValue) -> bool {
150 if self.classical_value_reads().contains(&value) {
151 return true;
152 }
153
154 self.bodies().any(|body| body.reads_value(value))
155 }
156}
157
158impl fmt::Display for ClassicalControlOp {
159 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
160 f.write_str(self.name())
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use super::ClassicalControlOp;
167 use crate::circuit::{
168 Circuit, CircuitId, ClassicalExpr, ClassicalType, ClassicalValue, ClassicalVar,
169 ControlBody, ForOp, IfOp, Instruction, Qubit, SwitchCase, SwitchOp, WhileOp,
170 };
171
172 fn assert_bodies(control: &ClassicalControlOp, expected: &[&ControlBody]) {
173 let actual = control.bodies().collect::<Vec<_>>();
174 assert_eq!(actual.len(), expected.len());
175 for (actual, expected) in actual.into_iter().zip(expected) {
176 assert!(std::ptr::eq(actual, *expected));
177 }
178 }
179
180 #[test]
181 fn bodies_preserve_optional_empty_branches_without_recursing() {
182 let mut nested = Circuit::new(1);
183 nested
184 .while_(ClassicalExpr::bool_literal(false), |body| {
185 body.x(Qubit::new(0))
186 })
187 .unwrap();
188 for else_body in [None, Some(ControlBody::new(vec![]))] {
189 let control = ClassicalControlOp::If(
190 IfOp::new(
191 ClassicalExpr::bool_literal(false),
192 ControlBody::new(nested.operations().to_vec()),
193 else_body,
194 )
195 .unwrap(),
196 );
197 let ClassicalControlOp::If(op) = &control else {
198 unreachable!()
199 };
200 if let Some(else_body) = op.else_body() {
201 assert_bodies(&control, &[op.then_body(), else_body]);
202 } else {
203 assert_bodies(&control, &[op.then_body()]);
204 }
205 assert!(matches!(
206 control.bodies().next().unwrap().operations()[0].instruction,
207 Instruction::ClassicalControl(ClassicalControlOp::While(_))
208 ));
209 }
210 }
211
212 #[test]
213 fn bodies_preserve_switch_source_order_and_optional_default() {
214 let empty = ControlBody::new(vec![]);
215 for cases in [
216 vec![],
217 vec![
218 SwitchCase::new(2, empty.clone()),
219 SwitchCase::new(0, empty.clone()),
220 ],
221 ] {
222 for default in [None, Some(empty.clone())] {
223 let control = ClassicalControlOp::Switch(
224 SwitchOp::new(
225 ClassicalExpr::uint_literal(2, 0).unwrap(),
226 cases.clone(),
227 default,
228 )
229 .unwrap(),
230 );
231 let ClassicalControlOp::Switch(op) = &control else {
232 unreachable!()
233 };
234 let mut expected = Vec::new();
235 if !op.cases().is_empty() {
236 expected.push(op.cases()[0].body());
238 expected.push(op.cases()[1].body());
239 }
240 if let Some(default) = op.default() {
241 expected.push(default);
242 }
243 assert_bodies(&control, &expected);
244 }
245 }
246 }
247
248 #[test]
249 fn bodies_enumerate_loop_structure_independently_of_iteration_count() {
250 let var = ClassicalVar::new(CircuitId::new(), 0, ClassicalType::uint(8).unwrap());
251 for stop in [0, 3] {
252 let control = ClassicalControlOp::For(
253 ForOp::new(
254 var,
255 ClassicalExpr::uint_literal(8, 0).unwrap(),
256 ClassicalExpr::uint_literal(8, stop).unwrap(),
257 ClassicalExpr::uint_literal(8, 1).unwrap(),
258 ControlBody::new(vec![]),
259 )
260 .unwrap(),
261 );
262 let ClassicalControlOp::For(op) = &control else {
263 unreachable!()
264 };
265 assert_bodies(&control, &[op.body()]);
266 }
267 let control = ClassicalControlOp::While(
268 WhileOp::new(ClassicalExpr::bool_literal(false), ControlBody::new(vec![])).unwrap(),
269 );
270 let ClassicalControlOp::While(op) = &control else {
271 unreachable!()
272 };
273 assert_bodies(&control, &[op.body()]);
274 assert_bodies(&ClassicalControlOp::Break, &[]);
275 assert_bodies(&ClassicalControlOp::Continue, &[]);
276 }
277
278 #[test]
279 fn break_and_continue_have_no_resource_dependencies() {
280 assert!(ClassicalControlOp::Break.classical_var_reads().is_empty());
281 assert!(ClassicalControlOp::Break.classical_value_reads().is_empty());
282 assert!(ClassicalControlOp::Break.classical_writes().is_empty());
283 assert!(ClassicalControlOp::Break.used_qubits().is_empty());
284 assert!(
285 ClassicalControlOp::Continue
286 .classical_var_reads()
287 .is_empty()
288 );
289 assert!(
290 ClassicalControlOp::Continue
291 .classical_value_reads()
292 .is_empty()
293 );
294 assert!(ClassicalControlOp::Continue.classical_writes().is_empty());
295 assert!(ClassicalControlOp::Continue.used_qubits().is_empty());
296 }
297
298 #[test]
299 fn control_op_forwards_classical_var_and_value_reads() {
300 let circuit_id = CircuitId::new();
301 let bit = ClassicalVar::new(circuit_id, 1, ClassicalType::Bit);
302 let value = ClassicalValue::new(circuit_id, 2, ClassicalType::Bit);
303 let condition = ClassicalExpr::try_and(
304 ClassicalExpr::bit_to_bool(bit.expr()).unwrap(),
305 ClassicalExpr::bit_to_bool(value.expr()).unwrap(),
306 )
307 .unwrap();
308 let op = IfOp::new(condition, ControlBody::new(vec![]), None).unwrap();
309 let op = ClassicalControlOp::If(op);
310
311 assert!(op.classical_var_reads().contains(&bit));
312 assert!(op.classical_value_reads().contains(&value));
313 }
314
315 #[test]
316 fn display_reports_operation_kind() {
317 let condition = ClassicalExpr::bool_literal(true);
318 let body = ControlBody::new(vec![]);
319 let if_op = IfOp::new(condition.clone(), body.clone(), None).unwrap();
320 let while_op = WhileOp::new(condition, body.clone()).unwrap();
321 let loop_var = ClassicalVar::new(CircuitId::new(), 0, ClassicalType::uint(8).unwrap());
322 let for_op = ForOp::new(
323 loop_var,
324 ClassicalExpr::uint_literal(8, 0).unwrap(),
325 ClassicalExpr::uint_literal(8, 2).unwrap(),
326 ClassicalExpr::uint_literal(8, 1).unwrap(),
327 body.clone(),
328 )
329 .unwrap();
330 let switch_op = SwitchOp::new(
331 ClassicalExpr::uint_literal(2, 1).unwrap(),
332 vec![SwitchCase::new(1, body)],
333 None,
334 )
335 .unwrap();
336
337 assert_eq!(ClassicalControlOp::If(if_op).to_string(), "if");
338 assert_eq!(ClassicalControlOp::While(while_op).to_string(), "while");
339 assert_eq!(ClassicalControlOp::For(for_op).to_string(), "for");
340 assert_eq!(ClassicalControlOp::Switch(switch_op).to_string(), "switch");
341 assert_eq!(ClassicalControlOp::Break.to_string(), "break");
342 assert_eq!(ClassicalControlOp::Continue.to_string(), "continue");
343 }
344}