Allow repeated measurement keys

cirq-core/cirq/circuits/circuit.py

@@ -1980,11 +1980,11 @@ def earliest_available_moment(
         while k > 0:
             k -= 1
             moment = self._moments[k]
-            # This should also validate that measurement keys are disjoint once we allow repeated
-            # measurements. Search for same message in raw_types.py.
+            moment_measurement_keys = protocols.measurement_key_objs(moment)
             if (
                 moment.operates_on(op_qubits)
-                or not op_control_keys.isdisjoint(protocols.measurement_key_objs(moment))
+                or not op_measurement_keys.isdisjoint(moment_measurement_keys)
+                or not op_control_keys.isdisjoint(moment_measurement_keys)
                 or not protocols.control_keys(moment).isdisjoint(op_measurement_keys)
             ):
                 return last_available

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -834,24 +834,28 @@ def test_mapped_circuit_keeps_keys_under_parent_path():
     assert cirq.measurement_key_names(op2.mapped_circuit()) == {'X:A', 'X:B', 'X:C', 'X:D'}
 
 
-def test_keys_conflict_no_repetitions():
+def test_mapped_circuit_allows_repeated_keys():
     q = cirq.LineQubit(0)
     op1 = cirq.CircuitOperation(
         cirq.FrozenCircuit(
             cirq.measure(q, key='A'),
         )
     )
     op2 = cirq.CircuitOperation(cirq.FrozenCircuit(op1, op1))
-    with pytest.raises(ValueError, match='Conflicting measurement keys found: A'):
-        _ = op2.mapped_circuit(deep=True)
-
-
-def test_keys_conflict_locally():
-    q = cirq.LineQubit(0)
+    circuit = op2.mapped_circuit(deep=True)
+    cirq.testing.assert_has_diagram(
+        circuit,
+        "0: ───M('A')───M('A')───",
+        use_unicode_characters=True,
+    )
     op1 = cirq.measure(q, key='A')
     op2 = cirq.CircuitOperation(cirq.FrozenCircuit(op1, op1))
-    with pytest.raises(ValueError, match='Conflicting measurement keys found: A'):
-        _ = op2.mapped_circuit()
+    circuit = op2.mapped_circuit()
+    cirq.testing.assert_has_diagram(
+        circuit,
+        "0: ───M('A')───M('A')───",
+        use_unicode_characters=True,
+    )
 
 
 # TODO: Operation has a "gate" property. What is this for a CircuitOperation?

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -259,6 +259,45 @@ def test_key_set(sim):
     assert result.measurements['b'] == 1
 
 
+@pytest.mark.parametrize('sim', ALL_SIMULATORS)
+def test_repeated_measurement_unset(sim):
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.X(q0),
+        cirq.measure(q0, key='a'),
+        cirq.X(q1).with_classical_controls(cirq.KeyCondition(cirq.MeasurementKey('a'), index=-2)),
+        cirq.measure(q1, key='b'),
+        cirq.X(q1).with_classical_controls(cirq.KeyCondition(cirq.MeasurementKey('a'), index=-1)),
+        cirq.measure(q1, key='c'),
+    )
+    result = sim.run(circuit)
+    assert result.records['a'][0][0][0] == 0
+    assert result.records['a'][0][1][0] == 1
+    assert result.records['b'][0][0][0] == 0
+    assert result.records['c'][0][0][0] == 1
+
+
+@pytest.mark.parametrize('sim', ALL_SIMULATORS)
+def test_repeated_measurement_set(sim):
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.X(q0),
+        cirq.measure(q0, key='a'),
+        cirq.X(q0),
+        cirq.measure(q0, key='a'),
+        cirq.X(q1).with_classical_controls(cirq.KeyCondition(cirq.MeasurementKey('a'), index=-2)),
+        cirq.measure(q1, key='b'),
+        cirq.X(q1).with_classical_controls(cirq.KeyCondition(cirq.MeasurementKey('a'), index=-1)),
+        cirq.measure(q1, key='c'),
+    )
+    result = sim.run(circuit)
+    assert result.records['a'][0][0][0] == 1
+    assert result.records['a'][0][1][0] == 0
+    assert result.records['b'][0][0][0] == 1
+    assert result.records['c'][0][0][0] == 1
+
+
 @pytest.mark.parametrize('sim', ALL_SIMULATORS)
 def test_subcircuit_key_unset(sim):
     q0, q1 = cirq.LineQubit.range(2)

cirq-core/cirq/ops/measurement_gate_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import cast
 import numpy as np
 import pytest
 
@@ -323,10 +324,13 @@ def test_act_on_state_vector():
         dtype=np.complex64,
     )
     cirq.act_on(m, args)
+    datastore = cast(cirq.ClassicalDataDictionaryStore, args.classical_data)
+    out = cirq.MeasurementKey('out')
     assert args.log_of_measurement_results == {'out': [0, 1]}
-
-    with pytest.raises(ValueError, match="already logged to key"):
-        cirq.act_on(m, args)
+    assert datastore.records[out] == [(0, 1)]
+    cirq.act_on(m, args)
+    assert args.log_of_measurement_results == {'out': [0, 1]}
+    assert datastore.records[out] == [(0, 1), (0, 1)]
 
 
 def test_act_on_clifford_tableau():
@@ -361,10 +365,13 @@ def test_act_on_clifford_tableau():
         log_of_measurement_results={},
     )
     cirq.act_on(m, args)
+    datastore = cast(cirq.ClassicalDataDictionaryStore, args.classical_data)
+    out = cirq.MeasurementKey('out')
     assert args.log_of_measurement_results == {'out': [0, 1]}
-
-    with pytest.raises(ValueError, match="already logged to key"):
-        cirq.act_on(m, args)
+    assert datastore.records[out] == [(0, 1)]
+    cirq.act_on(m, args)
+    assert args.log_of_measurement_results == {'out': [0, 1]}
+    assert datastore.records[out] == [(0, 1), (0, 1)]
 
 
 def test_act_on_stabilizer_ch_form():
@@ -399,10 +406,13 @@ def test_act_on_stabilizer_ch_form():
         initial_state=10,
     )
     cirq.act_on(m, args)
+    datastore = cast(cirq.ClassicalDataDictionaryStore, args.classical_data)
+    out = cirq.MeasurementKey('out')
     assert args.log_of_measurement_results == {'out': [0, 1]}
-
-    with pytest.raises(ValueError, match="already logged to key"):
-        cirq.act_on(m, args)
+    assert datastore.records[out] == [(0, 1)]
+    cirq.act_on(m, args)
+    assert args.log_of_measurement_results == {'out': [0, 1]}
+    assert datastore.records[out] == [(0, 1), (0, 1)]
 
 
 def test_act_on_qutrit():

cirq-core/cirq/ops/raw_types.py

@@ -572,11 +572,13 @@ def _commutes_(
         if not isinstance(other, Operation):
             return NotImplemented
 
-        # This should also validate that measurement keys are disjoint once we allow repeated
-        # measurements. Search for same message in circuit.py.
-        if not protocols.control_keys(self).isdisjoint(
-            protocols.measurement_key_objs(other)
-        ) or not protocols.control_keys(other).isdisjoint(protocols.measurement_key_objs(self)):
+        self_keys = protocols.measurement_key_objs(self)
+        other_keys = protocols.measurement_key_objs(other)
+        if (
+            not self_keys.isdisjoint(other_keys)
+            or not protocols.control_keys(self).isdisjoint(other_keys)
+            or not protocols.control_keys(other).isdisjoint(self_keys)
+        ):
             return False
 
         if hasattr(other, 'qubits') and set(self.qubits).isdisjoint(other.qubits):

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

@@ -1,13 +1,13 @@
 {
   "cirq_type": "ClassicalDataDictionaryStore",
-  "measurements": [
+  "records": [
     [
       {
         "cirq_type": "MeasurementKey",
         "name": "m",
         "path": []
       },
-      [0, 1]
+      [[0, 1]]
     ]
   ],
   "measured_qubits": [
@@ -17,7 +17,7 @@
         "name": "m",
         "path": []
       },
-      [
+      [[
         {
           "cirq_type": "LineQubit",
           "x": 0
@@ -26,17 +26,17 @@
           "cirq_type": "LineQubit",
           "x": 1
         }
-      ]
+      ]]
     ]
   ],
-  "channel_measurements": [
+  "channel_records": [
     [
       {
         "cirq_type": "MeasurementKey",
         "name": "c",
         "path": []
       },
-      3
+      [3]
     ]
   ],
   "measurement_types": [

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

@@ -1 +1 @@
-cirq.ClassicalDataDictionaryStore(_measurements={cirq.MeasurementKey('m'): [0, 1]}, _measured_qubits={cirq.MeasurementKey('m'): [cirq.LineQubit(0), cirq.LineQubit(1)]}, _channel_measurements={cirq.MeasurementKey('c'): 3}, _measurement_types={cirq.MeasurementKey('m'): cirq.MeasurementType.MEASUREMENT, cirq.MeasurementKey('c'): cirq.MeasurementType.CHANNEL})
+cirq.ClassicalDataDictionaryStore(_records={cirq.MeasurementKey('m'): [[0, 1]]}, _measured_qubits={cirq.MeasurementKey('m'): [[cirq.LineQubit(0), cirq.LineQubit(1)]]}, _channel_records={cirq.MeasurementKey('c'): [3]}, _measurement_types={cirq.MeasurementKey('m'): cirq.MeasurementType.MEASUREMENT, cirq.MeasurementKey('c'): cirq.MeasurementType.CHANNEL})

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

@@ -7,15 +7,17 @@
         "cirq_type": "MeasurementKey",
         "name": "a",
         "path": []
-      }
+      },
+      "index": -1
     },
     {
       "cirq_type": "KeyCondition",
       "key": {
         "cirq_type": "MeasurementKey",
         "name": "b",
         "path": []
-      }
+      },
+      "index": -1
     }
   ],
   "sub_operation": {

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

@@ -4,5 +4,6 @@
     "cirq_type": "MeasurementKey",
     "name": "a",
     "path": []
-  }
+  },
+  "index": -1
 }

cirq-core/cirq/sim/act_on_args.py

@@ -76,8 +76,8 @@ def __init__(
         self._set_qubits(qubits)
         self.prng = prng
         self._classical_data = classical_data or value.ClassicalDataDictionaryStore(
-            _measurements={
-                value.MeasurementKey.parse_serialized(k): tuple(v)
+            _records={
+                value.MeasurementKey.parse_serialized(k): [tuple(v)]
                 for k, v in (log_of_measurement_results or {}).items()
             }
         )

cirq-core/cirq/sim/act_on_args_container.py

@@ -72,8 +72,8 @@ def __init__(
         self._qubits = tuple(qubits)
         self.split_untangled_states = split_untangled_states
         self._classical_data = classical_data or value.ClassicalDataDictionaryStore(
-            _measurements={
-                value.MeasurementKey.parse_serialized(k): tuple(v)
+            _records={
+                value.MeasurementKey.parse_serialized(k): [tuple(v)]
                 for k, v in (log_of_measurement_results or {}).items()
             }
         )

cirq-core/cirq/sim/clifford/act_on_stabilizer_args.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import abc
-from typing import Any, Dict, Generic, Optional, Sequence, TYPE_CHECKING, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, Sequence, TYPE_CHECKING, TypeVar, Union
 
 import numpy as np
 
@@ -38,7 +38,7 @@ def __init__(
         self,
         state: TStabilizerState,
         prng: Optional[np.random.RandomState] = None,
-        log_of_measurement_results: Optional[Dict[str, Any]] = None,
+        log_of_measurement_results: Optional[Dict[str, List[int]]] = None,
         qubits: Optional[Sequence['cirq.Qid']] = None,
         classical_data: Optional['cirq.ClassicalDataStore'] = None,
     ):

cirq-core/cirq/sim/clifford/act_on_stabilizer_ch_form_args.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional, Sequence, TYPE_CHECKING, Union
+from typing import Dict, List, Optional, Sequence, TYPE_CHECKING, Union
 
 import numpy as np
 
-from cirq import _compat, value, ops, protocols
+from cirq import _compat, value
 from cirq.sim.clifford import stabilizer_state_ch_form
 from cirq.sim.clifford.act_on_stabilizer_args import ActOnStabilizerArgs
 
@@ -39,7 +39,7 @@ def __init__(
         self,
         state: Optional['cirq.StabilizerStateChForm'] = None,
         prng: Optional[np.random.RandomState] = None,
-        log_of_measurement_results: Optional[Dict[str, Any]] = None,
+        log_of_measurement_results: Optional[Dict[str, List[int]]] = None,
         qubits: Optional[Sequence['cirq.Qid']] = None,
         initial_state: Union[int, 'cirq.StabilizerStateChForm'] = 0,
         classical_data: Optional['cirq.ClassicalDataStore'] = None,
@@ -107,16 +107,10 @@ def sample(
         repetitions: int = 1,
         seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
     ) -> np.ndarray:
-        measurements = value.ClassicalDataDictionaryStore()
         prng = value.parse_random_state(seed)
-        for i in range(repetitions):
-            op = ops.measure(*qubits, key=str(i))
+        axes = self.get_axes(qubits)
+        measurements = []
+        for _ in range(repetitions):
             state = self.state.copy()
-            ch_form_args = ActOnStabilizerCHFormArgs(
-                classical_data=measurements,
-                prng=prng,
-                qubits=self.qubits,
-                initial_state=state,
-            )
-            protocols.act_on(op, ch_form_args)
-        return np.array(list(measurements.measurements.values()), dtype=bool)
+            measurements.append([state._measure(i, prng) for i in axes])
+        return np.array(measurements, dtype=bool)

cirq-core/cirq/sim/clifford/clifford_simulator.py

@@ -293,4 +293,4 @@ def apply_measurement(
             initial_state=state.ch_form,
         )
         act_on(op, ch_form_args)
-        measurements.update({str(k): list(v) for k, v in classical_data.measurements.items()})
+        measurements.update({str(k): list(v[-1]) for k, v in classical_data.records.items()})

cirq-core/cirq/sim/clifford/stabilizer_state_ch_form_test.py

@@ -73,6 +73,6 @@ def test_run():
                 initial_state=state,
             )
             cirq.act_on(op, args)
-        measurements = {str(k): list(v) for k, v in classical_data.measurements.items()}
+        measurements = {str(k): list(v[-1]) for k, v in classical_data.records.items()}
         assert measurements['1'] == [1]
         assert measurements['0'] != measurements['2']

cirq-core/cirq/sim/operation_target.py

@@ -14,7 +14,6 @@
 """An interface for quantum states as targets for operations."""
 import abc
 from typing import (
-    Any,
     Dict,
     Generic,
     Iterator,
@@ -86,7 +85,7 @@ def qubits(self) -> Tuple['cirq.Qid', ...]:
         """Gets the qubit order maintained by this target."""
 
     @property
-    def log_of_measurement_results(self) -> Dict[str, Any]:
+    def log_of_measurement_results(self) -> Dict[str, List[int]]:
         """Gets the log of measurement results."""
         return {str(k): list(self.classical_data.get_digits(k)) for k in self.classical_data.keys()}