From 4c64a645602c0defdc6c549244176a56fe137f01 Mon Sep 17 00:00:00 2001
From: Kevin Hartman <kevin@hart.mn>
Date: Thu, 19 Sep 2024 11:32:31 -0400
Subject: [PATCH] Add `Operation::blocks` method. (#13056)

* Add PyInstruction::blocks method.

This gives us a way to get the blocks of a control flow
operation as an iterator of CircuitData. If called on
a non-control-flow instruction, returns None.

This is not intended as a final API for this, which
will likely instead return an iterator of &CircuitData
once control flow operations have been fully ported to
Rust and own their blocks without requiring conversion.

* Expose blocks at Operation level.

Also addresses other review comments.

* Add comment to explain unwraps.
---
 crates/circuit/src/operations.rs         | 53 +++++++++++++++++++++---
 crates/circuit/src/packed_instruction.rs |  4 ++
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/crates/circuit/src/operations.rs b/crates/circuit/src/operations.rs
index ad76e9d44008..a73d81ad7180 100644
--- a/crates/circuit/src/operations.rs
+++ b/crates/circuit/src/operations.rs
@@ -12,6 +12,7 @@
 
 use approx::relative_eq;
 use std::f64::consts::PI;
+use std::vec;
 
 use crate::circuit_data::CircuitData;
 use crate::circuit_instruction::ExtraInstructionAttributes;
@@ -163,6 +164,7 @@ pub trait Operation {
     fn num_clbits(&self) -> u32;
     fn num_params(&self) -> u32;
     fn control_flow(&self) -> bool;
+    fn blocks(&self) -> Vec<CircuitData>;
     fn matrix(&self, params: &[Param]) -> Option<Array2<Complex64>>;
     fn definition(&self, params: &[Param]) -> Option<CircuitData>;
     fn standard_gate(&self) -> Option<StandardGate>;
@@ -228,6 +230,15 @@ impl<'a> Operation for OperationRef<'a> {
         }
     }
     #[inline]
+    fn blocks(&self) -> Vec<CircuitData> {
+        match self {
+            OperationRef::Standard(standard) => standard.blocks(),
+            OperationRef::Gate(gate) => gate.blocks(),
+            OperationRef::Instruction(instruction) => instruction.blocks(),
+            OperationRef::Operation(operation) => operation.blocks(),
+        }
+    }
+    #[inline]
     fn matrix(&self, params: &[Param]) -> Option<Array2<Complex64>> {
         match self {
             Self::Standard(standard) => standard.matrix(params),
@@ -530,20 +541,20 @@ impl Operation for StandardGate {
         STANDARD_GATE_NUM_QUBITS[*self as usize]
     }
 
-    fn num_params(&self) -> u32 {
-        STANDARD_GATE_NUM_PARAMS[*self as usize]
-    }
-
     fn num_clbits(&self) -> u32 {
         0
     }
 
+    fn num_params(&self) -> u32 {
+        STANDARD_GATE_NUM_PARAMS[*self as usize]
+    }
+
     fn control_flow(&self) -> bool {
         false
     }
 
-    fn directive(&self) -> bool {
-        false
+    fn blocks(&self) -> Vec<CircuitData> {
+        vec![]
     }
 
     fn matrix(&self, params: &[Param]) -> Option<Array2<Complex64>> {
@@ -2043,6 +2054,10 @@ impl Operation for StandardGate {
     fn standard_gate(&self) -> Option<StandardGate> {
         Some(*self)
     }
+
+    fn directive(&self) -> bool {
+        false
+    }
 }
 
 const FLOAT_ZERO: Param = Param::Float(0.0);
@@ -2145,6 +2160,26 @@ impl Operation for PyInstruction {
     fn control_flow(&self) -> bool {
         self.control_flow
     }
+    fn blocks(&self) -> Vec<CircuitData> {
+        if !self.control_flow {
+            return vec![];
+        }
+        Python::with_gil(|py| -> Vec<CircuitData> {
+            // We expect that if PyInstruction::control_flow is true then the operation WILL
+            // have a 'blocks' attribute which is a tuple of the Python QuantumCircuit.
+            let raw_blocks = self.instruction.getattr(py, "blocks").unwrap();
+            let blocks: &Bound<PyTuple> = raw_blocks.downcast_bound::<PyTuple>(py).unwrap();
+            blocks
+                .iter()
+                .map(|b| {
+                    b.getattr(intern!(py, "_data"))
+                        .unwrap()
+                        .extract::<CircuitData>()
+                        .unwrap()
+                })
+                .collect()
+        })
+    }
     fn matrix(&self, _params: &[Param]) -> Option<Array2<Complex64>> {
         None
     }
@@ -2211,6 +2246,9 @@ impl Operation for PyGate {
     fn control_flow(&self) -> bool {
         false
     }
+    fn blocks(&self) -> Vec<CircuitData> {
+        vec![]
+    }
     fn matrix(&self, _params: &[Param]) -> Option<Array2<Complex64>> {
         Python::with_gil(|py| -> Option<Array2<Complex64>> {
             match self.gate.getattr(py, intern!(py, "to_matrix")) {
@@ -2287,6 +2325,9 @@ impl Operation for PyOperation {
     fn control_flow(&self) -> bool {
         false
     }
+    fn blocks(&self) -> Vec<CircuitData> {
+        vec![]
+    }
     fn matrix(&self, _params: &[Param]) -> Option<Array2<Complex64>> {
         None
     }
diff --git a/crates/circuit/src/packed_instruction.rs b/crates/circuit/src/packed_instruction.rs
index 82c678031a15..af72b3226a7c 100644
--- a/crates/circuit/src/packed_instruction.rs
+++ b/crates/circuit/src/packed_instruction.rs
@@ -395,6 +395,10 @@ impl Operation for PackedOperation {
         self.view().control_flow()
     }
     #[inline]
+    fn blocks(&self) -> Vec<CircuitData> {
+        self.view().blocks()
+    }
+    #[inline]
     fn matrix(&self, params: &[Param]) -> Option<Array2<Complex64>> {
         self.view().matrix(params)
     }