Scoping for control keys

cirq-core/cirq/__init__.py

@@ -597,6 +597,7 @@
     with_key_path,
     with_key_path_prefix,
     with_measurement_key_mapping,
+    with_rescoped_keys,
 )
 
 from cirq.ion import (

cirq-core/cirq/circuits/circuit.py

@@ -908,17 +908,21 @@ def _with_key_path_(self, path: Tuple[str, ...]):
             [protocols.with_key_path(moment, path) for moment in self.moments]
         )
 
-    def _with_key_path_prefix_(
+    def _with_key_path_prefix_(self, prefix: Tuple[str, ...]):
+        return self._with_sliced_moments(
+            [protocols.with_key_path_prefix(moment, prefix) for moment in self.moments]
+        )
+
+    def _with_rescoped_keys_(
         self,
         path: Tuple[str, ...],
-        local_keys: FrozenSet[value.MeasurementKey],
-        extern_keys: FrozenSet[value.MeasurementKey],
+        bindable_keys: FrozenSet['cirq.MeasurementKey'],
     ):
         moments = []
         for moment in self.moments:
-            moment_keys = protocols.measurement_key_objs(moment)
-            moments.append(protocols.with_key_path_prefix(moment, path, local_keys, extern_keys))
-            local_keys = local_keys.union(moment_keys)
+            new_moment = protocols.with_rescoped_keys(moment, path, bindable_keys)
+            moments.append(new_moment)
+            bindable_keys |= protocols.measurement_key_objs(new_moment)
         return self._with_sliced_moments(moments)
 
     def _qid_shape_(self) -> Tuple[int, ...]:
@@ -1534,10 +1538,8 @@ def factorize(self: CIRCUIT_TYPE) -> Iterable[CIRCUIT_TYPE]:
         )
 
     def _control_keys_(self) -> FrozenSet[value.MeasurementKey]:
-        all_keys: FrozenSet[value.MeasurementKey] = frozenset()
-        for op in self.all_operations():
-            all_keys = all_keys.union(protocols.control_keys(op))
-        return all_keys.difference(protocols.measurement_key_objs(self))
+        controls = frozenset(k for op in self.all_operations() for k in protocols.control_keys(op))
+        return controls - protocols.measurement_key_objs(self)
 
 
 def _overlap_collision_time(

cirq-core/cirq/circuits/circuit_operation.py

@@ -84,7 +84,12 @@ class CircuitOperation(ops.Operation):
             initialized to strings for numbers in `range(repetitions)`.
         parent_path: A tuple of identifiers for any parent CircuitOperations
             containing this one.
-        extern_keys: The set of measurement keys defined at extern scope.
+        extern_keys: The set of measurement keys defined at extern scope. The
+            values here are used by decomposition and simulation routines to
+            cache which external measurement keys exist as possible binding
+            targets for unbound `ClassicallyControlledOperation` keys. This
+            field is not intended to be set or changed manually, and should be
+            empty in circuits that aren't in the middle of decomposition.
     """
 
     _hash: Optional[int] = dataclasses.field(default=None, init=False)
@@ -99,7 +104,7 @@ class CircuitOperation(ops.Operation):
     param_resolver: study.ParamResolver = study.ParamResolver()
     repetition_ids: Optional[List[str]] = dataclasses.field(default=None)
     parent_path: Tuple[str, ...] = dataclasses.field(default_factory=tuple)
-    extern_keys: FrozenSet[value.MeasurementKey] = dataclasses.field(default_factory=frozenset)
+    extern_keys: FrozenSet['cirq.MeasurementKey'] = dataclasses.field(default_factory=frozenset)
 
     def __post_init__(self):
         if not isinstance(self.circuit, circuits.FrozenCircuit):
@@ -188,9 +193,7 @@ def _measurement_key_objs_(self) -> AbstractSet[value.MeasurementKey]:
                     for repetition_id in self.repetition_ids
                     for key in circuit_keys
                 }
-            circuit_keys = {
-                protocols.with_key_path_prefix(key, self.parent_path) for key in circuit_keys
-            }
+            circuit_keys = {key.with_key_path_prefix(*self.parent_path) for key in circuit_keys}
             object.__setattr__(
                 self,
                 '_cached_measurement_key_objs',
@@ -206,7 +209,7 @@ def _measurement_key_names_(self) -> AbstractSet[str]:
 
     def _control_keys_(self) -> AbstractSet[value.MeasurementKey]:
         if not protocols.control_keys(self.circuit):
-            return set()
+            return frozenset()
         return protocols.control_keys(self.mapped_circuit())
 
     def _parameter_names_(self) -> AbstractSet[str]:
@@ -244,10 +247,10 @@ def mapped_circuit(self, deep: bool = False) -> 'cirq.Circuit':
                 circuit = circuit * abs(self.repetitions)
             else:
                 circuit = circuits.Circuit(
-                    protocols.with_key_path_prefix(circuit, (rep,)) for rep in self.repetition_ids
+                    protocols.with_rescoped_keys(circuit, (rep,)) for rep in self.repetition_ids
                 )
-        circuit = protocols.with_key_path_prefix(
-            circuit, self.parent_path, extern_keys=self.extern_keys
+        circuit = protocols.with_rescoped_keys(
+            circuit, self.parent_path, bindable_keys=self.extern_keys
         )
         if deep:
             circuit = circuit.map_operations(
@@ -438,22 +441,23 @@ def __pow__(self, power: int) -> 'CircuitOperation':
     def _with_key_path_(self, path: Tuple[str, ...]):
         return dataclasses.replace(self, parent_path=path)
 
-    def _with_key_path_prefix_(
+    def _with_key_path_prefix_(self, prefix: Tuple[str, ...]):
+        return dataclasses.replace(self, parent_path=prefix + self.parent_path)
+
+    def _with_rescoped_keys_(
         self,
         path: Tuple[str, ...],
-        local_keys: FrozenSet[value.MeasurementKey],
-        extern_keys: FrozenSet[value.MeasurementKey],
+        bindable_keys: FrozenSet['cirq.MeasurementKey'],
     ):
-        new_keys = set(extern_keys)
-        for key in local_keys:
-            new_keys.add(protocols.with_key_path_prefix(key, path))
-        for key in self.extern_keys:
-            new_keys.add(protocols.with_key_path_prefix(key, path))
-        return dataclasses.replace(
-            self,
-            parent_path=path + self.parent_path,
-            extern_keys=frozenset(new_keys),
-        )
+        # The following line prevents binding to measurement keys in previous repeated subcircuits
+        # "just because their repetition ids matched". If we eventually decide to change that
+        # requirement and allow binding across subcircuits (possibly conditionally upon the key or
+        # the subcircuit having some 'allow_cross_circuit_binding' field set), this is the line to
+        # change or remove.
+        bindable_keys = frozenset(k for k in bindable_keys if len(k.path) <= len(path))
+        bindable_keys |= {k.with_key_path_prefix(*path) for k in self.extern_keys}
+        path += self.parent_path
+        return dataclasses.replace(self, parent_path=path, extern_keys=bindable_keys)
 
     def with_key_path(self, path: Tuple[str, ...]):
         return self._with_key_path_(path)
@@ -550,11 +554,8 @@ def with_measurement_key_mapping(self, key_map: Dict[str, str]) -> 'CircuitOpera
             if k_new != k:
                 new_map[k] = k_new
         new_op = self.replace(measurement_key_map=new_map)
-        if (
-            len(protocols.measurement_key_objs(new_op)) != len(protocols.measurement_key_objs(self))
-            or len(protocols.control_keys(new_op)) != len(protocols.control_keys(self))
-            or len(protocols.measurement_keys_touched(new_op))
-            != len(protocols.measurement_keys_touched(self))
+        if len(protocols.measurement_keys_touched(new_op)) != len(
+            protocols.measurement_keys_touched(self)
         ):
             raise ValueError(
                 f'Collision in measurement key map composition. Original map:\n'

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -844,19 +844,6 @@ def test_keys_conflict_no_repetitions():
         _ = op2.mapped_circuit(deep=True)
 
 
-def test_keys_conflict_with_repetitions():
-    q = cirq.LineQubit(0)
-    op1 = cirq.CircuitOperation(
-        cirq.FrozenCircuit(
-            cirq.measure(q, key='A'),
-        ),
-        repetitions=2,
-    )
-    op2 = cirq.CircuitOperation(cirq.FrozenCircuit(op1, op1), repetitions=2)
-    with pytest.raises(ValueError, match='Conflicting measurement keys found: 0:0:A'):
-        _ = op2.mapped_circuit(deep=True)
-
-
 def test_keys_conflict_locally():
     q = cirq.LineQubit(0)
     op1 = cirq.measure(q, key='A')

cirq-core/cirq/ops/classically_controlled_operation.py

@@ -169,31 +169,31 @@ def not_zero(measurement):
     def _with_measurement_key_mapping_(
         self, key_map: Dict[str, str]
     ) -> 'ClassicallyControlledOperation':
-        return ClassicallyControlledOperation(
-            self._sub_operation,
-            [protocols.with_measurement_key_mapping(k, key_map) for k in self._control_keys],
+        sub_operation = protocols.with_measurement_key_mapping(self._sub_operation, key_map)
+        sub_operation = self._sub_operation if sub_operation is NotImplemented else sub_operation
+        return sub_operation.with_classical_controls(
+            *[protocols.with_measurement_key_mapping(k, key_map) for k in self._control_keys]
         )
 
-    def _with_key_path_prefix_(
+    def _with_key_path_prefix_(self, path: Tuple[str, ...]) -> 'ClassicallyControlledOperation':
+        keys = [protocols.with_key_path_prefix(k, path) for k in self._control_keys]
+        return self._sub_operation.with_classical_controls(*keys)
+
+    def _with_rescoped_keys_(
         self,
         path: Tuple[str, ...],
-        local_keys: FrozenSet[value.MeasurementKey],
-        extern_keys: FrozenSet[value.MeasurementKey],
+        bindable_keys: FrozenSet['cirq.MeasurementKey'],
     ) -> 'ClassicallyControlledOperation':
         def map_key(key: value.MeasurementKey) -> value.MeasurementKey:
-            if key in local_keys:
-                return protocols.with_key_path_prefix(key, path)
-            for i in range(len(path)):
-                back_path = path[0 : len(path) - i]
-                new_key = protocols.with_key_path_prefix(key, back_path)
-                if new_key in extern_keys:
+            for i in range(len(path) + 1):
+                back_path = path[: len(path) - i]
+                new_key = key.with_key_path_prefix(*back_path)
+                if new_key in bindable_keys:
                     return new_key
             return key
 
-        return ClassicallyControlledOperation(
-            self._sub_operation,
-            [map_key(k) for k in self._control_keys],
-        )
+        sub_operation = protocols.with_rescoped_keys(self._sub_operation, path, bindable_keys)
+        return sub_operation.with_classical_controls(*[map_key(k) for k in self._control_keys])
 
     def _control_keys_(self) -> FrozenSet[value.MeasurementKey]:
         return frozenset(self._control_keys).union(protocols.control_keys(self._sub_operation))

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -429,6 +429,7 @@ def test_scope_local():
 """,
         use_unicode_characters=True,
     )
+    assert circuit == cirq.Circuit(cirq.decompose(outer_subcircuit))
 
 
 def test_scope_extern():
@@ -471,6 +472,45 @@ def test_scope_extern():
 """,
         use_unicode_characters=True,
     )
+    assert circuit == cirq.Circuit(cirq.decompose(outer_subcircuit))
+
+
+def test_scope_extern_wrapping_with_non_repeating_subcircuits():
+    def wrap(*ops):
+        return cirq.CircuitOperation(cirq.FrozenCircuit(*ops))
+
+    def wrap_frozen(*ops):
+        return cirq.FrozenCircuit(wrap(*ops))
+
+    q = cirq.LineQubit(0)
+    inner = wrap_frozen(
+        wrap(cirq.measure(q, key='a')),
+        wrap(cirq.X(q).with_classical_controls('b')),
+    )
+    middle = wrap_frozen(
+        wrap(cirq.measure(q, key=cirq.MeasurementKey('b'))),
+        wrap(cirq.CircuitOperation(inner, repetitions=2)),
+    )
+    outer_subcircuit = cirq.CircuitOperation(middle, 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:b', '0:b', '1:b', '1:b']
+    assert not cirq.control_keys(outer_subcircuit)
+    assert not cirq.control_keys(circuit)
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ─────M───M('0:0:a')───X───M('0:1:a')───X───M───M('1:0:a')───X───M('1:1:a')───X───
+        ║                ║                ║   ║                ║                ║
+0:b: ═══@════════════════^════════════════^═══╬════════════════╬════════════════╬═══
+                                              ║                ║                ║
+1:b: ═════════════════════════════════════════@════════════════^════════════════^═══
+""",
+        use_unicode_characters=True,
+    )
+    assert circuit == cirq.Circuit(cirq.decompose(outer_subcircuit))
 
 
 def test_scope_root():
@@ -513,6 +553,7 @@ def test_scope_root():
 """,
         use_unicode_characters=True,
     )
+    assert circuit == cirq.Circuit(cirq.decompose(outer_subcircuit))
 
 
 def test_scope_extern_mismatch():
@@ -555,31 +596,7 @@ def test_scope_extern_mismatch():
 """,
         use_unicode_characters=True,
     )
-
-
-def test_scope_conflict():
-    q = cirq.LineQubit(0)
-    inner = cirq.Circuit(
-        cirq.measure(q, key='a'),
-        cirq.X(q).with_classical_controls('a'),
-    )
-    middle = cirq.Circuit(
-        cirq.measure(q, key=cirq.MeasurementKey('a', ('0',))),
-        cirq.CircuitOperation(inner.freeze(), repetitions=2),
-    )
-    op = cirq.CircuitOperation(middle.freeze(), repetitions=2)
-    assert not cirq.control_keys(op)
-    cirq.testing.assert_has_diagram(
-        cirq.Circuit(op),
-        """
-      [                  [ 0: ───M───X─── ]             ]
-0: ───[ 0: ───M('0:a')───[       ║   ║    ]──────────── ]────────────
-      [                  [ a: ═══@═══^═══ ](loops=2)    ](loops=2)
-""",
-        use_unicode_characters=True,
-    )
-    with pytest.raises(ValueError, match='Conflicting measurement keys found: 0:0:a'):
-        _ = op.mapped_circuit(deep=True)
+    assert circuit == cirq.Circuit(cirq.decompose(outer_subcircuit))
 
 
 def test_repr():
@@ -608,3 +625,47 @@ def test_unmeasured_condition():
         ),
     ):
         _ = cirq.Simulator().simulate(bad_circuit)
+
+
+def test_layered_circuit_operations_with_controls_in_between():
+    q = cirq.LineQubit(0)
+    outer_subcircuit = cirq.CircuitOperation(
+        cirq.Circuit(
+            cirq.CircuitOperation(
+                cirq.FrozenCircuit(
+                    cirq.X(q),
+                    cirq.Y(q),
+                )
+            ).with_classical_controls('m')
+        ).freeze()
+    )
+    circuit = outer_subcircuit.mapped_circuit(deep=True)
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(outer_subcircuit),
+        """
+      [ 0: ───[ 0: ───X───Y─── ].with_classical_controls(m)─── ]
+0: ───[       ║                                                ]───
+      [ m: ═══╩═══════════════════════════════════════════════ ]
+      ║
+m: ═══╩════════════════════════════════════════════════════════════
+""",
+        use_unicode_characters=True,
+    )
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───[ 0: ───X───Y─── ].with_classical_controls(m)───
+      ║
+m: ═══╩═══════════════════════════════════════════════
+""",
+        use_unicode_characters=True,
+    )
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(cirq.decompose(outer_subcircuit)),
+        """
+0: ───X───Y───
+      ║   ║
+m: ═══^═══^═══
+""",
+        use_unicode_characters=True,
+    )