Allow flattening of subcircuits

cirq-core/cirq/circuits/circuit_operation.py

@@ -108,6 +108,7 @@ class CircuitOperation(ops.Operation):
     repetition_ids: Optional[List[str]] = dataclasses.field(default=None)
     parent_path: Tuple[str, ...] = dataclasses.field(default_factory=tuple)
     extern_keys: FrozenSet['cirq.MeasurementKey'] = dataclasses.field(default_factory=frozenset)
+    flatten_repetitions: bool = False
 
     def __post_init__(self):
         if not isinstance(self.circuit, circuits.FrozenCircuit):
@@ -168,6 +169,7 @@ def __eq__(self, other) -> bool:
             and self.repetitions == other.repetitions
             and self.repetition_ids == other.repetition_ids
             and self.parent_path == other.parent_path
+            and self.flatten_repetitions == other.flatten_repetitions
         )
 
     # Methods for getting post-mapping properties of the contained circuit.
@@ -190,7 +192,7 @@ def _is_measurement_(self) -> bool:
     def _measurement_key_objs_(self) -> AbstractSet['cirq.MeasurementKey']:
         if self._cached_measurement_key_objs is None:
             circuit_keys = protocols.measurement_key_objs(self.circuit)
-            if self.repetition_ids is not None:
+            if self.repetition_ids is not None and not self.flatten_repetitions:
                 circuit_keys = {
                     key.with_key_path_prefix(repetition_id)
                     for repetition_id in self.repetition_ids
@@ -251,7 +253,7 @@ def mapped_circuit(self, deep: bool = False) -> 'cirq.Circuit':
         if self.param_resolver:
             circuit = protocols.resolve_parameters(circuit, self.param_resolver, recursive=False)
         if self.repetition_ids:
-            if not protocols.is_measurement(circuit):
+            if self.flatten_repetitions or not protocols.is_measurement(circuit):
                 circuit = circuit * abs(self.repetitions)
             else:
                 circuit = circuits.Circuit(
@@ -295,6 +297,8 @@ def __repr__(self):
         if self.repetition_ids != self._default_repetition_ids():
             # Default repetition_ids need not be specified.
             args += f'repetition_ids={proper_repr(self.repetition_ids)},\n'
+        if self.flatten_repetitions:
+            args += 'flatten_repetitions=True,\n'
         indented_args = args.replace('\n', '\n    ')
         return f'cirq.CircuitOperation({indented_args[:-4]})'
 
@@ -325,6 +329,8 @@ def dict_str(d: Dict) -> str:
         elif self.repetitions != 1:
             # Only add loops if we haven't added repetition_ids.
             args.append(f'loops={self.repetitions}')
+        if self.flatten_repetitions:
+            args.append('flat')
         if not args:
             return circuit_msg
         return f'{circuit_msg}({", ".join(args)})'
@@ -343,13 +349,14 @@ def __hash__(self):
                         self.param_resolver,
                         self.parent_path,
                         tuple([] if self.repetition_ids is None else self.repetition_ids),
+                        self.flatten_repetitions,
                     )
                 ),
             )
         return self._hash
 
     def _json_dict_(self):
-        return {
+        resp = {
             'circuit': self.circuit,
             'repetitions': self.repetitions,
             # JSON requires mappings to have keys of basic types.
@@ -360,6 +367,9 @@ def _json_dict_(self):
             'repetition_ids': self.repetition_ids,
             'parent_path': self.parent_path,
         }
+        if self.flatten_repetitions:
+            resp['flatten_repetitions'] = True
+        return resp
 
     @classmethod
     def _from_json_dict_(
@@ -371,10 +381,11 @@ def _from_json_dict_(
         param_resolver,
         repetition_ids,
         parent_path=(),
+        flatten_repetitions=False,
         **kwargs,
     ):
         return (
-            cls(circuit)
+            cls(circuit, flatten_repetitions=flatten_repetitions)
             .with_qubit_mapping(dict(qubit_map))
             .with_measurement_key_mapping(measurement_key_map)
             .with_params(param_resolver)

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -457,6 +457,26 @@ def test_string_format():
     ]),
 )"""
     )
+    op6 = cirq.CircuitOperation(fc5, flatten_repetitions=True)
+    assert (
+        repr(op6)
+        == """\
+cirq.CircuitOperation(
+    circuit=cirq.FrozenCircuit([
+        cirq.Moment(
+            cirq.X(cirq.LineQubit(0)),
+            cirq.CircuitOperation(
+                circuit=cirq.FrozenCircuit([
+                    cirq.Moment(
+                        cirq.X(cirq.LineQubit(1)),
+                    ),
+                ]),
+            ),
+        ),
+    ]),
+    flatten_repetitions=True,
+)"""
+    )
 
 
 def test_json_dict():

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -482,6 +482,123 @@ def test_scope_local():
     assert circuit == cirq.Circuit(cirq.decompose(outer_subcircuit))
 
 
+def test_scope_flatten_both():
+    q = cirq.LineQubit(0)
+    inner = cirq.Circuit(
+        cirq.measure(q, key='a'),
+        cirq.X(q).with_classical_controls('a'),
+    )
+    middle = cirq.Circuit(
+        cirq.CircuitOperation(inner.freeze(), repetitions=2, flatten_repetitions=True)
+    )
+    outer_subcircuit = cirq.CircuitOperation(
+        middle.freeze(), repetitions=2, flatten_repetitions=True
+    )
+    circuit = outer_subcircuit.mapped_circuit(deep=True)
+    internal_control_keys = [
+        str(condition) for op in circuit.all_operations() for condition in cirq.control_keys(op)
+    ]
+    assert internal_control_keys == ['a', 'a', 'a', 'a']
+    assert not cirq.control_keys(outer_subcircuit)
+    assert not cirq.control_keys(circuit)
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(outer_subcircuit),
+        """
+      [       [ 0: ───M───X─── ]                   ]
+0: ───[ 0: ───[       ║   ║    ]────────────────── ]──────────────────
+      [       [ a: ═══@═══^═══ ](loops=2, flat)    ](loops=2, flat)
+""",
+        use_unicode_characters=True,
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───M───X───M───X───M───X───M───X───
+      ║   ║   ║   ║   ║   ║   ║   ║
+a: ═══@═══^═══@═══^═══@═══^═══@═══^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_scope_flatten_inner():
+    q = cirq.LineQubit(0)
+    inner = cirq.Circuit(
+        cirq.measure(q, key='a'),
+        cirq.X(q).with_classical_controls('a'),
+    )
+    middle = cirq.Circuit(
+        cirq.CircuitOperation(inner.freeze(), repetitions=2, flatten_repetitions=True)
+    )
+    outer_subcircuit = cirq.CircuitOperation(middle.freeze(), repetitions=2)
+    circuit = outer_subcircuit.mapped_circuit(deep=True)
+    internal_control_keys = [
+        str(condition) for op in circuit.all_operations() for condition in cirq.control_keys(op)
+    ]
+    assert internal_control_keys == ['0:a', '0:a', '1:a', '1:a']
+    assert not cirq.control_keys(outer_subcircuit)
+    assert not cirq.control_keys(circuit)
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(outer_subcircuit),
+        """
+      [       [ 0: ───M───X─── ]                   ]
+0: ───[ 0: ───[       ║   ║    ]────────────────── ]────────────
+      [       [ a: ═══@═══^═══ ](loops=2, flat)    ](loops=2)
+""",
+        use_unicode_characters=True,
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ─────M───X───M───X───M───X───M───X───
+        ║   ║   ║   ║   ║   ║   ║   ║
+0:a: ═══@═══^═══@═══^═══╬═══╬═══╬═══╬═══
+                        ║   ║   ║   ║
+1:a: ═══════════════════@═══^═══@═══^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
+def test_scope_flatten_outer():
+    q = cirq.LineQubit(0)
+    inner = cirq.Circuit(
+        cirq.measure(q, key='a'),
+        cirq.X(q).with_classical_controls('a'),
+    )
+    middle = cirq.Circuit(cirq.CircuitOperation(inner.freeze(), repetitions=2))
+    outer_subcircuit = cirq.CircuitOperation(
+        middle.freeze(), repetitions=2, flatten_repetitions=True
+    )
+    circuit = outer_subcircuit.mapped_circuit(deep=True)
+    internal_control_keys = [
+        str(condition) for op in circuit.all_operations() for condition in cirq.control_keys(op)
+    ]
+    assert internal_control_keys == ['0:a', '1:a', '0:a', '1:a']
+    assert not cirq.control_keys(outer_subcircuit)
+    assert not cirq.control_keys(circuit)
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(outer_subcircuit),
+        """
+      [       [ 0: ───M───X─── ]             ]
+0: ───[ 0: ───[       ║   ║    ]──────────── ]──────────────────
+      [       [ a: ═══@═══^═══ ](loops=2)    ](loops=2, flat)
+""",
+        use_unicode_characters=True,
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ─────M───X───M───X───M───X───M───X───
+        ║   ║   ║   ║   ║   ║   ║   ║
+0:a: ═══@═══^═══╬═══╬═══@═══^═══╬═══╬═══
+                ║   ║           ║   ║
+1:a: ═══════════@═══^═══════════@═══^═══
+""",
+        use_unicode_characters=True,
+    )
+
+
 def test_scope_extern():
     q = cirq.LineQubit(0)
     inner = cirq.Circuit(

cirq-core/cirq/protocols/json_test_data/CircuitOperation.json

@@ -292,7 +292,8 @@
           ]
         },
         "parent_path": [],
-        "repetition_ids": null
+        "repetition_ids": null,
+        "flatten_repetitions": true
       }
     ]
   ]

cirq-core/cirq/protocols/json_test_data/CircuitOperation.repr

@@ -34,4 +34,5 @@ cirq.CircuitOperation(circuit=cirq.FrozenCircuit([
         (cirq.X**sympy.Symbol('theta')).on(cirq.LineQubit(0)),
     ),
 ]),
-param_resolver={sympy.Symbol('theta'): 1.5})]
+param_resolver={sympy.Symbol('theta'): 1.5},
+flatten_repetitions=True)]