Add stimcirq support for feedback controlled by sympy Xor expressions (#1058)

Fixes #1053

Author: Strilanc

glue/cirq/stimcirq/_cirq_to_stim.py

@@ -5,6 +5,7 @@
 
 import cirq
 import stim
+import sympy
 
 from ._ii_gate import IIGate
 
@@ -152,10 +153,15 @@ def _stim_append_classically_controlled_gate(
 
     if len(op.classical_controls) != 1:
         raise NotImplementedError(f'Stim only supports single-control Pauli feedback, but got {op=}')
-    control, = op.classical_controls
-    if not isinstance(control, cirq.KeyCondition):
-        raise NotImplementedError(f'Stim only supports single-control Pauli feedback (i.e. a `cirq.KeyCondition` control), but got {control=}')
-    control: cirq.KeyCondition
+    controls: list[cirq.KeyCondition] = []
+    single_control, = op.classical_controls
+    if isinstance(single_control, cirq.KeyCondition):
+        controls.append(single_control)
+    elif isinstance(single_control, cirq.SympyCondition) and isinstance(single_control.expr, sympy.Xor) and all(isinstance(e, sympy.Symbol) for e in single_control.expr.args):
+        for symbol in single_control.expr.args:
+            controls.append(cirq.KeyCondition(key=cirq.MeasurementKey(str(symbol)), index=-1))
+    else:
+        raise NotImplementedError(f'Stim only supports single-control Pauli feedback (i.e. a `cirq.KeyCondition` control), but got {single_control=}')
     gate = op.without_classical_controls().gate
 
     if gate == cirq.X:
@@ -168,24 +174,25 @@ def _stim_append_classically_controlled_gate(
         raise NotImplementedError(f'Stim only supports Pauli feedback, but got {op=}')
     assert len(targets) == 1
 
-    skips_left = control.index
-    for offset in range(len(measurement_key_lengths)):
-        m_key, m_len = measurement_key_lengths[-1 - offset]
-        if m_len != 1:
-            raise NotImplementedError(f"multi-qubit measurement {m_key!r}")
-        if m_key == control.key:
-            if skips_left > 0:
-                skips_left -= 1
-            else:
-                rec_target = stim.target_rec(-1 - offset)
-                break
-    else:
-        raise ValueError(
-            f"{control!r} was processed before the measurement it referenced."
-            f" Make sure the referenced measurements keys are actually in the circuit, and come"
-            f" in an earlier moment (or earlier in the same moment's operation order)."
-        )
-    circuit.append(f"C{stim_gate}", [rec_target, targets[0]], tag=tag)
+    for control in controls:
+        skips_left = control.index
+        for offset in range(len(measurement_key_lengths)):
+            m_key, m_len = measurement_key_lengths[-1 - offset]
+            if m_len != 1:
+                raise NotImplementedError(f"multi-qubit measurement {m_key!r}")
+            if m_key == control.key:
+                if skips_left > 0:
+                    skips_left -= 1
+                else:
+                    rec_target = stim.target_rec(-1 - offset)
+                    break
+        else:
+            raise ValueError(
+                f"{control!r} was processed before the measurement it referenced."
+                f" Make sure the referenced measurements keys are actually in the circuit, and come"
+                f" in an earlier moment (or earlier in the same moment's operation order)."
+            )
+        circuit.append(f"C{stim_gate}", [rec_target, targets[0]], tag=tag)
 
 
 @functools.lru_cache(maxsize=1)

glue/cirq/stimcirq/_cirq_to_stim_test.py

@@ -6,6 +6,7 @@
 import pytest
 import stim
 import stimcirq
+import sympy
 from stimcirq._cirq_to_stim import cirq_circuit_to_stim_data, gate_to_stim_append_func
 
 
@@ -425,3 +426,30 @@ def test_round_trip_example_circuit():
     cirq_circuit = stimcirq.stim_circuit_to_cirq_circuit(stim_circuit.flattened())
     circuit_back = stimcirq.cirq_circuit_to_stim_circuit(cirq_circuit)
     assert len(circuit_back.shortest_graphlike_error()) == 3
+
+
+def test_xor_feedback():
+    a, b, c, d, e = cirq.LineQubit.range(5)
+    cirq_circuit = cirq.Circuit([
+        cirq.Moment(
+            cirq.measure(a, key='a'),
+            cirq.measure(b, key='b'),
+            cirq.measure(c, key='c'),
+            cirq.measure(d, key='d'),
+        ),
+        cirq.Moment(
+            cirq.X(e).with_classical_controls(cirq.SympyCondition(sympy.Xor(
+                sympy.Symbol('a'),
+                sympy.Symbol('b'),
+                sympy.Symbol('c'),
+                sympy.Symbol('d'),
+            ))),
+        ),
+    ])
+    stim_circuit = stimcirq.cirq_circuit_to_stim_circuit(cirq_circuit)
+    assert stim_circuit == stim.Circuit("""
+        M 0 1 2 3
+        TICK
+        CX rec[-4] 4 rec[-3] 4 rec[-2] 4 rec[-1] 4
+        TICK
+    """)