Optimize is_measurement protocol (#4140)

* Optimize is_measurement protocol

* Optimize is_measurement protocol

* No shortcut in circuit

* Change from recursive to iterative

* add test and remove vals_to_decompose arg

* FrozenCircuit uses evaluates fully for memoization

* Coverage for tagged non-GateOperation

Author: smitsanghavi

cirq-core/cirq/circuits/circuit.py

@@ -757,7 +757,7 @@ def findall_operations_with_gate_type(
             yield index, gate_op, cast(T_DESIRED_GATE_TYPE, gate_op.gate)
 
     def has_measurements(self):
-        return any(self.findall_operations(protocols.is_measurement))
+        return protocols.is_measurement(self)
 
     def are_all_measurements_terminal(self) -> bool:
         """Whether all measurement gates are at the end of the circuit."""

cirq-core/cirq/circuits/circuit_operation.py

@@ -160,6 +160,9 @@ def _default_repetition_ids(self) -> Optional[List[str]]:
     def _qid_shape_(self) -> Tuple[int, ...]:
         return tuple(q.dimension for q in self.qubits)
 
+    def _is_measurement_(self) -> bool:
+        return self.circuit._is_measurement_()
+
     def _measurement_keys_(self) -> AbstractSet[str]:
         circuit_keys = [
             value.MeasurementKey.parse_serialized(key_str)

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -81,6 +81,16 @@ def test_repetitions_and_ids_length_mismatch():
         _ = cirq.CircuitOperation(circuit, repetitions=2, repetition_ids=['a', 'b', 'c'])
 
 
+def test_is_measurement_memoization():
+    a = cirq.LineQubit(0)
+    circuit = cirq.FrozenCircuit(cirq.measure(a, key='m'))
+    c_op = cirq.CircuitOperation(circuit)
+    assert circuit._has_measurements is None
+    # Memoize `_has_measurements` in the circuit.
+    assert cirq.is_measurement(c_op)
+    assert circuit._has_measurements is True
+
+
 def test_invalid_measurement_keys():
     a = cirq.LineQubit(0)
     circuit = cirq.FrozenCircuit(cirq.measure(a, key='m'))

cirq-core/cirq/circuits/frozen_circuit.py

@@ -109,6 +109,11 @@ def _unitary_(self) -> Union[np.ndarray, NotImplementedType]:
             self._unitary = super()._unitary_()
         return self._unitary
 
+    def _is_measurement_(self) -> bool:
+        if self._has_measurements is None:
+            self._has_measurements = protocols.is_measurement(self.unfreeze())
+        return self._has_measurements
+
     def all_qubits(self) -> FrozenSet['cirq.Qid']:
         if self._all_qubits is None:
             self._all_qubits = super().all_qubits()

cirq-core/cirq/ops/gate_operation.py

@@ -209,6 +209,13 @@ def _kraus_(self) -> Union[Tuple[np.ndarray], NotImplementedType]:
             return getter()
         return NotImplemented
 
+    def _is_measurement_(self) -> Optional[bool]:
+        getter = getattr(self.gate, '_is_measurement_', None)
+        if getter is not None:
+            return getter()
+        # Let the protocol handle the fallback.
+        return NotImplemented
+
     def _measurement_key_(self) -> Optional[str]:
         getter = getattr(self.gate, '_measurement_key_', None)
         if getter is not None:

cirq-core/cirq/ops/measurement_gate.py

@@ -122,6 +122,9 @@ def full_invert_mask(self):
         mask += (False,) * deficit
         return mask
 
+    def _is_measurement_(self) -> bool:
+        return True
+
     def _measurement_key_(self):
         return self.key
 

cirq-core/cirq/ops/raw_types.py

@@ -664,6 +664,12 @@ def _kraus_(self) -> Union[Tuple[np.ndarray], NotImplementedType]:
     def _measurement_key_(self) -> str:
         return protocols.measurement_key(self.sub_operation, NotImplemented)
 
+    def _is_measurement_(self) -> bool:
+        sub = getattr(self.sub_operation, "_is_measurement_", None)
+        if sub is not None:
+            return sub()
+        return NotImplemented
+
     def _is_parameterized_(self) -> bool:
         return protocols.is_parameterized(self.sub_operation) or any(
             protocols.is_parameterized(tag) for tag in self.tags

cirq-core/cirq/ops/raw_types_test.py

@@ -431,14 +431,19 @@ def test_tagged_operation():
     assert op.gate == cirq.X
     assert op.with_qubits(q2) == cirq.X(q2).with_tags('tag1')
     assert op.with_qubits(q2).qubits == (q2,)
+    assert not cirq.is_measurement(op)
 
 
 def test_tagged_measurement():
+    assert not cirq.is_measurement(cirq.GlobalPhaseOperation(coefficient=-1.0).with_tags('tag0'))
+
     a = cirq.LineQubit(0)
     op = cirq.measure(a, key='m').with_tags('tag')
+    assert cirq.is_measurement(op)
 
     remap_op = cirq.with_measurement_key_mapping(op, {'m': 'k'})
     assert remap_op.tags == ('tag',)
+    assert cirq.is_measurement(remap_op)
     assert cirq.measurement_keys(remap_op) == {'k'}
     assert cirq.with_measurement_key_mapping(op, {'x': 'k'}) == op
 

cirq-core/cirq/protocols/measurement_key_protocol.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Protocol for object that have measurement keys."""
 
-from typing import AbstractSet, Any, Dict, Iterable, Tuple
+from typing import AbstractSet, Any, Dict, List, Iterable, Optional, Tuple
 
 from typing_extensions import Protocol
 
@@ -46,6 +46,10 @@ class SupportsMeasurementKey(Protocol):
     conditional on the measurement outcome being $k$.
     """
 
+    @doc_private
+    def _is_measurement_(self) -> str:
+        """Return if this object is (or contains) a measurement."""
+
     @doc_private
     def _measurement_key_(self) -> str:
         """Return the key that will be used to identify this measurement.
@@ -106,6 +110,22 @@ def measurement_key(val: Any, default: Any = RaiseTypeErrorIfNotProvided):
     raise TypeError(f"Object of type '{type(val)}' had no measurement keys.")
 
 
+def _measurement_keys_from_magic_methods(val: Any) -> Optional[AbstractSet[str]]:
+    """Uses the measurement key related magic methods to get the keys for this object."""
+
+    getter = getattr(val, '_measurement_keys_', None)
+    result = NotImplemented if getter is None else getter()
+    if result is not NotImplemented and result is not None:
+        return set(result)
+
+    getter = getattr(val, '_measurement_key_', None)
+    result = NotImplemented if getter is None else getter()
+    if result is not NotImplemented and result is not None:
+        return {result}
+
+    return result
+
+
 def measurement_keys(val: Any, *, allow_decompose: bool = True) -> AbstractSet[str]:
     """Gets the measurement keys of measurements within the given value.
 
@@ -122,15 +142,9 @@ def measurement_keys(val: Any, *, allow_decompose: bool = True) -> AbstractSet[s
         The measurement keys of the value. If the value has no measurement,
         the result is the empty tuple.
     """
-    getter = getattr(val, '_measurement_keys_', None)
-    result = NotImplemented if getter is None else getter()
-    if result is not NotImplemented and result is not None:
-        return set(result)
-
-    getter = getattr(val, '_measurement_key_', None)
-    result = NotImplemented if getter is None else getter()
+    result = _measurement_keys_from_magic_methods(val)
     if result is not NotImplemented and result is not None:
-        return {result}
+        return result
 
     if allow_decompose:
         operations, _, _ = _try_decompose_into_operations_and_qubits(val)
@@ -140,13 +154,66 @@ def measurement_keys(val: Any, *, allow_decompose: bool = True) -> AbstractSet[s
     return set()
 
 
-def is_measurement(val: Any) -> bool:
-    """Determines whether or not the given value is a measurement.
+def _is_measurement_from_magic_method(val: Any) -> Optional[bool]:
+    """Uses `is_measurement` magic method to determine if this object is a measurement."""
+    getter = getattr(val, '_is_measurement_', None)
+    return NotImplemented if getter is None else getter()
+
 
-    Measurements are identified by the fact that `cirq.measurement_keys` returns
-    a non-empty result for them.
+def _is_any_measurement(vals: List[Any], allow_decompose: bool) -> bool:
+    """Given a list of objects, returns True if any of them is a measurement.
+
+    If `allow_decompose` is True, decomposes the objects and runs the measurement checks on the
+    constituent decomposed operations. But a decompose operation is only called if all cheaper
+    checks are done. A BFS for searching measurements, where "depth" is each level of decompose.
+    """
+    vals_to_decompose = []  # type: List[Any]
+    while vals:
+        val = vals.pop(0)
+        result = _is_measurement_from_magic_method(val)
+        if result is not NotImplemented:
+            if result is True:
+                return True
+            if result is False:
+                # Do not try any other strategies if `val` was explicitly marked as
+                # "not measurement".
+                continue
+
+        keys = _measurement_keys_from_magic_methods(val)
+        if keys is not NotImplemented and bool(keys) is True:
+            return True
+
+        if allow_decompose:
+            vals_to_decompose.append(val)
+
+        # If vals has finished iterating over, keep decomposing from vals_to_decompose until vals
+        # is populated with something.
+        while not vals:
+            if not vals_to_decompose:
+                # Nothing left to process, this is not a measurement.
+                return False
+            operations, _, _ = _try_decompose_into_operations_and_qubits(vals_to_decompose.pop(0))
+            if operations:
+                # Reverse the decomposed operations because measurements are typically at later
+                # moments.
+                vals = operations[::-1]
+
+    return False
+
+
+def is_measurement(val: Any, allow_decompose: bool = True) -> bool:
+    """Determines whether or not the given value is a measurement (or contains one).
+
+    Measurements are identified by the fact that any of them may have an `_is_measurement_` method
+    or `cirq.measurement_keys` returns a non-empty result for them.
+
+    Args:
+        val: The value which to evaluate.
+        allow_decompose: Defaults to True. When true, composite operations that
+            don't directly specify their `_is_measurement_` property will be decomposed in
+            order to find any measurements keys within the decomposed operations.
     """
-    return bool(measurement_keys(val))
+    return _is_any_measurement([val], allow_decompose)
 
 
 def with_measurement_key_mapping(val: Any, key_map: Dict[str, str]):

cirq-core/cirq/protocols/measurement_key_protocol_test.py

@@ -22,6 +22,7 @@ class ReturnsStr:
         def _measurement_key_(self):
             return 'door locker'
 
+    assert cirq.is_measurement(ReturnsStr())
     assert cirq.measurement_key(ReturnsStr()) == 'door locker'
 
     assert cirq.measurement_key(ReturnsStr(), None) == 'door locker'
@@ -84,6 +85,36 @@ def qubits(self):
     assert not cirq.is_measurement(NotImplementedOperation())
 
 
+def test_measurement_without_key():
+    class MeasurementWithoutKey:
+        def _is_measurement_(self):
+            return True
+
+    with pytest.raises(TypeError, match='no measurement keys'):
+        _ = cirq.measurement_key(MeasurementWithoutKey())
+
+    assert cirq.is_measurement(MeasurementWithoutKey())
+
+
+def test_non_measurement_with_key():
+    class NonMeasurementGate(cirq.Gate):
+        def _is_measurement_(self):
+            return False
+
+        def _decompose_(self, qubits):
+            # Decompose should not be called by `is_measurement`
+            assert False
+
+        def _measurement_key_(self):
+            # `measurement_key`` should not be called by `is_measurement`
+            assert False
+
+        def num_qubits(self) -> int:
+            return 2  # coverage: ignore
+
+    assert not cirq.is_measurement(NonMeasurementGate())
+
+
 def test_measurement_keys():
     class Composite(cirq.Gate):
         def _decompose_(self, qubits):
@@ -102,8 +133,10 @@ def num_qubits(self) -> int:
             return 1
 
     a, b = cirq.LineQubit.range(2)
+    assert cirq.is_measurement(Composite())
     assert cirq.measurement_keys(Composite()) == {'inner1', 'inner2'}
     assert cirq.measurement_keys(Composite().on(a, b)) == {'inner1', 'inner2'}
+    assert not cirq.is_measurement(Composite(), allow_decompose=False)
     assert cirq.measurement_keys(Composite(), allow_decompose=False) == set()
     assert cirq.measurement_keys(Composite().on(a, b), allow_decompose=False) == set()