enable simulation of controlled gates in classical simulator

cirq-core/cirq/sim/classical_simulator.py

@@ -117,12 +117,24 @@ def _act_on_fallback_(self, action, qubits: Sequence['cirq.Qid'], allow_decompos
 
         Raises:
             ValueError: If initial_state shape for type np.ndarray is not equal to 1.
-                        If gate is not one of X, CNOT, SWAP, CCNOT, or a measurement.
+                        If gate is not one of X, SWAP, a controlled version of X or SWAP, or a measurement.
         """
         if isinstance(self._state.basis, np.ndarray) and len(self._state.basis.shape) != 1:
             raise ValueError('initial_state shape for type np.ndarray is not equal to 1')
         gate = action.gate if isinstance(action, ops.Operation) else action
         mapped_qubits = [self.qubit_map[i] for i in qubits]
+
+        if isinstance(gate, ops.ControlledGate):
+            control_qubits = mapped_qubits[: gate.num_controls()]
+            mapped_qubits = mapped_qubits[gate.num_controls() :]
+
+            for c, v in zip(control_qubits, gate.control_values):
+                if self._state.basis[c] not in v:
+                    # gate has no effect; controls were off
+                    return True
+
+            gate = gate.sub_gate
+
         if _is_identity(gate):
             pass
         elif gate == ops.X:
@@ -138,7 +150,10 @@ def _act_on_fallback_(self, action, qubits: Sequence['cirq.Qid'], allow_decompos
             c1, c2, q = mapped_qubits
             self._state.basis[q] ^= self._state.basis[c1] & self._state.basis[c2]
         else:
-            raise ValueError(f'{gate} is not one of X, CNOT, SWAP, CCNOT, or a measurement')
+            raise ValueError(
+                f'{gate} is not one of X, SWAP; a controlled version '
+                'of X or SWAP; or a measurement'
+            )
         return True
 
 

cirq-core/cirq/sim/classical_simulator_test.py

@@ -78,6 +78,87 @@ def test_CCNOT():
     np.testing.assert_equal(results, expected_results)
 
 
+def test_CCCX():
+    CCCX = cirq.CCNOT.controlled()
+    qubits = cirq.LineQubit.range(4)
+    circuit = cirq.Circuit()
+
+    for i in range(8):
+        not_idxs = []
+        tmp = i
+        for j in range(3):
+            if tmp & 1:
+                not_idxs.append(j)
+            tmp >>= 1
+
+        circuit.append(cirq.X(qubits[j]) for j in not_idxs)
+        circuit.append(CCCX(*qubits))
+        circuit.append(cirq.measure(qubits, key='key'))
+        circuit.append(cirq.X(qubits[j]) for j in not_idxs)
+
+    expected_results = {
+        'key': np.array(
+            [
+                [
+                    [0, 0, 0, 0],
+                    [1, 0, 0, 0],
+                    [0, 1, 0, 0],
+                    [1, 1, 0, 0],
+                    [0, 0, 1, 0],
+                    [1, 0, 1, 0],
+                    [0, 1, 1, 0],
+                    [1, 1, 1, 1],
+                ]
+            ],
+            dtype=np.uint8,
+        )
+    }
+    sim = cirq.ClassicalStateSimulator()
+    results = sim.run(circuit, param_resolver=None, repetitions=1).records
+    np.testing.assert_equal(results, expected_results)
+
+
+def test_CSWAP():
+    CSWAP = cirq.SWAP.controlled()
+    qubits = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit()
+
+    for i in range(8):
+        not_idxs = []
+        tmp = i
+        for j in range(3):
+            if tmp & 1:
+                not_idxs.append(j)
+            tmp >>= 1
+
+        circuit.append(cirq.X(qubits[j]) for j in not_idxs)
+        circuit.append(CSWAP(*qubits))
+        circuit.append(cirq.measure(qubits, key='key'))
+        circuit.append(CSWAP(*qubits))
+        circuit.append(cirq.X(qubits[j]) for j in not_idxs)
+
+    expected_results = {
+        'key': np.array(
+            [
+                [
+                    [0, 0, 0],
+                    [1, 0, 0],
+                    [0, 1, 0],
+                    [1, 0, 1],
+                    [0, 0, 1],
+                    [1, 1, 0],
+                    [0, 1, 1],
+                    [1, 1, 1],
+                ]
+            ],
+            dtype=np.uint8,
+        )
+    }
+    sim = cirq.ClassicalStateSimulator()
+    results = sim.run(circuit, param_resolver=None, repetitions=1).records
+    np.testing.assert_equal(results, expected_results)
+
+
 def test_measurement_gate():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit()