feat: Add Program.wrap_in_loop() method

quil-py/quil/program/__init__.pyi

@@ -12,8 +12,10 @@ from quil.instructions import (
     Instruction,
     MeasureCalibrationDefinition,
     Measurement,
+    MemoryReference,
     Qubit,
     Sharing,
+    Target,
     Vector,
     Waveform,
     TargetPlaceholder,
@@ -120,6 +122,27 @@ class Program:
         Convert the instruction to a Quil string. If any part of the instruction can't
         be converted to valid Quil, it will be printed in a human-readable debug format.
         """
+    def wrap_in_loop(
+        self, loop_count_reference: MemoryReference, start_target: Target, end_target: Target, iterations: int
+    ) -> "Program":
+        """
+        Return a copy of the `Program` wrapped in a loop that repeats ``iterations`` times.
+
+        The loop is constructed by wrapping the body of the program in classical Quil instructions.
+        The given ``loop_count_reference`` must refer to an INTEGER memory region. The value at the
+        reference given will be set to ``iterations`` and decremented in the loop. The loop will
+        terminate when the reference reaches 0. For this reason your program should not itself
+        modify the value at the reference unless you intend to modify the remaining number of
+        iterations (i.e. to break the loop).
+
+        The given ``start_target`` and ``end_target`` will be used as the entry and exit points for
+        the loop, respectively. You should provide unique `quil.instructions.Target`s that won't be
+        used elsewhere in the program.
+
+        If `iterations` is 0, then a copy of the program is returned without any changes. Raises a
+        `TypeError` if `iterations` is negative.
+        """
+        ...
     def resolve_placeholders(self) -> None:
         """
         Resolve ``TargetPlaceholder``s and ``QubitPlaceholder``s within the program using default resolvers.

quil-py/src/program/mod.rs

@@ -24,7 +24,10 @@ use rigetti_pyo3::{
 
 use crate::{
     impl_eq, impl_to_quil,
-    instruction::{PyDeclaration, PyGateDefinition, PyInstruction, PyQubit, PyWaveform},
+    instruction::{
+        PyDeclaration, PyGateDefinition, PyInstruction, PyMemoryReference, PyQubit, PyTarget,
+        PyWaveform,
+    },
 };
 
 pub use self::{calibration::PyCalibrationSet, frame::PyFrameSet, memory::PyMemoryRegion};
@@ -238,6 +241,21 @@ impl PyProgram {
         self.as_inner_mut().resolve_placeholders();
     }
 
+    pub fn wrap_in_loop(
+        &self,
+        loop_count_reference: PyMemoryReference,
+        start_target: PyTarget,
+        end_target: PyTarget,
+        iterations: u32,
+    ) -> Self {
+        PyProgram(self.as_inner().wrap_in_loop(
+            loop_count_reference.into_inner(),
+            start_target.into_inner(),
+            end_target.into_inner(),
+            iterations,
+        ))
+    }
+
     // Because we can't bubble up an error from inside the closures, they panic when the given
     // Python functions return an error or an unexpected type. This is unusual, but in a Python
     // program, this function will only raise because [`pyo3`] wraps Rust panics in a

quil-rs/src/program/mod.rs

@@ -21,8 +21,10 @@ use ndarray::Array2;
 use nom_locate::LocatedSpan;
 
 use crate::instruction::{
-    Declaration, FrameDefinition, FrameIdentifier, GateDefinition, GateError, Instruction, Matrix,
-    Qubit, QubitPlaceholder, Target, TargetPlaceholder, Waveform, WaveformDefinition,
+    Arithmetic, ArithmeticOperand, ArithmeticOperator, Declaration, FrameDefinition,
+    FrameIdentifier, GateDefinition, GateError, Instruction, Jump, JumpUnless, Label, Matrix,
+    MemoryReference, Move, Qubit, QubitPlaceholder, ScalarType, Target, TargetPlaceholder, Vector,
+    Waveform, WaveformDefinition,
 };
 use crate::parser::{lex, parse_instructions, ParseError};
 use crate::quil::Quil;
@@ -379,6 +381,76 @@ impl Program {
         Ok(expanded_program)
     }
 
+    /// Return a copy of the [`Program`] wrapped in a loop that repeats `iterations` times.
+    ///
+    /// The loop is constructed by wrapping the body of the program in classical Quil instructions.
+    /// The given `loop_count_reference` must refer to an INTEGER memory region. The value at the
+    /// reference given will be set to `iterations` and decremented in the loop. The loop will
+    /// terminate when the reference reaches 0. For this reason your program should not itself
+    /// modify the value at the reference unless you intend to modify the remaining number of
+    /// iterations (i.e. to break the loop).
+    ///
+    /// The given `start_target` and `end_target` will be used as the entry and exit points for the
+    /// loop, respectively. You should provide unique [`Target`]s that won't be used elsewhere in
+    /// the program.
+    ///
+    /// If `iterations` is 0, then a copy of the program is returned without any changes.
+    pub fn wrap_in_loop(
+        &self,
+        loop_count_reference: MemoryReference,
+        start_target: Target,
+        end_target: Target,
+        iterations: u32,
+    ) -> Self {
+        if iterations == 0 {
+            return self.clone();
+        }
+
+        let mut looped_program = self.clone_without_body_instructions();
+
+        looped_program.add_instructions(
+            vec![
+                Instruction::Declaration(Declaration {
+                    name: loop_count_reference.name.clone(),
+                    size: Vector {
+                        data_type: ScalarType::Integer,
+                        length: 1,
+                    },
+                    sharing: None,
+                }),
+                Instruction::Move(Move {
+                    destination: loop_count_reference.clone(),
+                    source: ArithmeticOperand::LiteralInteger(iterations.into()),
+                }),
+                Instruction::Label(Label {
+                    target: start_target.clone(),
+                }),
+            ]
+            .into_iter()
+            .chain(self.body_instructions().cloned())
+            .chain(vec![
+                Instruction::Arithmetic(Arithmetic {
+                    operator: ArithmeticOperator::Subtract,
+                    destination: ArithmeticOperand::MemoryReference(MemoryReference {
+                        name: loop_count_reference.name.clone(),
+                        index: 0,
+                    }),
+                    source: ArithmeticOperand::LiteralInteger(1),
+                }),
+                Instruction::JumpUnless(JumpUnless {
+                    target: end_target.clone(),
+                    condition: loop_count_reference,
+                }),
+                Instruction::Jump(Jump {
+                    target: start_target,
+                }),
+            ])
+            .collect::<Vec<Instruction>>(),
+        );
+
+        looped_program
+    }
+
     /// Resolve [`LabelPlaceholder`]s and [`QubitPlaceholder`]s within the program using default resolvers.
     ///
     /// See [`resolve_placeholders_with_custom_resolvers`](Self::resolve_placeholders_with_custom_resolvers),
@@ -622,7 +694,7 @@ mod tests {
         real,
     };
     use approx::assert_abs_diff_eq;
-    use insta::assert_debug_snapshot;
+    use insta::{assert_debug_snapshot, assert_snapshot};
     use ndarray::{array, linalg::kron, Array2};
     use num_complex::Complex64;
     use once_cell::sync::Lazy;
@@ -1243,4 +1315,34 @@ CNOT 2 3
             ]
         );
     }
+
+    #[test]
+    fn test_wrap_in_loop() {
+        let input = "DECLARE ro BIT
+DECLARE shot_count INTEGER
+MEASURE q ro
+JUMP-UNLESS @end-reset ro
+X q
+LABEL @end-reset
+
+DEFCAL I 0:
+    DELAY 0 1.0
+DEFFRAME 0 \"rx\":
+    HARDWARE-OBJECT: \"hardware\"
+DEFWAVEFORM custom:
+    1,2
+I 0
+";
+        let program = Program::from_str(input).unwrap().wrap_in_loop(
+            MemoryReference {
+                name: "shot_count".to_string(),
+                index: 0,
+            },
+            Target::Fixed("loop-start".to_string()),
+            Target::Fixed("loop-end".to_string()),
+            10,
+        );
+
+        assert_snapshot!(program.to_quil().unwrap())
+    }
 }

quil-rs/src/program/snapshots/quil_rs__program__tests__wrap_in_loop.snap

@@ -0,0 +1,23 @@
+---
+source: quil-rs/src/program/mod.rs
+expression: program.to_quil().unwrap()
+---
+DECLARE ro BIT[1]
+DECLARE shot_count INTEGER[1]
+DEFFRAME 0 "rx":
+	HARDWARE-OBJECT: "hardware"
+DEFWAVEFORM custom:
+	1, 2
+DEFCAL I 0:
+	DELAY 0 1
+MOVE shot_count[0] 10
+LABEL @loop-start
+MEASURE q ro[0]
+JUMP-UNLESS @end-reset ro[0]
+X q
+LABEL @end-reset
+I 0
+SUB shot_count[0] 1
+JUMP-UNLESS @loop-end shot_count[0]
+JUMP @loop-start
+