Allow quantum if with multiple effects

unitary/alpha/quantum_effect.py

@@ -100,7 +100,7 @@ class QuantumThen(QuantumEffect):
     def __init__(self, *objects: "QuantumObject"):
         self.control_objects = list(objects)
         self.condition = [1] * len(self.control_objects)
-        self.then_effect = None
+        self.then_effects = None
 
     def equals(
         self, *conditions: Union[enum.Enum, int, Sequence[Union[enum.Enum, int]]]
@@ -124,9 +124,13 @@ def then(self, effect: "QuantumEffect"):
         """Use `apply(effect)` instead."""
         return self.apply(effect)
 
-    def apply(self, effect: "QuantumEffect"):
-        """Applies a QuantumEffect conditionally to the specified qubits."""
-        self.then_effect = effect
+    def apply(self, *effects: "QuantumEffect"):
+        """Applies a QuantumEffect conditionally to the specified qubits.
+
+        If multiple effects are specified, they will all be applied
+        (to the same qubits).
+        """
+        self.then_effects = list(effects)
         return self
 
     def effect(self, *objects: "QuantumObject"):
@@ -136,8 +140,9 @@ def effect(self, *objects: "QuantumObject"):
             if cond == 0 and self.control_objects[idx].num_states == 2:
                 yield cirq.X(self.control_objects[idx].qubit)
 
-        for op in self.then_effect.effect(*objects):
-            yield op.controlled_by(*[q.qubit for q in self.control_objects])
+        for effect in self.then_effects:
+            for op in effect.effect(*objects):
+                yield op.controlled_by(*[q.qubit for q in self.control_objects])
 
         # For anti-controls, add an X after the controlled operation
         # to revert its state back to what it was.

unitary/alpha/quantum_effect_test.py

@@ -68,6 +68,22 @@ def test_anti_control(simulator, compile_to_qubits):
     assert (result[0] == 1 for result in results)
 
 
+@pytest.mark.parametrize("compile_to_qubits", [False, True])
+@pytest.mark.parametrize("simulator", [cirq.Simulator, SparseSimulator])
+def test_multiple_effects_quantum_if(simulator, compile_to_qubits):
+    board = alpha.QuantumWorld(sampler=simulator(), compile_to_qubits=compile_to_qubits)
+    piece = alpha.QuantumObject("q0", 1)
+    piece2 = alpha.QuantumObject("q1", 0)
+    board.add_object(piece)
+    board.add_object(piece2)
+    alpha.quantum_if(piece).apply(
+        alpha.Flip(effect_fraction=0.5), alpha.Flip(effect_fraction=0.5)
+    )(piece2)
+
+    # Test that results are as expected
+    results = board.peek([piece, piece2], count=100)
+
+
 def test_no_world():
     piece = alpha.QuantumObject("q0", 1)
     with pytest.raises(ValueError, match="must be added to a QuantumWorld"):