Add Expr support to DAGCircuit.compose

The remapping facilities are (except for some esoteric behaviour for
mapping registers to others in a different order) the same as for
`QuantumCircuit`, so it makes sense to unify the handling.  As part of
switching the old `DAGCircuit` classical-resource-mapping methods over
to the new general-purpose forms, this does most of the necessary work
to upgrade `substitute_node_with_dag` as well.

qiskit/circuit/_classical_resource_map.py

@@ -0,0 +1,148 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Shared helper utility for mapping classical resources from one circuit or DAG to another."""
+
+from __future__ import annotations
+
+import typing
+
+from .bit import Bit
+from .classical import expr
+from .classicalregister import ClassicalRegister, Clbit
+
+
+class VariableMapper(expr.ExprVisitor[expr.Expr]):
+    """Stateful helper class that manages the mapping of variables in conditions and expressions.
+
+    This is designed to be used by both :class:`.QuantumCircuit` and :class:`.DAGCircuit` when
+    managing operations that need to map classical resources from one circuit to another.
+
+    The general usage is to initialise this at the start of a many-block mapping operation, then
+    call its :meth:`map_condition`, :meth:`map_target` or :meth:`map_expr` methods as appropriate,
+    which will return the new object that should be used.
+
+    If an ``add_register`` callable is given to the initialiser, the mapper will use it to attempt
+    to add new aliasing registers to the outer circuit object, if there is not already a suitable
+    register for the mapping available in the circuit.  If this parameter is not given, the mapping
+    function will raise an exception of type ``exc_type`` instead.
+    """
+
+    __slots__ = ("target_cregs", "register_map", "bit_map", "add_register", "exc_type")
+
+    def __init__(
+        self,
+        target_cregs: typing.Iterable[ClassicalRegister],
+        bit_map: typing.Mapping[Bit, Bit],
+        add_register: typing.Callable[[ClassicalRegister], None] | None = None,
+        exc_type: typing.Type[Exception] = ValueError,
+    ):
+        self.target_cregs = tuple(target_cregs)
+        self.register_map = {}
+        self.bit_map = bit_map
+        self.add_register = add_register
+        self.exc_type = exc_type
+
+    def _map_register(self, theirs: ClassicalRegister) -> ClassicalRegister:
+        """Map the target's registers to suitable equivalents in the destination, adding an
+        extra one if there's no exact match."""
+        if (mapped_theirs := self.register_map.get(theirs.name)) is not None:
+            return mapped_theirs
+        mapped_bits = [self.bit_map[bit] for bit in theirs]
+        for ours in self.target_cregs:
+            if mapped_bits == list(ours):
+                mapped_theirs = ours
+                break
+        else:
+            if self.add_register is None:
+                raise self.exc_type(
+                    f"Register '{theirs.name}' has no counterpart in the destination."
+                )
+            mapped_theirs = ClassicalRegister(bits=mapped_bits)
+            self.add_register(mapped_theirs)
+        self.register_map[theirs.name] = mapped_theirs
+        return mapped_theirs
+
+    def map_condition(self, condition, /, *, allow_reorder=False):
+        """Map the given ``condition`` so that it only references variables in the destination
+        circuit (as given to this class on initialisation).
+
+        If ``allow_reorder`` is ``True``, then when a legacy condition (the two-tuple form) is made
+        on a register that has a counterpart in the destination with all the same (mapped) bits but
+        in a different order, then that register will be used and the value suitably modified to
+        make the equality condition work.  This is maintaining legacy (tested) behaviour of
+        :meth:`.DAGCircuit.compose`; nowhere else does this, and in general this would require *far*
+        more complex classical rewriting than Terra needs to worry about in the full expression era.
+        """
+        if condition is None:
+            return None
+        if isinstance(condition, expr.Expr):
+            return self.map_expr(condition)
+        target, value = condition
+        if isinstance(target, Clbit):
+            return (self.bit_map[target], value)
+        if not allow_reorder:
+            return (self._map_register(target), value)
+        # This is maintaining the legacy behaviour of `DAGCircuit.compose`.  We don't attempt to
+        # speed-up this lookup with a cache, since that would just make the more standard cases more
+        # annoying to deal with.
+        mapped_bits_order = [self.bit_map[bit] for bit in target]
+        mapped_bits_set = set(mapped_bits_order)
+        for register in self.target_cregs:
+            if mapped_bits_set == set(register):
+                mapped_theirs = register
+                break
+        else:
+            if self.add_register is None:
+                raise self.exc_type(
+                    f"Register '{target.name}' has no counterpart in the destination."
+                )
+            mapped_theirs = ClassicalRegister(bits=mapped_bits_order)
+            self.add_register(mapped_theirs)
+        new_order = {bit: i for i, bit in enumerate(mapped_bits_order)}
+        value_bits = f"{value:0{len(target)}b}"[::-1]  # Little-index-indexed binary bitstring.
+        mapped_value = int("".join(value_bits[new_order[bit]] for bit in mapped_theirs)[::-1], 2)
+        return (mapped_theirs, mapped_value)
+
+    def map_target(self, target, /):
+        """Map the runtime variables in a ``target`` of a :class:`.SwitchCaseOp` to the new circuit,
+        as defined in the ``circuit`` argument of the initialiser of this class."""
+        if isinstance(target, Clbit):
+            return self.bit_map[target]
+        if isinstance(target, ClassicalRegister):
+            return self._map_register(target)
+        return self.map_expr(target)
+
+    def map_expr(self, node: expr.Expr, /) -> expr.Expr:
+        """Map the variables in an :class:`~.expr.Expr` node to the new circuit."""
+        return node.accept(self)
+
+    def visit_var(self, node, /):
+        if isinstance(node.var, Clbit):
+            return expr.Var(self.bit_map[node.var], node.type)
+        if isinstance(node.var, ClassicalRegister):
+            return expr.Var(self._map_register(node.var), node.type)
+        # Defensive against the expansion of the variable system; we don't want to silently do the
+        # wrong thing (which would be `return node` without mapping, right now).
+        raise RuntimeError(f"unhandled variable in 'compose': {node}")  # pragma: no cover
+
+    def visit_value(self, node, /):
+        return expr.Value(node.value, node.type)
+
+    def visit_unary(self, node, /):
+        return expr.Unary(node.op, node.operand.accept(self), node.type)
+
+    def visit_binary(self, node, /):
+        return expr.Binary(node.op, node.left.accept(self), node.right.accept(self), node.type)
+
+    def visit_cast(self, node, /):
+        return expr.Cast(node.operand.accept(self), node.type, implicit=node.implicit)

qiskit/circuit/quantumcircuit.py

@@ -48,6 +48,7 @@
 from qiskit.circuit.parameter import Parameter
 from qiskit.qasm.exceptions import QasmError
 from qiskit.circuit.exceptions import CircuitError
+from . import _classical_resource_map
 from .classical import expr
 from .parameterexpression import ParameterExpression, ParameterValueType
 from .quantumregister import QuantumRegister, Qubit, AncillaRegister, AncillaQubit
@@ -959,7 +960,9 @@ def compose(
                 )
             edge_map.update(zip(other.clbits, dest.cbit_argument_conversion(clbits)))
 
-        variable_mapper = _ComposeVariableMapper(dest, edge_map)
+        variable_mapper = _classical_resource_map.VariableMapper(
+            dest.cregs, edge_map, dest.add_register
+        )
         mapped_instrs: list[CircuitInstruction] = []
         for instr in other.data:
             n_qargs: list[Qubit] = [edge_map[qarg] for qarg in instr.qubits]
@@ -5196,79 +5199,3 @@ def _bit_argument_conversion_scalar(specifier, bit_sequence, bit_set, type_):
         else f"Invalid bit index: '{specifier}' of type '{type(specifier)}'"
     )
     raise CircuitError(message)
-
-
-class _ComposeVariableMapper(expr.ExprVisitor[expr.Expr]):
-    """Stateful helper class that manages the mapping of variables in conditions and expressions to
-    items in the destination ``circuit``.
-
-    This mutates ``circuit`` by adding registers as required."""
-
-    __slots__ = ("circuit", "register_map", "bit_map")
-
-    def __init__(self, circuit, bit_map):
-        self.circuit = circuit
-        self.register_map = {}
-        self.bit_map = bit_map
-
-    def _map_register(self, theirs):
-        """Map the target's registers to suitable equivalents in the destination, adding an
-        extra one if there's no exact match."""
-        if (mapped_theirs := self.register_map.get(theirs.name)) is not None:
-            return mapped_theirs
-        mapped_bits = [self.bit_map[bit] for bit in theirs]
-        for ours in self.circuit.cregs:
-            if mapped_bits == list(ours):
-                mapped_theirs = ours
-                break
-        else:
-            mapped_theirs = ClassicalRegister(bits=mapped_bits)
-            self.circuit.add_register(mapped_theirs)
-        self.register_map[theirs.name] = mapped_theirs
-        return mapped_theirs
-
-    def map_condition(self, condition, /):
-        """Map the given ``condition`` so that it only references variables in the destination
-        circuit (as given to this class on initialisation)."""
-        if condition is None:
-            return None
-        if isinstance(condition, expr.Expr):
-            return self.map_expr(condition)
-        target, value = condition
-        if isinstance(target, Clbit):
-            return (self.bit_map[target], value)
-        return (self._map_register(target), value)
-
-    def map_target(self, target, /):
-        """Map the runtime variables in a ``target`` of a :class:`.SwitchCaseOp` to the new circuit,
-        as defined in the ``circuit`` argument of the initialiser of this class."""
-        if isinstance(target, Clbit):
-            return self.bit_map[target]
-        if isinstance(target, ClassicalRegister):
-            return self._map_register(target)
-        return self.map_expr(target)
-
-    def map_expr(self, node: expr.Expr, /) -> expr.Expr:
-        """Map the variables in an :class:`~.expr.Expr` node to the new circuit."""
-        return node.accept(self)
-
-    def visit_var(self, node, /):
-        if isinstance(node.var, Clbit):
-            return expr.Var(self.bit_map[node.var], node.type)
-        if isinstance(node.var, ClassicalRegister):
-            return expr.Var(self._map_register(node.var), node.type)
-        # Defensive against the expansion of the variable system; we don't want to silently do the
-        # wrong thing (which would be `return node` without mapping, right now).
-        raise CircuitError(f"unhandled variable in 'compose': {node}")  # pragma: no cover
-
-    def visit_value(self, node, /):
-        return expr.Value(node.value, node.type)
-
-    def visit_unary(self, node, /):
-        return expr.Unary(node.op, node.operand.accept(self), node.type)
-
-    def visit_binary(self, node, /):
-        return expr.Binary(node.op, node.left.accept(self), node.right.accept(self), node.type)
-
-    def visit_cast(self, node, /):
-        return expr.Cast(node.operand.accept(self), node.type, implicit=node.implicit)