Fix mapped measurement keys of repeat_until fields in CircuitOperations (#6881)

* Fix repeat_until scoped keys

Create a _mapped_repeat_until cached property that ensures repeat_until conditions have the path and key mappings of the subcircuit applied during execution time or when calculating control_keys.

* Remove param resolver support from repeat_until

Conditions don't support params, as any unresolved symbols are expected to be measurement keys. Perhaps a future feature to consider, but outside the scope of this PR.

* Undo previous change, just add dummy test for resolve_parameters

Built-in conditions don't support param resolvers, but custom ones plausibly could. May as well keep the functionality, and add a dummy test for coverage.

* Improve tests

* coverage

* fix test after merge from "use_repetition_ids" changes

* check for symbols in repeat_until

* doc

Author: daxfohl

cirq-core/cirq/circuits/circuit_operation.py

@@ -204,11 +204,12 @@ def __init__(
                 )
 
         self._repeat_until = repeat_until
-        if self._repeat_until:
+        mapped_repeat_until = self._mapped_repeat_until
+        if mapped_repeat_until:
             if self._use_repetition_ids or self._repetitions != 1:
                 raise ValueError('Cannot use repetitions with repeat_until')
             if protocols.measurement_key_objs(self._mapped_single_loop()).isdisjoint(
-                self._repeat_until.keys
+                mapped_repeat_until.keys
             ):
                 raise ValueError('Infinite loop: condition is not modified in subcircuit.')
 
@@ -349,8 +350,9 @@ def _control_keys(self) -> FrozenSet['cirq.MeasurementKey']:
             if not protocols.control_keys(self.circuit)
             else protocols.control_keys(self._mapped_single_loop())
         )
-        if self.repeat_until is not None:
-            keys |= frozenset(self.repeat_until.keys) - self._measurement_key_objs_()
+        mapped_repeat_until = self._mapped_repeat_until
+        if mapped_repeat_until is not None:
+            keys |= frozenset(mapped_repeat_until.keys) - self._measurement_key_objs_()
         return keys
 
     def _control_keys_(self) -> FrozenSet['cirq.MeasurementKey']:
@@ -364,11 +366,8 @@ def _parameter_names_(self) -> FrozenSet[str]:
 
     def _parameter_names_generator(self) -> Iterator[str]:
         yield from protocols.parameter_names(self.repetitions)
-        for symbol in protocols.parameter_symbols(self.circuit):
-            for name in protocols.parameter_names(
-                protocols.resolve_parameters(symbol, self.param_resolver, recursive=False)
-            ):
-                yield name
+        yield from protocols.parameter_names(self._mapped_repeat_until)
+        yield from protocols.parameter_names(self._mapped_any_loop)
 
     @cached_property
     def _mapped_any_loop(self) -> 'cirq.Circuit':
@@ -391,6 +390,26 @@ def _mapped_single_loop(self, repetition_id: Optional[str] = None) -> 'cirq.Circ
             circuit, self.parent_path, bindable_keys=self._extern_keys
         )
 
+    @cached_property
+    def _mapped_repeat_until(self) -> Optional['cirq.Condition']:
+        """Applies measurement_key_map, param_resolver, and current scope to repeat_until."""
+        repeat_until = self.repeat_until
+        if not repeat_until:
+            return repeat_until
+        if self.measurement_key_map:
+            repeat_until = protocols.with_measurement_key_mapping(
+                repeat_until, self.measurement_key_map
+            )
+        if self.param_resolver:
+            repeat_until = protocols.resolve_parameters(
+                repeat_until, self.param_resolver, recursive=False
+            )
+        return protocols.with_rescoped_keys(
+            repeat_until,
+            self.parent_path,
+            bindable_keys=self._extern_keys | self._measurement_key_objs,
+        )
+
     def mapped_circuit(self, deep: bool = False) -> 'cirq.Circuit':
         """Applies all maps to the contained circuit and returns the result.
 
@@ -427,12 +446,13 @@ def _decompose_(self) -> Iterator['cirq.Operation']:
         return self.mapped_circuit(deep=False).all_operations()
 
     def _act_on_(self, sim_state: 'cirq.SimulationStateBase') -> bool:
-        if self.repeat_until:
+        mapped_repeat_until = self._mapped_repeat_until
+        if mapped_repeat_until:
             circuit = self._mapped_single_loop()
             while True:
                 for op in circuit.all_operations():
                     protocols.act_on(op, sim_state)
-                if self.repeat_until.resolve(sim_state.classical_data):
+                if mapped_repeat_until.resolve(sim_state.classical_data):
                     break
         else:
             for op in self._decompose_():
@@ -808,7 +828,9 @@ def with_params(
                 by param_values.
         """
         new_params = {}
-        for k in protocols.parameter_symbols(self.circuit):
+        for k in protocols.parameter_symbols(self.circuit) | protocols.parameter_symbols(
+            self.repeat_until
+        ):
             v = self.param_resolver.value_of(k, recursive=False)
             v = protocols.resolve_parameters(v, param_values, recursive=recursive)
             if v != k:

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -1239,4 +1239,57 @@ def test_repeat_until_error():
         )
 
 
+def test_repeat_until_protocols():
+    q = cirq.LineQubit(0)
+    op = cirq.CircuitOperation(
+        cirq.FrozenCircuit(cirq.H(q) ** sympy.Symbol('p'), cirq.measure(q, key='a')),
+        repeat_until=cirq.SympyCondition(sympy.Eq(sympy.Symbol('a'), 0)),
+    )
+    scoped = cirq.with_rescoped_keys(op, ('0',))
+    # Ensure the _repeat_until has been mapped, the measurement has been mapped to the same key,
+    # and the control keys of the subcircuit is empty (because the control key of the condition is
+    # bound to the measurement).
+    assert scoped._mapped_repeat_until.keys == (cirq.MeasurementKey('a', ('0',)),)
+    assert cirq.measurement_key_objs(scoped) == {cirq.MeasurementKey('a', ('0',))}
+    assert not cirq.control_keys(scoped)
+    mapped = cirq.with_measurement_key_mapping(scoped, {'a': 'b'})
+    assert mapped._mapped_repeat_until.keys == (cirq.MeasurementKey('b', ('0',)),)
+    assert cirq.measurement_key_objs(mapped) == {cirq.MeasurementKey('b', ('0',))}
+    assert not cirq.control_keys(mapped)
+    prefixed = cirq.with_key_path_prefix(mapped, ('1',))
+    assert prefixed._mapped_repeat_until.keys == (cirq.MeasurementKey('b', ('1', '0')),)
+    assert cirq.measurement_key_objs(prefixed) == {cirq.MeasurementKey('b', ('1', '0'))}
+    assert not cirq.control_keys(prefixed)
+    setpath = cirq.with_key_path(prefixed, ('2',))
+    assert setpath._mapped_repeat_until.keys == (cirq.MeasurementKey('b', ('2',)),)
+    assert cirq.measurement_key_objs(setpath) == {cirq.MeasurementKey('b', ('2',))}
+    assert not cirq.control_keys(setpath)
+    resolved = cirq.resolve_parameters(setpath, {'p': 1})
+    assert resolved._mapped_repeat_until.keys == (cirq.MeasurementKey('b', ('2',)),)
+    assert cirq.measurement_key_objs(resolved) == {cirq.MeasurementKey('b', ('2',))}
+    assert not cirq.control_keys(resolved)
+
+
+def test_inner_repeat_until_simulate():
+    sim = cirq.Simulator()
+    q = cirq.LineQubit(0)
+    inner_loop = cirq.CircuitOperation(
+        cirq.FrozenCircuit(cirq.H(q), cirq.measure(q, key="inner_loop")),
+        repeat_until=cirq.SympyCondition(sympy.Eq(sympy.Symbol("inner_loop"), 0)),
+    )
+    outer_loop = cirq.Circuit(inner_loop, cirq.X(q), cirq.measure(q, key="outer_loop"))
+    circuit = cirq.Circuit(
+        cirq.CircuitOperation(
+            cirq.FrozenCircuit(outer_loop), repetitions=2, use_repetition_ids=True
+        )
+    )
+    result = sim.run(circuit, repetitions=1)
+    assert all(len(v) == 1 and v[0] == 1 for v in result.records['0:inner_loop'][0][:-1])
+    assert result.records['0:inner_loop'][0][-1] == [0]
+    assert result.records['0:outer_loop'] == [[[1]]]
+    assert all(len(v) == 1 and v[0] == 1 for v in result.records['1:inner_loop'][0][:-1])
+    assert result.records['1:inner_loop'][0][-1] == [0]
+    assert result.records['1:outer_loop'] == [[[1]]]
+
+
 # TODO: Operation has a "gate" property. What is this for a CircuitOperation?