Created ControlValues for controlled gates/operations, fix for quantumlib#4512

created control_values.py which contains the ControlValues class.
FreeVars and ConstrainedVars classes are provided for ease of use.
while the basic idea of ControlValues integrating it inside the code base was challening
the old way of using control_values assumed it's a tuple of tuples of ints and was used as thus (comparasion, hashing, slicing, fomatting, conditioning, and loops), the ControlValues class had to provide these functionalities
the trickiest part to get right was the support for formatting!

Author: NoureldinYosri

cirq-core/cirq/ops/control_values.py

@@ -48,77 +48,79 @@ def __init__(
         self, control_values: Sequence[Union[int, Collection[int], Type['ControlValues']]]
     ):
         if len(control_values) == 0:
-            self.vals = cast(Tuple[Tuple[int, ...], ...], (()))
-            self.num_variables = 0
-            self.nxt = None
-            self.itr = None
+            self._vals = cast(Tuple[Tuple[int, ...], ...], (()))
+            self._num_variables = 0
+            self._nxt = None
             return
-        self.itr = None
-        self.nxt = None
+        self._nxt = None
 
         if len(control_values) > 1:
-            self.nxt = ControlValues(control_values[1:])
+            self._nxt = ControlValues(control_values[1:])
 
         if isinstance(control_values[0], ControlValues):
             aux = control_values[0].copy()
-            aux.product(self.nxt)
-            self.vals, self.num_variables, self.nxt = aux.vals, aux.num_variables, aux.nxt
-            self.vals = cast(Tuple[Tuple[int, ...], ...], self.vals)
+            aux.product(self._nxt)
+            self._vals, self._num_variables, self._nxt = aux._vals, aux._num_variables, aux._nxt
+            self._vals = cast(Tuple[Tuple[int, ...], ...], self._vals)
             return
 
         val = control_values[0]
         if isinstance(val, int):
-            self.vals = _from_int(val)
+            self._vals = _from_int(val)
         elif isinstance(val, (list, tuple)):
             if isinstance(val[0], int):
-                self.vals = _from_sequence_int(val)
+                self._vals = _from_sequence_int(val)
             else:
-                self.vals = _from_sequence_sequence(val)
+                self._vals = _from_sequence_sequence(val)
         else:
             raise TypeError(f'{val} is of Unsupported type {type(val)}')
-        self.num_variables = len(self.vals[0])
+        self._num_variables = len(self._vals[0])
 
     def product(self, other):
-        # Cartesian product of all combinations in self x other
+        """Sets self to be the cartesian product of all combinations in self x other.
+
+        Args:
+            other: A ControlValues object
+        """
         if other is None:
             return
         other = other.copy()
         cur = self
-        while cur.nxt is not None:
-            cur = cur.nxt
-        cur.nxt = other
+        while cur._nxt is not None:
+            cur = cur._nxt
+        cur._nxt = other
 
     def __call__(self):
         return self.__iter__()
 
     def __iter__(self):
-        nxt = self.nxt if self.nxt else lambda: [()]
-        if self.num_variables:
-            self.itr = itertools.product(self.vals, nxt())
+        nxt = self._nxt if self._nxt else lambda: [()]
+        if self._num_variables:
+            return itertools.product(self._vals, nxt())
         else:
-            self.itr = itertools.product(*(), nxt())
-        return self.itr
+            return itertools.product(*(), nxt())
 
     def copy(self):
-        if self.num_variables == 0:
+        """Returns a deep copy of the object."""
+        if self._num_variables == 0:
             new_copy = ControlValues([])
         else:
             new_copy = ControlValues(
                 [
-                    copy.deepcopy(self.vals),
+                    copy.deepcopy(self._vals),
                 ]
             )
-        new_copy.nxt = None
-        if self.nxt:
-            new_copy.nxt = self.nxt.copy()
+        new_copy._nxt = None
+        if self._nxt:
+            new_copy._nxt = self._nxt.copy()
         return new_copy
 
     def __len__(self):
         cur = self
         num_variables = 0
         while cur is not None:
-            num_variables += cur.num_variables
-            cur = cur.nxt
+            num_variables += cur._num_variables
+            cur = cur._nxt
         return num_variables
 
     def __getitem__(self, key):
@@ -131,9 +133,9 @@ def __getitem__(self, key):
         if not 0 <= key < num_variables:
             key = key % num_variables
         cur = self
-        while cur.num_variables <= key:
-            key -= cur.num_variables
-            cur = cur.nxt
+        while cur._num_variables <= key:
+            key -= cur._num_variables
+            cur = cur._nxt
         return cur
 
     def __eq__(self, other):
@@ -147,33 +149,55 @@ def __eq__(self, other):
         return self_values == other_values
 
     def identifier(self, companions: Sequence[Union[int, 'cirq.Qid']]):
+        """Returns an identifier that pairs the control_values with their counterparts
+            in the companion sequence (e.g. sequence of the control qubits ids)
+
+        Args:
+            companions: sequence of the same length as the number of control qubits.
+
+        Returns:
+            Tuple of pairs of the control values paired with its corresponding companion.
+        """
         companions = tuple(companions)
         controls = []
         cur = cast(Optional[ControlValues], self)
         while cur is not None:
-            controls.append((cur.vals, companions[: cur.num_variables]))
-            companions = companions[cur.num_variables :]
-            cur = cur.nxt
+            controls.append((cur._vals, companions[: cur._num_variables]))
+            companions = companions[cur._num_variables :]
+            cur = cur._nxt
         return tuple(controls)
 
-    def check_dimentionality(
+    def check_dimensionality(
         self,
         qid_shape: Optional[Union[Tuple[int, ...], List[int]]] = None,
         controls: Optional[Union[Tuple['cirq.Qid', ...], List['cirq.Qid']]] = None,
         offset=0,
     ):
-        if self.num_variables == 0:
+        """Checks the dimentionality of control values with respect to qid_shape or controls.
+            *At least one of qid_shape and controls must be provided.
+            *if both are provided then controls is ignored*
+
+        Args:
+            qid_shape: Sequence of shapes of the control qubits.
+            controls: Sequence of Qids.
+            offset: starting index.
+        Raises:
+            ValueError:
+                - if none of qid_shape or controls are provided or both are empty.
+                - if one of the control values violates the shape of its qubit.
+        """
+        if self._num_variables == 0:
             return
         if qid_shape is None and controls is None:
             raise ValueError('At least one of qid_shape or controls has to be not given.')
         if controls is not None:
             controls = tuple(controls)
         if (qid_shape is None or len(qid_shape) == 0) and controls is not None:
-            qid_shape = tuple(q.dimension for q in controls[: self.num_variables])
+            qid_shape = tuple(q.dimension for q in controls[: self._num_variables])
         qid_shape = cast(Tuple[int], qid_shape)
-        for product in self.vals:
+        for product in self._vals:
             product = flatten(product)
-            for i in range(self.num_variables):
+            for i in range(self._num_variables):
                 if not 0 <= product[i] < qid_shape[i]:
                     message = (
                         'Control values <{!r}> outside of range ' 'for control qubit number <{!r}>.'
@@ -185,40 +209,59 @@ def check_dimentionality(
                         )
                     raise ValueError(message)
 
-        if self.nxt is not None:
-            self.nxt.check_dimentionality(
-                qid_shape=qid_shape[self.num_variables :],
-                controls=controls[self.num_variables :] if controls else None,
-                offset=offset + self.num_variables,
+        if self._nxt is not None:
+            self._nxt.check_dimensionality(
+                qid_shape=qid_shape[self._num_variables :],
+                controls=controls[self._num_variables :] if controls else None,
+                offset=offset + self._num_variables,
             )
 
     def are_same_value(self, value: int = 1):
-        for product in self.vals:
+        for product in self._vals:
             product = flatten(product)
             if not all(v == value for v in product):
                 return False
-        if self.nxt is not None:
-            return self.nxt.are_same_value(value)
+        if self._nxt is not None:
+            return self._nxt.are_same_value(value)
         return True
 
     def arrangements(self):
+        """Returns a list of the control values.
+
+        Returns:
+            lists containing the control values whose product is the list of all
+            possible combinations of the control values.
+        """
         _arrangements = []
         cur = self
         while cur is not None:
-            if cur.num_variables == 1:
-                _arrangements.append(flatten(cur.vals))
+            if cur._num_variables == 1:
+                _arrangements.append(flatten(cur._vals))
             else:
-                _arrangements.append(tuple(flatten(product) for product in cur.vals))
-            cur = cur.nxt
+                _arrangements.append(tuple(flatten(product) for product in cur._vals))
+            cur = cur._nxt
         return _arrangements
 
     def pop(self):
-        if self.nxt is None:
+        """Removes the last control values combination."""
+        if self._nxt is None:
             return None
-        self.nxt = self.nxt.pop()
+        self._nxt = self._nxt.pop()
         return self
 
 
+def to_control_values(
+    values: Union[ControlValues, Sequence[Union[int, Collection[int]]]]
+) -> ControlValues:
+    if not isinstance(values, ControlValues):
+        # Convert to sorted tuples
+        return ControlValues(
+            tuple((val,) if isinstance(val, int) else tuple(sorted(val)) for val in values)
+        )
+    else:
+        return values
+
+
 class FreeVars(ControlValues):
     pass
 

cirq-core/cirq/ops/control_values_test.py

@@ -78,13 +78,13 @@ def test_slicing_not_supported():
         _ = control_vals[0:1]
 
 
-def test_check_dimentionality():
+def test_check_dimensionality():
     empty_control_vals = cv.ControlValues([])
-    empty_control_vals.check_dimentionality()
+    empty_control_vals.check_dimensionality()
 
     control_values = cv.ControlValues([[0, 1], 1])
     with pytest.raises(ValueError):
-        control_values.check_dimentionality()
+        control_values.check_dimensionality()
 
 
 def test_pop():

cirq-core/cirq/ops/controlled_gate.py

@@ -79,18 +79,10 @@ def __init__(
             raise ValueError('len(control_qid_shape) != num_controls')
         self.control_qid_shape = tuple(control_qid_shape)
 
-        # Convert to sorted tuples
-        if not isinstance(control_values, cv.ControlValues):
-            self.control_values = cv.ControlValues(
-                tuple(
-                    (val,) if isinstance(val, int) else tuple(sorted(val)) for val in control_values
-                )
-            )
-        else:
-            self.control_values = control_values
+        self.control_values = cv.to_control_values(control_values)
 
         # Verify control values not out of bounds
-        self.control_values.check_dimentionality(self.control_qid_shape)
+        self.control_values.check_dimensionality(self.control_qid_shape)
 
         # Flatten nested ControlledGates.
         if isinstance(sub_gate, ControlledGate):
@@ -100,6 +92,16 @@ def __init__(
         else:
             self.sub_gate = sub_gate
 
+    @property
+    def control_values(self) -> cv.ControlValues:
+        return self._control_values
+
+    @control_values.setter
+    def control_values(
+        self, values: Union[cv.ControlValues, Sequence[Union[int, Collection[int]]]]
+    ) -> None:
+        self._control_values = cv.to_control_values(values)
+
     def num_controls(self) -> int:
         return len(self.control_qid_shape)
 

cirq-core/cirq/ops/controlled_gate_test.py

@@ -20,6 +20,7 @@
 
 import cirq
 from cirq.type_workarounds import NotImplementedType
+from cirq.ops import control_values as cv
 
 
 class GateUsingWorkspaceForApplyUnitary(cirq.SingleQubitGate):
@@ -158,6 +159,22 @@ def test_init2():
     assert gate.num_qubits() == 3
     assert cirq.qid_shape(gate) == (3, 3, 2)
 
+    gate = cirq.ControlledGate(cirq.Z, control_values=cv.FreeVars([0, (0, 1)]))
+    assert gate.sub_gate is cirq.Z
+    assert gate.num_controls() == 2
+    assert gate.control_values == ((0,), (0, 1))
+    assert gate.control_qid_shape == (2, 2)
+    assert gate.num_qubits() == 3
+    assert cirq.qid_shape(gate) == (2, 2, 2)
+
+    gate = cirq.ControlledGate(cirq.Z, control_values=cv.ConstrainedVars([(1, 0), (0, 1)]))
+    assert gate.sub_gate is cirq.Z
+    assert gate.num_controls() == 2
+    assert gate.control_values == [((1, 0), (0, 1))]
+    assert gate.control_qid_shape == (2, 2)
+    assert gate.num_qubits() == 3
+    assert cirq.qid_shape(gate) == (2, 2, 2)
+
 
 def test_validate_args():
     a = cirq.NamedQubit('a')

cirq-core/cirq/ops/controlled_operation.py

@@ -53,17 +53,11 @@ def __init__(
             control_values = ((1,),) * len(controls)
         if len(control_values) != len(controls):
             raise ValueError('len(control_values) != len(controls)')
-        if not isinstance(control_values, cv.ControlValues):
-            self.control_values = cv.ControlValues(
-                # Convert to sorted tuples
-                tuple(
-                    (val,) if isinstance(val, int) else tuple(sorted(val)) for val in control_values
-                )
-            )
-        else:
-            self.control_values = control_values
+
+        self.control_values = cv.to_control_values(control_values)
+
         # Verify control values not out of bounds
-        self.control_values.check_dimentionality(controls=tuple(controls))
+        self.control_values.check_dimensionality(controls=tuple(controls))
 
         if not isinstance(sub_operation, ControlledOperation):
             self.controls = tuple(controls)
@@ -88,6 +82,16 @@ def gate(self) -> Optional['cirq.ControlledGate']:
     def qubits(self):
         return self.controls + self.sub_operation.qubits
 
+    @property
+    def control_values(self) -> cv.ControlValues:
+        return self._control_values
+
+    @control_values.setter
+    def control_values(
+        self, values: Union[cv.ControlValues, Sequence[Union[int, Collection[int]]]]
+    ) -> None:
+        self._control_values = cv.to_control_values(values)
+
     def with_qubits(self, *new_qubits):
         n = len(self.controls)
         return ControlledOperation(