Add support for OpenQASM 3.0 switch statements in OQPy

oqpy/control_flow.py

@@ -18,7 +18,17 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Iterable, Iterator, Optional, TypeVar, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ContextManager,
+    Iterable,
+    Iterator,
+    Literal,
+    Optional,
+    TypeVar,
+    overload,
+)
 
 from openpulse import ast
 
@@ -38,7 +48,7 @@
     from oqpy.program import Program
 
 
-__all__ = ["If", "Else", "ForIn", "While", "Range"]
+__all__ = ["If", "Else", "ForIn", "While", "Range", "Switch", "Case", "Default"]
 
 
 @contextlib.contextmanager
@@ -176,3 +186,110 @@ def While(program: Program, condition: OQPyExpression) -> Iterator[None]:
     yield
     state = program._pop()
     program._add_statement(ast.WhileLoop(to_ast(program, condition), state.body))
+
+
+class Switch(ContextManager[None]):
+    """Context manager for switch statement control flow.
+
+    .. code-block:: python
+
+        selector = IntVar(0)
+        with Switch(program, selector):
+            with Case(program, 0):
+                program.increment(result, 1)
+            with Case(program, 1, 2):  # Multiple values in one case
+                program.increment(result, 2)
+            with Default(program):
+                program.increment(result, 100)
+
+    """
+
+    def __init__(self, program: "Program", target: OQPyExpression):
+        self.program = program
+        self.target = target
+        self.cases: list[tuple[list[ast.Expression], list[ast.Statement]]] = []
+        self.default: list[ast.Statement] | None = None
+
+    def __enter__(self) -> None:
+        self.program._push()
+        self.program._state.active_switch = self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> Literal[False]:
+        # Pop the switch context state
+        self.program._pop()
+        if exc_type is not None:
+            # Don't add statement if an exception occurred; propagate the exception
+            return False
+        # Build the case tuples as (list of expressions, CompoundStatement)
+        case_tuples = [(values, ast.CompoundStatement(body)) for values, body in self.cases]
+        default_stmt = ast.CompoundStatement(self.default) if self.default else None
+        stmt = ast.SwitchStatement(
+            to_ast(self.program, self.target),
+            case_tuples,
+            default_stmt,
+        )
+        self.program._add_statement(stmt)
+        # Return False to indicate exceptions should not be suppressed
+        return False
+
+
+@contextlib.contextmanager
+def Case(program: "Program", *values: AstConvertible) -> Iterator[None]:
+    """Context manager for a case within a switch statement.
+
+    Must be used inside a Switch context. Multiple values can be provided
+    for a single case block.
+
+    .. code-block:: python
+
+        with Switch(program, selector):
+            with Case(program, 0):
+                # Handle case 0
+                program.increment(result, 1)
+            with Case(program, 1, 2):
+                # Handle cases 1 and 2
+                program.increment(result, 2)
+
+    """
+    if not values:
+        raise ValueError("Case requires at least one value")
+    switch = program._state.active_switch
+    if switch is None:
+        raise RuntimeError("Case must be used inside a Switch context")
+    program._push()
+    yield
+    state = program._pop()
+    case_values = [to_ast(program, v) for v in values]
+    switch.cases.append((case_values, state.body))
+
+
+@contextlib.contextmanager
+def Default(program: "Program") -> Iterator[None]:
+    """Context manager for the default case within a switch statement.
+
+    Must be used inside a Switch context.
+
+    .. code-block:: python
+
+        with Switch(program, selector):
+            with Case(program, 0):
+                program.increment(result, 1)
+            with Default(program):
+                # Handle all other cases
+                program.increment(result, 100)
+
+    """
+    switch = program._state.active_switch
+    if switch is None:
+        raise RuntimeError("Default must be used inside a Switch context")
+    if switch.default is not None:
+        raise RuntimeError("Switch statement can only have one default case")
+    program._push()
+    yield
+    state = program._pop()
+    switch.default = state.body

oqpy/program.py

@@ -25,7 +25,7 @@
 
 import warnings
 from copy import deepcopy
-from typing import Any, Hashable, Iterable, Iterator, Optional
+from typing import TYPE_CHECKING, Any, Hashable, Iterable, Iterator, Optional
 
 from openpulse import ast
 from openpulse.printer import dumps
@@ -44,6 +44,9 @@
 from oqpy.pulse import FrameVar, PortVar, WaveformVar
 from oqpy.timing import convert_duration_to_float, convert_float_to_duration
 
+if TYPE_CHECKING:
+    from oqpy.control_flow import Switch
+
 __all__ = ["Program"]
 
 
@@ -59,6 +62,7 @@ def __init__(self) -> None:
         self.body: list[ast.Statement | ast.Pragma] = []
         self.if_clause: Optional[ast.BranchingStatement] = None
         self.annotations: list[ast.Annotation] = []
+        self.active_switch: Optional["Switch"] = None  # Set when inside a switch context
 
     def add_if_clause(self, condition: ast.Expression, if_clause: list[ast.Statement]) -> None:
         if_clause_annotations, self.annotations = self.annotations, []
@@ -82,6 +86,10 @@ def add_statement(self, stmt: ast.Statement | ast.Pragma) -> None:
         # it seems to conflict with the definition of ast.Program.
         # Issue raised in https://github.com/openqasm/openqasm/issues/468
         assert isinstance(stmt, (ast.Statement, ast.Pragma))
+        if self.active_switch is not None:
+            raise RuntimeError(
+                "Statements inside a Switch block must be within a Case or Default context"
+            )
         if isinstance(stmt, ast.Statement) and self.annotations:
             stmt.annotations = self.annotations + list(stmt.annotations)
             self.annotations = []
@@ -133,7 +141,9 @@ def __iadd__(self, other: Program) -> Program:
         self.defcals.update(other.defcals)
         for name, subroutine_stmt in other.subroutines.items():
             self._add_subroutine(
-                name, subroutine_stmt, needs_declaration=name not in other.declared_subroutines
+                name,
+                subroutine_stmt,
+                needs_declaration=name not in other.declared_subroutines,
             )
         for name, gate_stmt in other.gates.items():
             self._add_gate(name, gate_stmt, needs_declaration=name not in other.declared_gates)
@@ -418,7 +428,9 @@ def declare(
         return self
 
     def delay(
-        self, time: AstConvertible, qubits_or_frames: AstConvertible | Iterable[AstConvertible] = ()
+        self,
+        time: AstConvertible,
+        qubits_or_frames: AstConvertible | Iterable[AstConvertible] = (),
     ) -> Program:
         """Apply a delay to a set of qubits or frames."""
         if not isinstance(qubits_or_frames, Iterable):
@@ -608,7 +620,9 @@ def reset(self, qubit: quantum_types.Qubit) -> Program:
         return self
 
     def measure(
-        self, qubit: quantum_types.Qubit, output_location: classical_types.BitVar | None = None
+        self,
+        qubit: quantum_types.Qubit,
+        output_location: classical_types.BitVar | None = None,
     ) -> Program:
         """Measure a particular qubit.
 
@@ -709,6 +723,13 @@ def visit_BranchingStatement(self, node: ast.BranchingStatement, context: None =
         node.else_block = self.process_statement_list(node.else_block)
         self.generic_visit(node, context)
 
+    def visit_SwitchStatement(self, node: ast.SwitchStatement, context: None = None) -> None:
+        for _, case_block in node.cases:
+            case_block.statements = self.process_statement_list(case_block.statements)
+        if node.default is not None:
+            node.default.statements = self.process_statement_list(node.default.statements)
+        self.generic_visit(node, context)
+
     def visit_CalibrationStatement(
         self, node: ast.CalibrationStatement, context: None = None
     ) -> None: