Add Expr support to ControlFlowOp representation

qiskit/circuit/classical/expr/__init__.py

@@ -137,6 +137,11 @@
 that they wish to handle.  Any non-overridden methods will call :meth:`~ExprVisitor.visit_generic`,
 which unless overridden will raise a ``RuntimeError`` to ensure that you are aware if new nodes
 have been added to the expression tree that you are not yet handling.
+
+For the convenience of simple visitors that only need to inspect the variables in an expression and
+not the general structure, the iterator method :func:`iter_vars` is provided.
+
+.. autofunction:: iter_vars
 """
 
 __all__ = [
@@ -147,6 +152,7 @@
     "Unary",
     "Binary",
     "ExprVisitor",
+    "iter_vars",
     "lift",
     "cast",
     "bit_not",
@@ -166,7 +172,7 @@
 ]
 
 from .expr import Expr, Var, Value, Cast, Unary, Binary
-from .visitors import ExprVisitor
+from .visitors import ExprVisitor, iter_vars
 from .constructors import (
     lift,
     cast,

qiskit/circuit/classical/expr/constructors.py

@@ -116,7 +116,7 @@ def lift_legacy_condition(
 
             lifted = expr.lift_legacy_condition(instr.condition)
     """
-    from qiskit.circuit import Clbit
+    from qiskit.circuit import Clbit  # pylint: disable=cyclic-import
 
     target, value = condition
     if isinstance(target, Clbit):
@@ -159,7 +159,7 @@ def lift(value: typing.Any, /, type: types.Type | None = None) -> Expr:
         if type is not None:
             raise ValueError("use 'cast' to cast existing expressions, not 'lift'")
         return value
-    from qiskit.circuit import Clbit, ClassicalRegister
+    from qiskit.circuit import Clbit, ClassicalRegister  # pylint: disable=cyclic-import
 
     inferred: types.Type
     if value is True or value is False or isinstance(value, Clbit):

qiskit/circuit/classical/expr/expr.py

@@ -37,9 +37,13 @@
 if typing.TYPE_CHECKING:
     import qiskit
 
+
 _T_co = typing.TypeVar("_T_co", covariant=True)
 
 
+# If adding nodes, remember to update `visitors.ExprVisitor` as well.
+
+
 class Expr(abc.ABC):
     """Root base class of all nodes in the expression tree.  The base case should never be
     instantiated directly.

qiskit/circuit/classical/expr/visitors.py

@@ -16,6 +16,7 @@
 
 __all__ = [
     "ExprVisitor",
+    "iter_vars",
 ]
 
 import typing
@@ -30,6 +31,7 @@ class ExprVisitor(typing.Generic[_T_co]):
     the ``visit_*`` methods that they are able to handle, and should be organised such that
     non-existent methods will never be called."""
 
+    # The method names are self-explanatory and docstrings would just be noise.
     # pylint: disable=missing-function-docstring
 
     __slots__ = ()
@@ -51,3 +53,46 @@ def visit_binary(self, node: expr.Binary, /) -> _T_co:  # pragma: no cover
 
     def visit_cast(self, node: expr.Cast, /) -> _T_co:  # pragma: no cover
         return self.visit_generic(node)
+
+
+class _VarWalkerImpl(ExprVisitor[typing.Iterable[expr.Var]]):
+    __slots__ = ()
+
+    def visit_var(self, node, /):
+        yield node
+
+    def visit_value(self, node, /):
+        yield from ()
+
+    def visit_unary(self, node, /):
+        yield from node.operand.accept(self)
+
+    def visit_binary(self, node, /):
+        yield from node.left.accept(self)
+        yield from node.right.accept(self)
+
+    def visit_cast(self, node, /):
+        yield from node.operand.accept(self)
+
+
+_VAR_WALKER = _VarWalkerImpl()
+
+
+def iter_vars(node: expr.Expr) -> typing.Iterator[expr.Var]:
+    """Get an iterator over the :class:`~.expr.Var` nodes referenced at any level in the given
+    :class:`~.expr.Expr`.
+
+    Examples:
+        Print out the name of each :class:`.ClassicalRegister` encountered::
+
+            from qiskit.circuit import ClassicalRegister
+            from qiskit.circuit.classical import expr
+
+            cr1 = ClassicalRegister(3, "a")
+            cr2 = ClassicalRegister(3, "b")
+
+            for node in expr.iter_vars(expr.bit_and(expr.bit_not(cr1), cr2)):
+                if isinstance(node.var, ClassicalRegister):
+                    print(node.var.name)
+    """
+    yield from node.accept(_VAR_WALKER)

qiskit/circuit/controlflow/__init__.py

@@ -13,6 +13,7 @@
 """Instruction sub-classes for dynamic circuits."""
 
 
+from ._builder_utils import condition_resources, node_resources, LegacyResources
 from .control_flow import ControlFlowOp
 from .continue_loop import ContinueLoopOp
 from .break_loop import BreakLoopOp

qiskit/circuit/controlflow/_builder_utils.py

@@ -12,14 +12,92 @@
 
 """Private utility functions that are used by the builder interfaces."""
 
-from typing import Iterable, Tuple, Set
+from __future__ import annotations
 
+import dataclasses
+from typing import Iterable, Tuple, Set, Union, TypeVar
+
+from qiskit.circuit.classical import expr, types
 from qiskit.circuit.exceptions import CircuitError
 from qiskit.circuit.quantumcircuit import QuantumCircuit
 from qiskit.circuit.register import Register
-from qiskit.circuit.classicalregister import ClassicalRegister
+from qiskit.circuit.classicalregister import ClassicalRegister, Clbit
 from qiskit.circuit.quantumregister import QuantumRegister
 
+_ConditionT = TypeVar(
+    "_ConditionT", bound=Union[Tuple[ClassicalRegister, int], Tuple[Clbit, int], expr.Expr]
+)
+
+
+def validate_condition(condition: _ConditionT) -> _ConditionT:
+    """Validate that a condition is in a valid format and return it, but raise if it is invalid.
+
+    Args:
+        condition: the condition to be tested for validity.  Must be either the legacy 2-tuple
+            format, or a :class:`~.expr.Expr` that has `Bool` type.
+
+    Raises:
+        CircuitError: if the condition is not in a valid format.
+
+    Returns:
+        The same condition as passed, if it was valid.
+    """
+    if isinstance(condition, expr.Expr):
+        if condition.type.kind is not types.Bool:
+            raise CircuitError(
+                "Classical conditions must be expressions with the type 'Bool()',"
+                f" not '{condition.type}'."
+            )
+        return condition
+    try:
+        bits, value = condition
+        if isinstance(bits, (ClassicalRegister, Clbit)) and isinstance(value, int):
+            return (bits, value)
+    except (TypeError, ValueError):
+        pass
+    raise CircuitError(
+        "A classical condition should be a 2-tuple of `(ClassicalRegister | Clbit, int)`,"
+        f" but received '{condition!r}'."
+    )
+
+
+@dataclasses.dataclass
+class LegacyResources:
+    """A pair of the :class:`.Clbit` and :class:`.ClassicalRegister` resources used by some other
+    object (such as a legacy condition or :class:`.expr.Expr` node)."""
+
+    clbits: tuple[Clbit, ...]
+    cregs: tuple[ClassicalRegister, ...]
+
+
+def node_resources(node: expr.Expr) -> LegacyResources:
+    """Get the legacy classical resources (:class:`.Clbit` and :class:`.ClassicalRegister`)
+    referenced by an :class:`~.expr.Expr`."""
+    # It's generally convenient for us to ensure that the resources are returned in some
+    # deterministic order.  This uses the ordering of 'dict' objects to fake out an ordered set.
+    clbits = {}
+    cregs = {}
+    for var in expr.iter_vars(node):
+        if isinstance(var.var, Clbit):
+            clbits[var.var] = None
+        elif isinstance(var.var, ClassicalRegister):
+            clbits.update((bit, None) for bit in var.var)
+            cregs[var.var] = None
+    return LegacyResources(tuple(clbits), tuple(cregs))
+
+
+def condition_resources(
+    condition: tuple[ClassicalRegister, int] | tuple[Clbit, int] | expr.Expr
+) -> LegacyResources:
+    """Get the legacy classical resources (:class:`.Clbit` and :class:`.ClassicalRegister`)
+    referenced by a legacy condition or an :class:`~.expr.Expr`."""
+    if isinstance(condition, expr.Expr):
+        return node_resources(condition)
+    target, _ = condition
+    if isinstance(target, ClassicalRegister):
+        return LegacyResources(tuple(target), (target,))
+    return LegacyResources((target,), ())
+
 
 def partition_registers(
     registers: Iterable[Register],

qiskit/circuit/controlflow/builder.py

@@ -30,7 +30,7 @@
 from qiskit.circuit.quantumregister import Qubit, QuantumRegister
 from qiskit.circuit.register import Register
 
-from .condition import condition_registers
+from ._builder_utils import condition_resources
 
 if typing.TYPE_CHECKING:
     import qiskit  # pylint: disable=cyclic-import
@@ -447,7 +447,7 @@ def build(
                         self.add_register(register)
                         out.add_register(register)
             if getattr(instruction.operation, "condition", None) is not None:
-                for register in condition_registers(instruction.operation.condition):
+                for register in condition_resources(instruction.operation.condition).cregs:
                     if register not in self.registers:
                         self.add_register(register)
                         out.add_register(register)

qiskit/circuit/controlflow/if_else.py

@@ -10,20 +10,26 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
-"Circuit operation representing an ``if/else`` statement."
+"""Circuit operation representing an ``if/else`` statement."""
 
+from __future__ import annotations
 
 from typing import Optional, Tuple, Union, Iterable
 import itertools
 
 from qiskit.circuit import ClassicalRegister, Clbit, QuantumCircuit
+from qiskit.circuit.classical import expr
 from qiskit.circuit.instructionset import InstructionSet
 from qiskit.circuit.exceptions import CircuitError
 
 from .builder import ControlFlowBuilderBlock, InstructionPlaceholder, InstructionResources
-from .condition import validate_condition, condition_bits, condition_registers
 from .control_flow import ControlFlowOp
-from ._builder_utils import partition_registers, unify_circuit_resources
+from ._builder_utils import (
+    partition_registers,
+    unify_circuit_resources,
+    validate_condition,
+    condition_resources,
+)
 
 
 # This is just an indication of what's actually meant to be the public API.
@@ -71,10 +77,10 @@ class IfElseOp(ControlFlowOp):
 
     def __init__(
         self,
-        condition: Tuple[Union[ClassicalRegister, Clbit], int],
+        condition: tuple[ClassicalRegister, int] | tuple[Clbit, int] | expr.Expr,
         true_body: QuantumCircuit,
-        false_body: Optional[QuantumCircuit] = None,
-        label: Optional[str] = None,
+        false_body: QuantumCircuit | None = None,
+        label: str | None = None,
     ):
         # Type checking generally left to @params.setter, but required here for
         # finding num_qubits and num_clbits.
@@ -364,9 +370,10 @@ def in_loop(self) -> bool:
         return self._in_loop
 
     def __enter__(self):
+        resources = condition_resources(self._condition)
         self._circuit._push_scope(
-            clbits=condition_bits(self._condition),
-            registers=condition_registers(self._condition),
+            clbits=resources.clbits,
+            registers=resources.cregs,
             allow_jumps=self._in_loop,
         )
         return ElseContext(self)