feat: allow using instruction handler when simplifying program

quil-rs/src/instruction/mod.rs

@@ -21,6 +21,7 @@ use crate::expression::Expression;
 use crate::parser::lex;
 use crate::parser::parse_instructions;
 use crate::program::frame::{FrameMatchCondition, FrameMatchConditions};
+use crate::program::ProgramError;
 use crate::program::{MatchedFrames, MemoryAccesses};
 use crate::quil::{write_join_quil, Quil, ToQuilResult};
 use crate::Program;
@@ -923,6 +924,11 @@ impl InstructionHandler {
             .and_then(|f| f(instruction))
             .unwrap_or_else(|| instruction.get_memory_accesses())
     }
+
+    /// Like [`Program::into_simplified`], but using custom instruction handling.
+    pub fn simplify_program(&mut self, program: &Program) -> Result<Program, ProgramError> {
+        program.simplify_with_handler(self)
+    }
 }
 
 #[cfg(test)]
@@ -1061,4 +1067,37 @@ RX(%a) 0",
             assert_eq!(qubit_2, Qubit::Placeholder(placeholder_2));
         }
     }
+
+    mod instruction_handler {
+        use super::super::*;
+
+        #[test]
+        fn it_considers_custom_instruction_frames() {
+            let program = r#"DEFFRAME 0 "rf":
+    CENTER-FREQUENCY: 3e9
+
+PRAGMA USES-ALL-FRAMES
+"#
+            .parse::<Program>()
+            .unwrap();
+
+            // This test assumes that the default simplification behavior will not assign frames to
+            // `PRAGMA` instructions. This is verified below.
+            assert!(program.into_simplified().unwrap().frames.is_empty());
+
+            let mut handler =
+                InstructionHandler::default().set_matching_frames(|instruction, program| {
+                    if let Instruction::Pragma(_) = instruction {
+                        Some(Some(MatchedFrames {
+                            used: program.frames.get_keys().into_iter().collect(),
+                            blocked: HashSet::new(),
+                        }))
+                    } else {
+                        None
+                    }
+                });
+
+            assert_eq!(handler.simplify_program(&program).unwrap().frames.len(), 1);
+        }
+    }
 }

quil-rs/src/program/mod.rs

@@ -22,9 +22,9 @@ use nom_locate::LocatedSpan;
 
 use crate::instruction::{
     Arithmetic, ArithmeticOperand, ArithmeticOperator, Declaration, FrameDefinition,
-    FrameIdentifier, GateDefinition, GateError, Instruction, Jump, JumpUnless, Label, Matrix,
-    MemoryReference, Move, Qubit, QubitPlaceholder, ScalarType, Target, TargetPlaceholder, Vector,
-    Waveform, WaveformDefinition,
+    FrameIdentifier, GateDefinition, GateError, Instruction, InstructionHandler, 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;
@@ -349,16 +349,10 @@ impl Program {
         instructions
     }
 
-    /// Simplify this program into a new [`Program`] which contains only instructions
-    /// and definitions which are executed; effectively, perform dead code removal.
-    ///
-    /// Removes:
-    /// - All calibrations, following calibration expansion
-    /// - Frame definitions which are not used by any instruction such as `PULSE` or `CAPTURE`
-    /// - Waveform definitions which are not used by any instruction
-    ///
-    /// When a valid program is simplified, it remains valid.
-    pub fn into_simplified(&self) -> Result<Self> {
+    pub(crate) fn simplify_with_handler(
+        &self,
+        instruction_handler: &mut InstructionHandler,
+    ) -> Result<Self> {
         let mut expanded_program = self.expand_calibrations()?;
         // Remove calibrations such that the resulting program contains
         // only instructions. Calibrations have already been expanded, so
@@ -369,7 +363,9 @@ impl Program {
         let mut waveforms_used: HashSet<&String> = HashSet::new();
 
         for instruction in &expanded_program.instructions {
-            if let Some(matched_frames) = expanded_program.get_frames_for_instruction(instruction) {
+            if let Some(matched_frames) =
+                instruction_handler.matching_frames(instruction, &expanded_program)
+            {
                 frames_used.extend(matched_frames.used())
             }
 
@@ -386,6 +382,24 @@ impl Program {
         Ok(expanded_program)
     }
 
+    /// Simplify this program into a new [`Program`] which contains only instructions
+    /// and definitions which are executed; effectively, perform dead code removal.
+    ///
+    /// Removes:
+    /// - All calibrations, following calibration expansion
+    /// - Frame definitions which are not used by any instruction such as `PULSE` or `CAPTURE`
+    /// - Waveform definitions which are not used by any instruction
+    ///
+    /// When a valid program is simplified, it remains valid.
+    ///
+    /// # Note
+    ///
+    /// If you need custom instruction handling during simplification, using
+    /// [`InstructionHandler::simplify_program`] instead.
+    pub fn into_simplified(&self) -> Result<Self> {
+        self.simplify_with_handler(&mut InstructionHandler::default())
+    }
+
     /// 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.