feat: Track classical register indices for measurements

src/braket/circuits/braket_program_context.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
+from collections.abc import Iterable
 from typing import Optional, Union
 
 import numpy as np
@@ -161,16 +162,17 @@ def handle_parameter_value(
             return FreeParameterExpression(evaluated_value)
         return value
 
-    def add_measure(self, target: tuple[int]) -> None:
+    def add_measure(
+        self, target: tuple[int], classical_targets: Optional[Iterable[int]] = None
+    ) -> None:
         """Add a measure instruction to the circuit
 
         Args:
             target (tuple[int]): the target qubits to be measured.
+            classical_targets (Optional[Iterable[int]]): the classical registers
+                to use in the qubit measurement.
         """
-        for index, qubit in enumerate(target):
+        for iter, qubit in enumerate(target):
+            index = classical_targets[iter] if classical_targets else iter
             instruction = Instruction(Measure(index=index), qubit)
             self._circuit.add_instruction(instruction)
-            if self._circuit._measure_targets:
-                self._circuit._measure_targets.append(qubit)
-            else:
-                self._circuit._measure_targets = [qubit]

src/braket/circuits/circuit.py

@@ -504,6 +504,11 @@ def add_instruction(
         # Check if there is a measure instruction on the circuit
         self._check_if_qubit_measured(instruction, target, target_mapping)
 
+        # Update measure targets if instruction is a measurement
+        if isinstance(instruction.operator, Measure):
+            measure_target = target or instruction.target[0]
+            self._measure_targets = (self._measure_targets or []) + [measure_target]
+
         if not target_mapping and not target:
             # Nothing has been supplied, add instruction
             instructions_to_add = [instruction]
@@ -710,10 +715,6 @@ def _add_measure(self, target_qubits: QubitSetInput) -> None:
                     target=target,
                 )
             )
-            if self._measure_targets:
-                self._measure_targets.append(target)
-            else:
-                self._measure_targets = [target]
 
     def measure(self, target_qubits: QubitSetInput) -> Circuit:
         """