feature: support measure_ff and cc_prx gates

src/autoqasm/instructions/__init__.py

@@ -27,5 +27,5 @@ def bell():
 
 from .gates import *
 from .instructions import reset  # noqa: F401
-from .measurements import measure  # noqa: F401
+from .measurements import measure, measure_ff  # noqa: F401
 from .qubits import global_qubit_register  # noqa: F401

src/autoqasm/instructions/gates.py

@@ -40,6 +40,38 @@ def ccnot(
     _qubit_instruction("ccnot", [control_0, control_1, target], **kwargs)
 
 
+def cc_prx(
+    target: QubitIdentifierType,
+    angle_0: GateParameterType,
+    angle_1: GateParameterType,
+    feedback_key: int,
+    **kwargs,
+) -> None:
+    """Classically-controlled Phased Rx gate.
+
+    Applies :func:`prx` to ``target`` on the runtime branches where the
+    classical feedback bit identified by ``feedback_key`` is ``1``. The
+    feedback bit is produced by a prior :func:`measure_ff` with the same
+    ``feedback_key``.
+
+    This is an IQM experimental capability. See
+    :class:`braket.experimental_capabilities.iqm.classical_control.CCPRx`
+    for the corresponding Braket SDK surface.
+
+    Args:
+        target (QubitIdentifierType): Target qubit.
+        angle_0 (GateParameterType): First PRx angle in radians.
+        angle_1 (GateParameterType): Second PRx angle in radians.
+        feedback_key (int): Integer key identifying which prior
+            ``measure_ff`` result gates this operation. Must be the same
+            value passed to the corresponding :func:`measure_ff` call.
+
+    """
+    _qubit_instruction(
+        "cc_prx", [target], angle_0, angle_1, feedback_key, is_unitary=False, **kwargs
+    )
+
+
 def cnot(
     control: QubitIdentifierType,
     target: QubitIdentifierType,

src/autoqasm/instructions/measurements.py

@@ -28,6 +28,7 @@ def my_program():
 
 from autoqasm import program
 from autoqasm import types as aq_types
+from autoqasm.instructions.instructions import _qubit_instruction
 from autoqasm.instructions.qubits import GlobalQubitRegister, _qubit, global_qubit_register
 
 
@@ -67,3 +68,29 @@ def measure(
             oqpy_program.measure(qubit, bit_var[idx])
 
     return bit_var
+
+
+def measure_ff(
+    target: aq_types.QubitIdentifierType,
+    feedback_key: int,
+    **kwargs,
+) -> None:
+    """Measure a qubit and store its result under a classical feedback key.
+
+    The measurement result is not bound to a Python variable; instead it is
+    stored by the runtime under the integer ``feedback_key`` so that it can
+    be consumed later in the same program by a classically-controlled
+    operation such as :func:`autoqasm.instructions.cc_prx`.
+
+    This is an IQM experimental capability. See
+    :class:`braket.experimental_capabilities.iqm.classical_control.MeasureFF`
+    for the corresponding Braket SDK surface.
+
+    Args:
+        target (QubitIdentifierType): The qubit to measure.
+        feedback_key (int): Integer key under which the measurement result
+            is recorded. Must match the ``feedback_key`` passed to any
+            subsequent ``cc_prx`` call that depends on this measurement.
+
+    """
+    _qubit_instruction("measure_ff", [target], feedback_key, is_unitary=False, **kwargs)

src/autoqasm/simulator/native_interpreter.py

@@ -26,6 +26,7 @@
     BitType,
     BooleanLiteral,
     ClassicalDeclaration,
+    Identifier,
     IndexedIdentifier,
     IODeclaration,
     IOKeyword,
@@ -49,6 +50,7 @@ def __init__(
         self.simulation = simulation
         context = context or McmProgramContext()
         super().__init__(context, logger)
+        self._declared_feedback_keys: set[int] = set()
 
     def simulate(
         self,
@@ -80,6 +82,7 @@ def simulate(
         program = parse(source)
         for _ in range(shots):
             program_copy = deepcopy(program)
+            self._declared_feedback_keys.clear()
             self.visit(program_copy)
             self.context.save_output_values()
             self.context.num_qubits = 0
@@ -151,3 +154,79 @@ def _(self, node: IODeclaration) -> None:
             init_value = wrap_value_into_literal(self.context.inputs[node.identifier.name])
             declaration = ClassicalDeclaration(node.type, node.identifier, init_value)
             self.visit(declaration)
+
+    def handle_builtin_gate(
+        self,
+        gate_name: str,
+        arguments: list,
+        qubits: list,
+        modifiers: list,
+    ) -> None:
+        """Handle a call to a built-in quantum gate.
+
+        Intercepts the IQM classical-control gates ``measure_ff`` and
+        ``cc_prx`` and implements them directly against ``self.simulation``,
+        bypassing the upstream ``ProgramContext`` mid-circuit-measurement
+        pipeline (which is built around a one-pass, all-shots-at-once
+        branching model incompatible with this per-shot interpreter).
+        """
+        if gate_name == "measure_ff":
+            self._handle_measure_ff(arguments, qubits)
+            return
+        if gate_name == "cc_prx":
+            self._handle_cc_prx(arguments, qubits, modifiers)
+            return
+        super().handle_builtin_gate(gate_name, arguments, qubits, modifiers)
+
+    def _handle_measure_ff(self, arguments: list, qubits: list) -> None:
+        """Measure the target qubit and bind the outcome under a synthetic
+        bit variable ``__ff_<feedback_key>__`` in the context's symbol table.
+        """
+        feedback_key = int(arguments[0].value)
+        if feedback_key in self._declared_feedback_keys:
+            raise ValueError(
+                f"measure_ff feedback key {feedback_key} is already in use; "
+                "feedback keys must be unique within a program."
+            )
+        ff_name = _feedback_key_name(feedback_key)
+        self.simulation.evolve(self.context.pop_instructions())
+        targets = self.context.get_qubits(qubits[0])
+        outcome = self.simulation.measure(targets)
+        self._declared_feedback_keys.add(feedback_key)
+        self.context.declare_variable(ff_name, BitType(size=None))
+        self.context.update_value(
+            Identifier(name=ff_name),
+            BooleanLiteral(bool(outcome[0])),
+        )
+
+    def _handle_cc_prx(
+        self,
+        arguments: list,
+        qubits: list,
+        modifiers: list,
+    ) -> None:
+        """Apply ``prx(angle_0, angle_1) q`` only when the feedback bit
+        identified by ``feedback_key`` is ``1``.
+        """
+        feedback_key = int(arguments[2].value)
+        ff_name = _feedback_key_name(feedback_key)
+        try:
+            ff_value = self.context.get_value_by_identifier(Identifier(name=ff_name))
+        except KeyError as exc:
+            raise ValueError(
+                f"cc_prx references feedback key {feedback_key}, but no measure_ff "
+                "has been recorded for that key."
+            ) from exc
+        if ff_value is None or not bool(getattr(ff_value, "value", ff_value)):
+            return
+        # Dispatch through the normal builtin-gate path so gate modifiers
+        # (ctrl, pow, etc.) are honoured consistently with other gates.
+        super().handle_builtin_gate("prx", arguments[:2], qubits, modifiers)
+
+
+def _feedback_key_name(feedback_key: int) -> str:
+    """Synthetic bit-variable name used to store a feedback key's value.
+
+    Matches the convention used by ``braket.default_simulator``.
+    """
+    return f"__ff_{int(feedback_key)}__"

src/autoqasm/simulator/simulator.py

@@ -21,6 +21,16 @@
 
 
 class McmSimulator(StateVectorSimulator):
+    """AutoQASM-backed local simulator registered under the ``autoqasm`` device ID.
+
+    Deprecated: kept for now because ``braket.devices.LocalSimulator`` does
+    not yet execute OpenQASM ``output`` variable declarations, which
+    ``@aq.main`` functions emit for every Python return value. This class
+    will be removed once upstream ``LocalSimulator`` supports output
+    variables; after that, autoqasm programs can run directly on
+    ``braket.devices.LocalSimulator()``.
+    """
+
     DEVICE_ID = "autoqasm"
 
     def initialize_simulation(self, **kwargs) -> Simulation:

test/unit_tests/autoqasm/test_classical_control.py

@@ -0,0 +1,167 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+"""Tests for the IQM classical-control gates ``cc_prx`` and ``measure_ff``."""
+
+import math
+
+import pytest
+from braket.devices import LocalSimulator
+
+import autoqasm as aq
+from autoqasm.instructions import cc_prx, h, measure, measure_ff
+
+
+def test_measure_ff_emits_feedback_key() -> None:
+    @aq.main
+    def program():
+        h(0)
+        measure_ff(0, 0)
+
+    ir = program.build().to_ir()
+    assert "measure_ff(0) __qubits__[0];" in ir
+
+
+def test_cc_prx_emits_angles_and_feedback_key() -> None:
+    @aq.main
+    def program():
+        h(0)
+        measure_ff(0, 0)
+        cc_prx(1, 0.15, 0.25, 0)
+
+    ir = program.build().to_ir()
+    assert "cc_prx(0.15, 0.25, 0) __qubits__[1];" in ir
+
+
+def test_measure_ff_different_feedback_keys() -> None:
+    @aq.main
+    def program():
+        measure_ff(0, 0)
+        measure_ff(1, 5)
+
+    ir = program.build().to_ir()
+    assert "measure_ff(0) __qubits__[0];" in ir
+    assert "measure_ff(5) __qubits__[1];" in ir
+
+
+def test_cc_prx_symbolic_angles() -> None:
+    """``cc_prx`` should accept ``FreeParameterExpression``-style angles
+    exactly like other parameterised gates (e.g. ``prx``)."""
+
+    @aq.main
+    def program(theta: float):
+        measure_ff(0, 0)
+        cc_prx(1, theta, 0.0, 0)
+
+    ir = program.build().to_ir()
+    assert "input float theta;" in ir
+    assert "cc_prx(theta, 0.0, 0) __qubits__[1];" in ir
+
+
+def test_cc_prx_disallowed_inside_gate_definition() -> None:
+    """``cc_prx`` requires classical feedback; gate definitions must be
+    purely unitary. Using it inside ``@aq.gate`` should raise
+    ``InvalidGateDefinition``."""
+
+    @aq.gate
+    def bad_gate(q: aq.Qubit):
+        cc_prx(q, 0.1, 0.2, 0)
+
+    @aq.main
+    def program():
+        bad_gate(0)
+
+    with pytest.raises(aq.errors.InvalidGateDefinition):
+        program.build()
+
+
+def test_measure_ff_disallowed_inside_gate_definition() -> None:
+    @aq.gate
+    def bad_gate(q: aq.Qubit):
+        measure_ff(q, 0)
+
+    @aq.main
+    def program():
+        bad_gate(0)
+
+    with pytest.raises(aq.errors.InvalidGateDefinition):
+        program.build()
+
+
+def test_classical_control_runs_on_local_simulator() -> None:
+    """End-to-end: measure |+> on qubit 0 (50/50 outcome), and conditionally
+    X qubit 1 via ``cc_prx(pi, 0, key)`` on the measured-1 branch. Qubit 1
+    outcomes should track the qubit-0 feedback, giving ``00`` / ``11``."""
+
+    @aq.main
+    def teleport_like():
+        h(0)
+        measure_ff(0, 0)
+        cc_prx(1, math.pi, 0.0, 0)
+
+    result = LocalSimulator().run(teleport_like, shots=200).result()
+    counts = result.measurement_counts
+    # Outcomes should be only the correlated Bell-like pair.
+    for outcome in counts:
+        assert outcome in {"00", "11"}, f"unexpected outcome: {outcome}"
+    # With 200 shots we expect both outcomes to appear.
+    assert "00" in counts
+    assert "11" in counts
+
+
+def test_classical_control_runs_on_autoqasm_simulator() -> None:
+    """Same behaviour on the AutoQASM-backed simulator."""
+
+    @aq.main
+    def teleport_like():
+        h(0)
+        measure_ff(0, 0)
+        cc_prx(1, math.pi, 0.0, 0)
+        measure(1)
+
+    result = LocalSimulator("autoqasm").run(teleport_like, shots=100).result()
+    measurements = result.measurements
+    feedback = [bool(v) for v in measurements["__ff_0__"]]
+    qubit_1_key = next(k for k in measurements if k.startswith("__bit_"))
+    qubit_1 = [bool(v) for v in measurements[qubit_1_key]]
+    # Qubit 1 should match the feedback bit every time.
+    assert feedback == qubit_1, "cc_prx failed to conditionally flip qubit 1"
+    # Both outcomes should appear with 100 shots.
+    assert any(feedback)
+    assert not all(feedback)
+
+
+def test_cc_prx_missing_feedback_raises() -> None:
+    """If ``cc_prx`` is used before any ``measure_ff`` with the same
+    feedback key, the AutoQASM simulator raises a clean ValueError."""
+
+    @aq.main
+    def missing_key():
+        cc_prx(0, 0.1, 0.2, 42)
+
+    with pytest.raises(ValueError, match="feedback key 42"):
+        LocalSimulator("autoqasm").run(missing_key, shots=1).result()
+
+
+def test_measure_ff_duplicate_feedback_key_raises() -> None:
+    """IQM requires feedback keys to be unique within a program; the
+    AutoQASM simulator should raise ``ValueError`` when a feedback key
+    is reused within a single shot."""
+
+    @aq.main
+    def duplicate_key():
+        measure_ff(0, 7)
+        measure_ff(1, 7)
+
+    with pytest.raises(ValueError, match="feedback key 7"):
+        LocalSimulator("autoqasm").run(duplicate_key, shots=1).result()