Add Moment.from_ops to more efficiently construct moments (quantumlib#6078)

Review: @tanujkhattar

Author: maffoo

cirq/circuits/moment.py

@@ -19,6 +19,7 @@
     AbstractSet,
     Any,
     Callable,
+    cast,
     Dict,
     FrozenSet,
     Iterable,
@@ -80,17 +81,26 @@ class Moment:
             are no such operations, returns an empty Moment.
     """
 
-    def __init__(self, *contents: 'cirq.OP_TREE') -> None:
+    def __init__(self, *contents: 'cirq.OP_TREE', _flatten_contents: bool = True) -> None:
         """Constructs a moment with the given operations.
 
         Args:
             contents: The operations applied within the moment.
                 Will be flattened and frozen into a tuple before storing.
+            _flatten_contents: If True, use flatten_to_ops to convert
+                the OP_TREE of contents into a tuple of Operation. If False,
+                we skip flattening and assume that contents already consists
+                of individual operations. This is used internally by helper
+                methods to avoid unnecessary validation.
 
         Raises:
             ValueError: A qubit appears more than once.
         """
-        self._operations = tuple(op_tree.flatten_to_ops(contents))
+        self._operations = (
+            tuple(op_tree.flatten_to_ops(contents))
+            if _flatten_contents
+            else cast(Tuple['cirq.Operation'], contents)
+        )
         self._sorted_operations: Optional[Tuple['cirq.Operation', ...]] = None
 
         # An internal dictionary to support efficient operation access by qubit.
@@ -106,6 +116,20 @@ def __init__(self, *contents: 'cirq.OP_TREE') -> None:
         self._measurement_key_objs: Optional[FrozenSet['cirq.MeasurementKey']] = None
         self._control_keys: Optional[FrozenSet['cirq.MeasurementKey']] = None
 
+    @classmethod
+    def from_ops(cls, *ops: 'cirq.Operation') -> 'cirq.Moment':
+        """Construct a Moment from the given operations.
+
+        This avoids calling `flatten_to_ops` in the moment constructor, which
+        results in better performance in cases where the contents of the moment
+        are already in the form of a sequence of operations rather than an
+        arbitrary OP_TREE.
+
+        Args:
+            *ops: Operations to include in the Moment.
+        """
+        return cls(*ops, _flatten_contents=False)
+
     @property
     def operations(self) -> Tuple['cirq.Operation', ...]:
         return self._operations
@@ -164,7 +188,7 @@ def with_operation(self, operation: 'cirq.Operation') -> 'cirq.Moment':
             raise ValueError(f'Overlapping operations: {operation}')
 
         # Use private variables to facilitate a quick copy.
-        m = Moment()
+        m = Moment(_flatten_contents=False)
         m._operations = self._operations + (operation,)
         m._sorted_operations = None
         m._qubits = self._qubits.union(operation.qubits)
@@ -194,7 +218,7 @@ def with_operations(self, *contents: 'cirq.OP_TREE') -> 'cirq.Moment':
         if not flattened_contents:
             return self
 
-        m = Moment()
+        m = Moment(_flatten_contents=False)
         # Use private variables to facilitate a quick copy.
         m._qubit_to_op = self._qubit_to_op.copy()
         qubits = set(self._qubits)
@@ -483,7 +507,7 @@ def _json_dict_(self) -> Dict[str, Any]:
 
     @classmethod
     def _from_json_dict_(cls, operations, **kwargs):
-        return Moment(operations)
+        return cls.from_ops(*operations)
 
     def __add__(self, other: 'cirq.OP_TREE') -> 'cirq.Moment':
 

cirq/circuits/moment_test.py

@@ -170,6 +170,13 @@ def test_operation_at():
     assert cirq.Moment([cirq.CZ(a, b), cirq.X(c)]).operation_at(a) == cirq.CZ(a, b)
 
 
+def test_from_ops():
+    a = cirq.NamedQubit('a')
+    b = cirq.NamedQubit('b')
+
+    assert cirq.Moment.from_ops(cirq.X(a), cirq.Y(b)) == cirq.Moment(cirq.X(a), cirq.Y(b))
+
+
 def test_with_operation():
     a = cirq.NamedQubit('a')
     b = cirq.NamedQubit('b')