Skip to main content

prism_q/circuit/dynamic/
builder.rs

1//! Structured construction of a [`DynamicProgram`].
2
3use super::{
4    Action, BasicBlock, BlockId, ClassicalExpr, ClassicalType, ClassicalValue, DynamicProgram,
5    RotationKind, Terminator, VarId, Variable,
6};
7use crate::circuit::{Circuit, Instruction, SmallVec};
8use crate::error::{PrismError, Result};
9use crate::gates::Gate;
10
11struct PendingBlock {
12    instructions: Vec<Instruction>,
13    actions: Vec<Action>,
14    terminator: Option<Terminator>,
15    loop_id: Option<usize>,
16}
17
18enum Frame {
19    While {
20        loop_id: usize,
21        header: usize,
22        exit: usize,
23    },
24    If {
25        branch: usize,
26        join: usize,
27        in_else: bool,
28    },
29}
30
31/// Builds a [`DynamicProgram`] from structured control flow.
32///
33/// Instructions append to the block being built. [`begin_while`] and
34/// [`begin_if`] open a construct that [`end`] closes; [`begin_else`] switches
35/// the innermost `if` to its other arm, or reopens an `if` that [`end`] has
36/// just closed. Every index is checked when [`build`] validates the program,
37/// so no method here panics on a bad one.
38///
39/// # Examples
40///
41/// Repeat until success: prepare `|+>`, measure, and retry while the outcome
42/// is 1.
43///
44/// ```
45/// use prism_q::circuit::dynamic::{ClassicalExpr, DynamicProgramBuilder};
46/// use prism_q::{Gate, simulate_program};
47///
48/// let mut b = DynamicProgramBuilder::new(1, 1);
49/// b.add_gate(Gate::H, &[0]).add_measure(0, 0);
50/// b.begin_while("retry", ClassicalExpr::Bit(0));
51/// b.add_reset(0).add_gate(Gate::H, &[0]).add_measure(0, 0);
52/// b.end()?;
53/// let program = b.build()?;
54///
55/// let shots = simulate_program(&program).seed(42).shots(100)?;
56/// assert!(shots.shots.iter().all(|shot| !shot[0]));
57/// # Ok::<(), prism_q::PrismError>(())
58/// ```
59///
60/// [`begin_while`]: Self::begin_while
61/// [`begin_if`]: Self::begin_if
62/// [`begin_else`]: Self::begin_else
63/// [`end`]: Self::end
64/// [`build`]: Self::build
65pub struct DynamicProgramBuilder {
66    num_qubits: usize,
67    num_classical_bits: usize,
68    variables: Vec<Variable>,
69    blocks: Vec<PendingBlock>,
70    current: usize,
71    frames: Vec<Frame>,
72    loop_names: Vec<String>,
73    /// The `if` frame [`Self::end`] just closed, which [`Self::begin_else`] may
74    /// reopen until anything else is appended.
75    reopenable: Option<Frame>,
76}
77
78impl DynamicProgramBuilder {
79    pub fn new(num_qubits: usize, num_classical_bits: usize) -> Self {
80        Self {
81            num_qubits,
82            num_classical_bits,
83            variables: Vec::new(),
84            blocks: vec![PendingBlock {
85                instructions: Vec::new(),
86                actions: Vec::new(),
87                terminator: None,
88                loop_id: None,
89            }],
90            current: 0,
91            frames: Vec::new(),
92            loop_names: Vec::new(),
93            reopenable: None,
94        }
95    }
96
97    pub fn num_qubits(&self) -> usize {
98        self.num_qubits
99    }
100
101    pub fn num_classical_bits(&self) -> usize {
102        self.num_classical_bits
103    }
104
105    pub(crate) fn resize(&mut self, num_qubits: usize, num_classical_bits: usize) {
106        self.num_qubits = num_qubits;
107        self.num_classical_bits = num_classical_bits;
108    }
109
110    /// Declare a variable every shot starts at `initial`.
111    pub fn declare(
112        &mut self,
113        name: impl Into<String>,
114        ty: ClassicalType,
115        initial: ClassicalValue,
116    ) -> VarId {
117        self.variables.push(Variable {
118            name: name.into(),
119            ty,
120            initial,
121        });
122        VarId::new(self.variables.len() - 1)
123    }
124
125    /// The variable declared under `name`, the latest one when several share it.
126    pub fn variable(&self, name: &str) -> Option<VarId> {
127        self.variables
128            .iter()
129            .rposition(|variable| variable.name == name)
130            .map(VarId::new)
131    }
132
133    /// Parse an OpenQASM expression over this builder's variables, by name, and
134    /// its classical bits, read as `c[i]` for one bit and `c` for all of them as
135    /// an unsigned integer.
136    ///
137    /// # Errors
138    ///
139    /// [`PrismError::Parse`] for text that is not an expression, an unknown
140    /// name, or a builtin function over a runtime value.
141    pub fn expr(&self, text: &str) -> Result<ClassicalExpr> {
142        crate::circuit::openqasm::builder_expr(text, self)
143    }
144
145    pub fn add_instruction(&mut self, instruction: Instruction) -> &mut Self {
146        let at = self.circuit_block();
147        self.blocks[at].instructions.push(instruction);
148        self
149    }
150
151    pub fn add_gate(&mut self, gate: Gate, targets: &[usize]) -> &mut Self {
152        self.add_instruction(Instruction::Gate {
153            gate,
154            targets: SmallVec::from_slice(targets),
155        })
156    }
157
158    pub fn add_measure(&mut self, qubit: usize, classical_bit: usize) -> &mut Self {
159        self.add_instruction(Instruction::Measure {
160            qubit,
161            classical_bit,
162        })
163    }
164
165    pub fn add_reset(&mut self, qubit: usize) -> &mut Self {
166        self.add_instruction(Instruction::Reset { qubit })
167    }
168
169    /// Append every instruction of `circuit`.
170    pub fn append(&mut self, circuit: &Circuit) -> &mut Self {
171        let at = self.circuit_block();
172        self.blocks[at]
173            .instructions
174            .extend(circuit.instructions.iter().cloned());
175        self
176    }
177
178    /// Apply `kind` to `targets` at the angle `angle` evaluates to when the
179    /// block runs.
180    pub fn add_rotation(
181        &mut self,
182        kind: RotationKind,
183        targets: &[usize],
184        angle: ClassicalExpr,
185    ) -> &mut Self {
186        self.push_action(Action::Rotation {
187            kind,
188            targets: SmallVec::from_slice(targets),
189            angle,
190        })
191    }
192
193    /// Store `value` into `var`, converted to its declared type.
194    pub fn assign(&mut self, var: VarId, value: ClassicalExpr) -> &mut Self {
195        self.push_action(Action::Assign { var, value })
196    }
197
198    /// Open a loop that runs its body while `condition` is truthy, testing it
199    /// before each pass. `name` labels the loop in the step-limit error.
200    pub fn begin_while(&mut self, name: impl Into<String>, condition: ClassicalExpr) {
201        self.reopenable = None;
202        let enclosing = self.blocks[self.current].loop_id;
203        let loop_id = self.loop_names.len();
204        self.loop_names.push(name.into());
205        let header = self.new_block(Some(loop_id));
206        let body = self.new_block(Some(loop_id));
207        let exit = self.new_block(enclosing);
208        self.terminate(Terminator::Jump(BlockId::new(header)));
209        self.blocks[header].terminator = Some(Terminator::Branch {
210            condition,
211            then: BlockId::new(body),
212            otherwise: BlockId::new(exit),
213        });
214        self.frames.push(Frame::While {
215            loop_id,
216            header,
217            exit,
218        });
219        self.current = body;
220    }
221
222    /// Open a branch whose body runs when `condition` is truthy.
223    pub fn begin_if(&mut self, condition: ClassicalExpr) {
224        self.reopenable = None;
225        let loop_id = self.blocks[self.current].loop_id;
226        let then = self.new_block(loop_id);
227        let join = self.new_block(loop_id);
228        self.terminate(Terminator::Branch {
229            condition,
230            then: BlockId::new(then),
231            otherwise: BlockId::new(join),
232        });
233        self.frames.push(Frame::If {
234            branch: self.current,
235            join,
236            in_else: false,
237        });
238        self.current = then;
239    }
240
241    /// Switch the innermost open `if` to its `else` arm, or reopen the `if`
242    /// that [`Self::end`] closed last when nothing has been appended since.
243    ///
244    /// # Errors
245    ///
246    /// [`PrismError::InvalidParameter`] when there is no such `if`, or it
247    /// already has an `else`.
248    pub fn begin_else(&mut self) -> Result<()> {
249        let open = matches!(self.frames.last(), Some(Frame::If { in_else: false, .. }));
250        let (branch, join) = if open {
251            let Some(Frame::If { branch, join, .. }) = self.frames.pop() else {
252                unreachable!("checked above")
253            };
254            self.terminate(Terminator::Jump(BlockId::new(join)));
255            (branch, join)
256        } else {
257            match self.reopenable.take() {
258                Some(Frame::If {
259                    branch,
260                    join,
261                    in_else: false,
262                }) => (branch, join),
263                _ => return Err(misuse("`begin_else` with no `if` to attach to")),
264            }
265        };
266        let other = self.new_block(self.blocks[join].loop_id);
267        if let Some(Terminator::Branch { otherwise, .. }) = &mut self.blocks[branch].terminator {
268            *otherwise = BlockId::new(other);
269        }
270        self.frames.push(Frame::If {
271            branch,
272            join,
273            in_else: true,
274        });
275        self.current = other;
276        Ok(())
277    }
278
279    /// Close the innermost `while` or `if`.
280    ///
281    /// # Errors
282    ///
283    /// [`PrismError::InvalidParameter`] when nothing is open.
284    pub fn end(&mut self) -> Result<()> {
285        let frame = self
286            .frames
287            .pop()
288            .ok_or_else(|| misuse("`end` with no open `while` or `if`"))?;
289        match frame {
290            Frame::While { header, exit, .. } => {
291                self.terminate(Terminator::Jump(BlockId::new(header)));
292                self.current = exit;
293                self.reopenable = None;
294            }
295            Frame::If {
296                branch,
297                join,
298                in_else,
299            } => {
300                self.terminate(Terminator::Jump(BlockId::new(join)));
301                self.current = join;
302                self.reopenable = Some(Frame::If {
303                    branch,
304                    join,
305                    in_else,
306                });
307            }
308        }
309        Ok(())
310    }
311
312    /// Leave the innermost `while`.
313    ///
314    /// # Errors
315    ///
316    /// [`PrismError::InvalidParameter`] outside a loop.
317    pub fn break_loop(&mut self) -> Result<()> {
318        let exit = self
319            .innermost_loop()
320            .map(|(_, exit)| exit)
321            .ok_or_else(|| misuse("`break` outside a `while` loop"))?;
322        self.jump_away(exit);
323        Ok(())
324    }
325
326    /// Skip to the next test of the innermost `while` condition.
327    ///
328    /// # Errors
329    ///
330    /// [`PrismError::InvalidParameter`] outside a loop.
331    pub fn continue_loop(&mut self) -> Result<()> {
332        let header = self
333            .innermost_loop()
334            .map(|(header, _)| header)
335            .ok_or_else(|| misuse("`continue` outside a `while` loop"))?;
336        self.jump_away(header);
337        Ok(())
338    }
339
340    /// Finish the program, dropping unreachable blocks and merging each
341    /// straight-line run of blocks into one, so fusion sees the longest
342    /// circuits the control flow allows.
343    ///
344    /// # Errors
345    ///
346    /// [`PrismError::InvalidParameter`] while a `while` or `if` is still open,
347    /// and anything [`DynamicProgram::new`] rejects.
348    pub fn build(mut self) -> Result<DynamicProgram> {
349        if let Some(frame) = self.frames.last() {
350            let what = match frame {
351                Frame::While { loop_id, .. } => {
352                    format!("`while` loop `{}`", self.loop_names[*loop_id])
353                }
354                Frame::If { .. } => "`if`".to_string(),
355            };
356            return Err(misuse(format!("{what} is still open at `build`")));
357        }
358        for block in &mut self.blocks {
359            block.terminator.get_or_insert(Terminator::End);
360        }
361        let blocks = compact(self.blocks);
362        let blocks = blocks
363            .into_iter()
364            .map(|block| BasicBlock {
365                circuit: Circuit {
366                    num_qubits: self.num_qubits,
367                    num_classical_bits: self.num_classical_bits,
368                    instructions: block.instructions,
369                },
370                actions: block.actions,
371                terminator: block.terminator.expect("every block terminated above"),
372                loop_name: block.loop_id.map(|id| self.loop_names[id].clone()),
373            })
374            .collect();
375        DynamicProgram::new(
376            self.num_qubits,
377            self.num_classical_bits,
378            self.variables,
379            blocks,
380        )
381    }
382
383    fn new_block(&mut self, loop_id: Option<usize>) -> usize {
384        self.blocks.push(PendingBlock {
385            instructions: Vec::new(),
386            actions: Vec::new(),
387            terminator: None,
388            loop_id,
389        });
390        self.blocks.len() - 1
391    }
392
393    fn terminate(&mut self, terminator: Terminator) {
394        self.blocks[self.current]
395            .terminator
396            .get_or_insert(terminator);
397    }
398
399    /// The block instructions append to: the current one, unless it already
400    /// holds an action, which a circuit may not follow within one block.
401    fn circuit_block(&mut self) -> usize {
402        self.reopenable = None;
403        if !self.blocks[self.current].actions.is_empty() {
404            let next = self.new_block(self.blocks[self.current].loop_id);
405            self.terminate(Terminator::Jump(BlockId::new(next)));
406            self.current = next;
407        }
408        self.current
409    }
410
411    fn push_action(&mut self, action: Action) -> &mut Self {
412        self.reopenable = None;
413        self.blocks[self.current].actions.push(action);
414        self
415    }
416
417    fn innermost_loop(&self) -> Option<(usize, usize)> {
418        self.frames.iter().rev().find_map(|frame| match frame {
419            Frame::While { header, exit, .. } => Some((*header, *exit)),
420            Frame::If { .. } => None,
421        })
422    }
423
424    /// End the current block with a jump and continue in a fresh one, which
425    /// nothing reaches until a later construct joins into it.
426    fn jump_away(&mut self, target: usize) {
427        self.reopenable = None;
428        self.terminate(Terminator::Jump(BlockId::new(target)));
429        self.current = self.new_block(self.blocks[self.current].loop_id);
430    }
431}
432
433fn misuse(message: impl Into<String>) -> PrismError {
434    PrismError::InvalidParameter {
435        message: message.into(),
436    }
437}
438
439/// Thread jumps through empty blocks, merge every block into the one block
440/// that jumps to it alone, then drop what the entry cannot reach and renumber
441/// in reachable order.
442fn compact(mut blocks: Vec<PendingBlock>) -> Vec<PendingBlock> {
443    let forward = |blocks: &[PendingBlock], start: usize| {
444        let mut at = start;
445        for _ in 0..blocks.len() {
446            let block = &blocks[at];
447            match block.terminator {
448                Some(Terminator::Jump(next))
449                    if block.instructions.is_empty()
450                        && block.actions.is_empty()
451                        && next.index() != at =>
452                {
453                    at = next.index();
454                }
455                _ => break,
456            }
457        }
458        at
459    };
460    let threaded: Vec<usize> = (0..blocks.len()).map(|at| forward(&blocks, at)).collect();
461    for block in &mut blocks {
462        retarget(block, |id| BlockId::new(threaded[id.index()]));
463    }
464    let entry = threaded[0];
465
466    let mut predecessors = vec![0usize; blocks.len()];
467    for block in &blocks {
468        for target in successors(block) {
469            predecessors[target] += 1;
470        }
471    }
472    for at in 0..blocks.len() {
473        while let Some(Terminator::Jump(target)) = blocks[at].terminator {
474            let from = target.index();
475            if !blocks[at].actions.is_empty()
476                || from == at
477                || from == entry
478                || predecessors[from] != 1
479            {
480                break;
481            }
482            let absorbed = std::mem::replace(
483                &mut blocks[from],
484                PendingBlock {
485                    instructions: Vec::new(),
486                    actions: Vec::new(),
487                    terminator: Some(Terminator::End),
488                    loop_id: None,
489                },
490            );
491            let block = &mut blocks[at];
492            block.instructions.extend(absorbed.instructions);
493            block.actions = absorbed.actions;
494            block.terminator = absorbed.terminator;
495        }
496    }
497
498    let mut order = Vec::with_capacity(blocks.len());
499    let mut renumber = vec![usize::MAX; blocks.len()];
500    let mut stack = vec![entry];
501    while let Some(at) = stack.pop() {
502        if renumber[at] != usize::MAX {
503            continue;
504        }
505        renumber[at] = order.len();
506        order.push(at);
507        for target in successors(&blocks[at]).into_iter().rev() {
508            if renumber[target] == usize::MAX {
509                stack.push(target);
510            }
511        }
512    }
513    let mut slots: Vec<Option<PendingBlock>> = blocks.into_iter().map(Some).collect();
514    order
515        .iter()
516        .map(|&at| {
517            let mut block = slots[at].take().expect("each block visited once");
518            retarget(&mut block, |id| BlockId::new(renumber[id.index()]));
519            block
520        })
521        .collect()
522}
523
524fn successors(block: &PendingBlock) -> SmallVec<[usize; 2]> {
525    match block.terminator.as_ref().expect("every block terminated") {
526        Terminator::Jump(target) => SmallVec::from_slice(&[target.index()]),
527        Terminator::Branch {
528            then, otherwise, ..
529        } => SmallVec::from_slice(&[then.index(), otherwise.index()]),
530        Terminator::End => SmallVec::new(),
531    }
532}
533
534fn retarget(block: &mut PendingBlock, map: impl Fn(BlockId) -> BlockId) {
535    match block.terminator.as_mut().expect("every block terminated") {
536        Terminator::Jump(target) => *target = map(*target),
537        Terminator::Branch {
538            then, otherwise, ..
539        } => {
540            *then = map(*then);
541            *otherwise = map(*otherwise);
542        }
543        Terminator::End => {}
544    }
545}