fix: Record measurements into named classical registers

src/braket/default_simulator/openqasm/circuit.py

@@ -13,7 +13,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Sequence
 
 import numpy as np
 
@@ -36,6 +36,11 @@ def __init__(self, name: str | None, size: int, order: int):
     def size(self) -> int:
         return len(self.sources)
 
+    @property
+    def measured(self) -> bool:
+        """Whether any element holds a measurement."""
+        return any(source is not None for source in self.sources)
+
     def bind(self, element: int, qubit: int) -> None:
         """Record ``qubit`` as the measurement source of ``element``.
 
@@ -186,50 +191,32 @@ def measured_qubits(self) -> list[int]:
         """The measured qubit of each measurement slot, in ``measurement_slots`` order."""
         return [qubit for _, _, qubit in self.measurement_slots]
 
-    def add_measure_by_index(
-        self,
-        target: tuple[int, ...],
-        classical_targets: Iterable[int] | None = None,
-        allow_remeasure: bool = False,
-    ) -> None:
-        """Record a measurement by flat classical index into the anonymous register.
+    @property
+    def measured_registers(self) -> list[ClassicalRegister]:
+        """The registers holding at least one measurement, in declaration order."""
+        return [register for register in self.classical_registers if register.measured]
 
-        Temporary entry point for callers that still address measurements by a
-        program-wide classical index. It is removed once measurements are routed
-        through declared registers via ``add_measure``.
+    def validate_single_measured_register(self) -> None:
+        """Reject programs whose measurements span more than one register.
 
-        Args:
-            target (tuple[int, ...]): The qubits measured, in order.
-            classical_targets (Iterable[int] | None): The classical index of each
-                qubit. ``None`` appends one new index per qubit after the indices
-                already in use.
-            allow_remeasure (bool): Whether a qubit that is already measured into a
-                register element may be measured again (replacing the source of the
-                requested index).
-        """
-        register = self.anonymous_register()
-        classical_targets = list(classical_targets) if classical_targets else None
-        for index, qubit in enumerate(target):
-            if not allow_remeasure and qubit in self.measured_qubits:
-                raise ValueError(f"Qubit {qubit} is already measured or captured.")
-            element = (
-                classical_targets[index]
-                if classical_targets
-                else max(index, sum(source is not None for source in register.sources))
-            )
-            if element >= register.size:
-                register.grow(element + 1 - register.size)
-            register.bind(element, qubit)
-            self.qubit_set.add(qubit)
-
-    @property
-    def target_classical_indices(self) -> list[int]:
-        """The classical index of each measurement slot, in ``measurement_slots`` order.
+        This validation will be removed once `output` is supported.
 
-        Temporary companion of ``add_measure_by_index``: with a single anonymous
-        register the element index is the classical index. Removed with it.
+        Raises:
+            ValueError: If measurements were recorded into more than one register.
         """
-        return [element for _, element, _ in self.measurement_slots]
+        measured = self.measured_registers
+        if len(measured) > 1:
+            names = ", ".join(
+                f"`{register.name}`"
+                if register.name is not None
+                else "measurements without a destination"
+                for register in measured
+            )
+            raise ValueError(
+                "Measurement results can only be reported for a single classical register, "
+                f"but measurements were recorded into {len(measured)}: {names}. "
+                "Declare one bit register and measure into it."
+            )
 
     def add_result(self, result: Results) -> None:
         """

src/braket/default_simulator/openqasm/simulation_path.py

@@ -14,10 +14,13 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from braket.default_simulator.operation import GateOperation
 
+if TYPE_CHECKING:  # pragma: no cover
+    from braket.default_simulator.openqasm.circuit import ClassicalRegister
+
 
 class FramedVariable:
     """Variable with frame tracking for proper scoping.
@@ -74,7 +77,7 @@ def __init__(
         variables: dict[str, FramedVariable] | None = None,
         measurements: dict[int, list[int]] | None = None,
         frame_number: int = 0,
-        mcm_outcomes: dict[int, int] | None = None,
+        mcm_outcomes: dict[tuple[ClassicalRegister, int], int] | None = None,
     ):
         self._instructions = instructions if instructions is not None else []
         self._shots = shots
@@ -103,6 +106,10 @@ def variables(self) -> dict[str, FramedVariable]:
     def measurements(self) -> dict[int, list[int]]:
         return self._measurements
 
+    @property
+    def mcm_outcomes(self) -> dict[tuple[ClassicalRegister, int], int]:
+        return self._mcm_outcomes
+
     @property
     def frame_number(self) -> int:
         return self._frame_number

src/braket/default_simulator/simulator.py

@@ -765,15 +765,12 @@ def run_openqasm(
             # Multi-path execution for programs with mid-circuit measurements
             return self._run_branched(context, openqasm_ir, shots, batch_size, created_at)
 
-        # Single-path execution (current behavior, unchanged)
+        # Single-path execution
         circuit = context.circuit
+        circuit.validate_single_measured_register()
         qubit_map = BaseLocalSimulator._map_circuit_to_contiguous_qubits(circuit)
         qubit_count = circuit.num_qubits
-        classical_bit_positions = {b: i for i, b in enumerate(circuit.target_classical_indices)}
-        measured_qubits = [
-            circuit.measured_qubits[classical_bit_positions[i]]
-            for i in sorted(circuit.target_classical_indices)
-        ]
+        measured_qubits = circuit.measured_qubits
         mapped_measured_qubits = (
             [qubit_map[q] for q in measured_qubits] if measured_qubits else None
         )
@@ -876,6 +873,7 @@ def _run_branched(
             GateModelTaskResult: Aggregated result across all paths.
         """
         circuit = context.circuit
+        circuit.validate_single_measured_register()
         path_qubit_set = set()
         for path in context.active_paths:
             for ins in path.instructions:
@@ -887,12 +885,9 @@ def _run_branched(
         qubit_count = len(qubit_map)
         BaseLocalSimulator._map_circuit_qubits(circuit, qubit_map)
 
-        # Determine measured qubits from the circuit
-        classical_bit_positions = {b: i for i, b in enumerate(circuit.target_classical_indices)}
-        measured_qubits = [
-            circuit.measured_qubits[classical_bit_positions[i]]
-            for i in sorted(circuit.target_classical_indices)
-        ]
+        # One result bit per measurement slot, in slot order
+        slots = circuit.measurement_slots
+        measured_qubits = [qubit for _, _, qubit in slots]
         mapped_measured_qubits = (
             [qubit_map[q] for q in measured_qubits] if measured_qubits else None
         )
@@ -917,16 +912,17 @@ def _run_branched(
             qubits_not_in_circuit = mapped_arr[~in_circuit_mask]
             measurements_array = np.array(measurements)
             selected = measurements_array[:, qubits_in_circuit]
-            # Overlay per-classical-bit MCM outcomes over the final-state
-            # samples so repeat measurements aren't reported as the qubit's
-            # final resampled state.
+            # Slots whose register element was written by a mid-circuit
+            # measurement on a path take that path's recorded outcome, so a
+            # qubit that evolved after being measured isn't reported as its
+            # final resampled state. Other slots keep the final-state sample.
             shot_offset = 0
             for path in context.active_paths:
-                mcm_outcomes = path._mcm_outcomes
-                for shot_idx in range(shot_offset, shot_offset + path.shots):
-                    for col, classical_idx in enumerate(sorted(circuit.target_classical_indices)):
-                        if classical_idx in mcm_outcomes:
-                            selected[shot_idx, col] = str(mcm_outcomes[classical_idx])
+                for slot_index, (register, register_index, _) in enumerate(slots):
+                    outcome = path.mcm_outcomes.get((register, register_index))
+                    if outcome is not None:
+                        path_shots = slice(shot_offset, shot_offset + path.shots)
+                        selected[path_shots, slot_index] = str(outcome)
                 shot_offset += path.shots
             measurements = np.pad(selected, ((0, 0), (0, len(qubits_not_in_circuit)))).tolist()
 

test/unit_tests/braket/default_simulator/openqasm/test_circuit.py

@@ -49,6 +49,7 @@ def test_declare_register():
     assert (d.name, d.size, d.order, d.sources) == ("d", 1, 1, [None])
     assert circuit.measurement_slots == []
     assert circuit.measured_qubits == []
+    assert circuit.measured_registers == []
 
 
 def test_registers_with_same_name_are_distinct():
@@ -139,50 +140,41 @@ def test_anonymous_register_grows_on_demand():
     assert circuit.measured_qubits == [1, 0, 2]
 
 
-def test_add_measure_by_index_rejects_duplicate_qubit_by_default():
+def test_measured_registers_skips_unmeasured():
     circuit = Circuit()
-    circuit.add_measure_by_index((0,), [0])
-    with pytest.raises(ValueError, match="Qubit 0 is already measured or captured."):
-        circuit.add_measure_by_index((0,), [1])
-
+    circuit.declare_register("unused", 2)
+    c = circuit.declare_register("c", 2)
+    circuit.add_measure((1,), c, [1])
+    assert circuit.measured_registers == [c]
+    assert circuit.measured_qubits == [1]
+    circuit.validate_single_measured_register()
 
-def test_add_measure_by_index_binds_elements_of_anonymous_register():
-    """Explicit indices address elements of the anonymous register; the derived
-    columns come out in index order."""
-    circuit = Circuit()
-    circuit.add_measure_by_index((5, 7), [2, 0], allow_remeasure=True)
-    (anonymous,) = circuit.classical_registers
-    assert anonymous is circuit.anonymous_register()
-    assert anonymous.sources == [7, None, 5]
-    assert circuit.measured_qubits == [7, 5]
-    assert circuit.target_classical_indices == [0, 2]
-    assert circuit.qubit_set == {5, 7}
 
+def test_validate_rejects_two_declared_registers():
+    """Measurements into two declared registers cannot be reported as one bit string.
 
-def test_add_measure_by_index_appends_after_bound_elements():
+    Remove once per register results are reported through the ``output`` field.
+    """
     circuit = Circuit()
-    circuit.add_measure_by_index((3,), [0], allow_remeasure=True)
-    circuit.add_measure_by_index((1, 2), allow_remeasure=True)
-    assert circuit.anonymous_register().sources == [3, 1, 2]
-    assert circuit.target_classical_indices == [0, 1, 2]
-
+    c = circuit.declare_register("c", 2)
+    d = circuit.declare_register("d", 2)
+    circuit.add_measure((0,), c, [1])
+    circuit.add_measure((1,), d, [1])
+    with pytest.raises(ValueError, match="recorded into 2: `c`, `d`"):
+        circuit.validate_single_measured_register()
 
-def test_add_measure_by_index_remeasure_replaces_source():
-    circuit = Circuit()
-    circuit.add_measure_by_index((0,), [1], allow_remeasure=True)
-    circuit.add_measure_by_index((4,), [1], allow_remeasure=True)
-    assert circuit.anonymous_register().sources == [None, 4]
-    assert circuit.measured_qubits == [4]
-    assert circuit.target_classical_indices == [1]
 
+def test_validate_rejects_declared_plus_anonymous():
+    """A declared register plus destination-less measurements is also two registers.
 
-def test_add_measure_by_index_sparse_then_append_fills_next_bound_count():
-    """With elements {0, 2} bound, the next bare measurement takes index 2 and
-    replaces it, matching the flat-index bookkeeping it stands in for."""
+    Remove once per register results are reported through the ``output`` field.
+    """
     circuit = Circuit()
-    circuit.add_measure_by_index((0, 1), [0, 2], allow_remeasure=True)
-    circuit.add_measure_by_index((5,), allow_remeasure=True)
-    assert circuit.anonymous_register().sources == [0, None, 5]
+    c = circuit.declare_register("c", 1)
+    circuit.add_measure((0,), c)
+    circuit.add_measure((1,))
+    with pytest.raises(ValueError, match="`c`, measurements without a destination"):
+        circuit.validate_single_measured_register()
 
 
 def test_clear_measurement():