decompose_preserving_structure (#3899)

Part of #3634.

In order to send `CircuitOperation`s to QCS, we must be able to decompose the contents of those `CircuitOperation`s without decomposing the `CircuitOperation` itself. `decompose_preserving_structure` provides this behavior.

cirq/protocols/decompose_protocol.py

@@ -38,8 +38,6 @@
 if TYPE_CHECKING:
     import cirq
 
-TValue = TypeVar('TValue')
-
 TDefault = TypeVar('TDefault')
 
 TError = TypeVar('TError', bound=Exception)
@@ -140,20 +138,21 @@ def decompose(
     intercepting_decomposer: Optional[OpDecomposer] = None,
     fallback_decomposer: Optional[OpDecomposer] = None,
     keep: Optional[Callable[['cirq.Operation'], bool]] = None,
-    on_stuck_raise: Optional[Union[TError, Callable[['cirq.Operation'], TError]]],
+    on_stuck_raise: Union[None, TError, Callable[['cirq.Operation'], Optional[TError]]],
 ) -> List['cirq.Operation']:
     pass
 
 
 def decompose(
-    val: TValue,
+    val: Any,
     *,
     intercepting_decomposer: Optional[OpDecomposer] = None,
     fallback_decomposer: Optional[OpDecomposer] = None,
     keep: Optional[Callable[['cirq.Operation'], bool]] = None,
     on_stuck_raise: Union[
         None, Exception, Callable[['cirq.Operation'], Union[None, Exception]]
     ] = _value_error_describing_bad_operation,
+    preserve_structure: bool = False,
 ) -> List['cirq.Operation']:
     """Recursively decomposes a value into `cirq.Operation`s meeting a criteria.
 
@@ -182,6 +181,9 @@ def decompose(
             returns `None`, non-decomposable operations are simply silently
             kept. `on_stuck_raise` defaults to a `ValueError` describing the
             unwanted non-decomposable operation.
+        preserve_structure: Prevents subcircuits (i.e. `CircuitOperation`s)
+            from being decomposed, but decomposes their contents. If this is
+            True, 'intercepting_decomposer' cannot be specified.
 
     Returns:
         A list of operations that the given value was decomposed into. If
@@ -209,6 +211,16 @@ def decompose(
             "acceptable to keep."
         )
 
+    if preserve_structure:
+        if intercepting_decomposer is not None:
+            raise ValueError('Cannot specify intercepting_decomposer while preserving structure.')
+        return _decompose_preserving_structure(
+            val,
+            fallback_decomposer=fallback_decomposer,
+            keep=keep,
+            on_stuck_raise=on_stuck_raise,
+        )
+
     def try_op_decomposer(val: Any, decomposer: Optional[OpDecomposer]) -> DecomposeResult:
         if decomposer is None or not isinstance(val, ops.Operation):
             return None
@@ -317,10 +329,7 @@ def decompose_once_with_qubits(val: Any, qubits: Iterable['cirq.Qid']) -> List['
 def decompose_once_with_qubits(
     val: Any,
     qubits: Iterable['cirq.Qid'],
-    # NOTE: should be TDefault instead of Any, but
-    # mypy has false positive errors when setting
-    # default to None.
-    default: Any,
+    default: Optional[TDefault],
 ) -> Union[TDefault, List['cirq.Operation']]:
     pass
 
@@ -387,3 +396,60 @@ def _try_decompose_into_operations_and_qubits(
         return result, qubits, tuple(qid_shape_dict[q] for q in qubits)
 
     return None, (), ()
+
+
+def _decompose_preserving_structure(
+    val: Any,
+    *,
+    fallback_decomposer: Optional[OpDecomposer] = None,
+    keep: Optional[Callable[['cirq.Operation'], bool]] = None,
+    on_stuck_raise: Union[
+        None, Exception, Callable[['cirq.Operation'], Union[None, Exception]]
+    ] = _value_error_describing_bad_operation,
+) -> List['cirq.Operation']:
+    """Preserves structure (e.g. subcircuits) while decomposing ops.
+
+    This can be used to reduce a circuit to a particular gateset without
+    increasing its serialization size. See tests for examples.
+    """
+
+    # This method provides a generated 'keep' to its decompose() calls.
+    # If the user-provided keep is not set, on_stuck_raise must be unset to
+    # ensure that failure to decompose does not generate errors.
+    on_stuck_raise = on_stuck_raise if keep is not None else None
+
+    from cirq.circuits import CircuitOperation, FrozenCircuit
+
+    visited_fcs = set()
+
+    def keep_structure(op: 'cirq.Operation'):
+        circuit = getattr(op.untagged, 'circuit', None)
+        if circuit is not None:
+            return circuit in visited_fcs
+        if keep is not None and keep(op):
+            return True
+
+    def dps_interceptor(op: 'cirq.Operation'):
+        if not isinstance(op.untagged, CircuitOperation):
+            return NotImplemented
+
+        new_fc = FrozenCircuit(
+            decompose(
+                op.untagged.circuit,
+                intercepting_decomposer=dps_interceptor,
+                fallback_decomposer=fallback_decomposer,
+                keep=keep_structure,
+                on_stuck_raise=on_stuck_raise,
+            )
+        )
+        visited_fcs.add(new_fc)
+        new_co = op.untagged.replace(circuit=new_fc)
+        return new_co if not op.tags else new_co.with_tags(*op.tags)
+
+    return decompose(
+        val,
+        intercepting_decomposer=dps_interceptor,
+        fallback_decomposer=fallback_decomposer,
+        keep=keep_structure,
+        on_stuck_raise=on_stuck_raise,
+    )

cirq/protocols/decompose_protocol_test.py

@@ -222,3 +222,100 @@ def test_decompose_intercept():
         intercepting_decomposer=lambda _: NotImplemented,
     )
     assert actual == [cirq.CNOT(a, b), cirq.CNOT(b, a), cirq.CNOT(a, b)]
+
+
+def test_decompose_preserving_structure():
+    a, b = cirq.LineQubit.range(2)
+    fc1 = cirq.FrozenCircuit(cirq.SWAP(a, b), cirq.FSimGate(0.1, 0.2).on(a, b))
+    cop1_1 = cirq.CircuitOperation(fc1).with_tags('test_tag')
+    cop1_2 = cirq.CircuitOperation(fc1).with_qubit_mapping({a: b, b: a})
+    fc2 = cirq.FrozenCircuit(cirq.X(a), cop1_1, cop1_2)
+    cop2 = cirq.CircuitOperation(fc2)
+
+    circuit = cirq.Circuit(cop2, cirq.measure(a, b, key='m'))
+    actual = cirq.Circuit(cirq.decompose(circuit, preserve_structure=True))
+
+    # This should keep the CircuitOperations but decompose their SWAPs.
+    fc1_decomp = cirq.FrozenCircuit(cirq.decompose(fc1))
+    expected = cirq.Circuit(
+        cirq.CircuitOperation(
+            cirq.FrozenCircuit(
+                cirq.X(a),
+                cirq.CircuitOperation(fc1_decomp).with_tags('test_tag'),
+                cirq.CircuitOperation(fc1_decomp).with_qubit_mapping({a: b, b: a}),
+            )
+        ),
+        cirq.measure(a, b, key='m'),
+    )
+    assert actual == expected
+
+
+def test_decompose_preserving_structure_forwards_args():
+    a, b = cirq.LineQubit.range(2)
+    fc1 = cirq.FrozenCircuit(cirq.SWAP(a, b), cirq.FSimGate(0.1, 0.2).on(a, b))
+    cop1_1 = cirq.CircuitOperation(fc1).with_tags('test_tag')
+    cop1_2 = cirq.CircuitOperation(fc1).with_qubit_mapping({a: b, b: a})
+    fc2 = cirq.FrozenCircuit(cirq.X(a), cop1_1, cop1_2)
+    cop2 = cirq.CircuitOperation(fc2)
+
+    circuit = cirq.Circuit(cop2, cirq.measure(a, b, key='m'))
+
+    def keep_func(op: 'cirq.Operation'):
+        # Only decompose SWAP and X.
+        return not isinstance(op.gate, (cirq.SwapPowGate, cirq.XPowGate))
+
+    def x_to_hzh(op: 'cirq.Operation'):
+        if isinstance(op.gate, cirq.XPowGate) and op.gate.exponent == 1:
+            return [
+                cirq.H(*op.qubits),
+                cirq.Z(*op.qubits),
+                cirq.H(*op.qubits),
+            ]
+
+    actual = cirq.Circuit(
+        cirq.decompose(
+            circuit,
+            keep=keep_func,
+            fallback_decomposer=x_to_hzh,
+            preserve_structure=True,
+        ),
+    )
+
+    # This should keep the CircuitOperations but decompose their SWAPs.
+    fc1_decomp = cirq.FrozenCircuit(
+        cirq.decompose(
+            fc1,
+            keep=keep_func,
+            fallback_decomposer=x_to_hzh,
+        )
+    )
+    expected = cirq.Circuit(
+        cirq.CircuitOperation(
+            cirq.FrozenCircuit(
+                cirq.H(a),
+                cirq.Z(a),
+                cirq.H(a),
+                cirq.CircuitOperation(fc1_decomp).with_tags('test_tag'),
+                cirq.CircuitOperation(fc1_decomp).with_qubit_mapping({a: b, b: a}),
+            )
+        ),
+        cirq.measure(a, b, key='m'),
+    )
+    assert actual == expected
+
+
+def test_decompose_preserving_structure_no_interceptor():
+    a, b = cirq.LineQubit.range(2)
+    fc1 = cirq.FrozenCircuit(cirq.SWAP(a, b), cirq.FSimGate(0.1, 0.2).on(a, b))
+    cop1_1 = cirq.CircuitOperation(fc1).with_tags('test_tag')
+    cop1_2 = cirq.CircuitOperation(fc1).with_qubit_mapping({a: b, b: a})
+    fc2 = cirq.FrozenCircuit(cirq.X(a), cop1_1, cop1_2)
+    cop2 = cirq.CircuitOperation(fc2)
+
+    circuit = cirq.Circuit(cop2, cirq.measure(a, b, key='m'))
+    with pytest.raises(ValueError, match='Cannot specify intercepting_decomposer'):
+        cirq.decompose(
+            circuit,
+            intercepting_decomposer=lambda x: [],
+            preserve_structure=True,
+        )