Repository: rigetti/quil-rs (language: python)

Issue:
Add method to `Program` that will wrap it in a Quil based num-shots loop
The QVM and QCS API accept a `trials`/`num_shots` parameter that allows users to configure how many times their programs run. However, it is possible, and sometimes desirable, to do this by just using Quil's classical instructions. We can add a method to `Program` that applies a generic template to the contained program to wrap it in a "num shots" loop.

Example template:
```
# User declarations go here
DECLARE shot_count INTEGER[1]
MOVE shot_count[0] 100 # Configurable via method paramter
LABEL @START
# Rest of the user program goes here
SUB shot_count[0] 1
JUMP-UNLESS @END shot_count[0]
JUMP @START
LABEL @END
```

We can then use this utility to implement rigetti/pyQuil#1712 upstream.



Discussion / hints:

Fix patch:
diff --git a/quil-py/quil/program/__init__.pyi b/quil-py/quil/program/__init__.pyi
index 1037ee78..d8ccdb16 100644
--- a/quil-py/quil/program/__init__.pyi
+++ b/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.
diff --git a/quil-py/src/program/mod.rs b/quil-py/src/program/mod.rs
index 3b434a5b..6757e13a 100644
--- a/quil-py/src/program/mod.rs
+++ b/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
diff --git a/quil-rs/src/program/mod.rs b/quil-rs/src/program/mod.rs
index 180a441f..2904f8ec 100644
--- a/quil-rs/src/program/mod.rs
+++ b/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())
+    }
 }


Test patch:
diff --git a/quil-rs/src/program/snapshots/quil_rs__program__tests__wrap_in_loop.snap b/quil-rs/src/program/snapshots/quil_rs__program__tests__wrap_in_loop.snap
new file mode 100644
index 00000000..695ae589
--- /dev/null
+++ b/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
+