From a0b032140344d4fcf0f094fd31fb3a0a3715242a Mon Sep 17 00:00:00 2001
From: Craig Gidney <craiggidney@google.com>
Date: Mon, 27 Jan 2025 03:31:32 -0800
Subject: [PATCH] Fix stimcirq failing to round trip stim's example surface
 code circuit (#874)

- Fix MeasureAndOrResetGate not putting its key into the record
---
 glue/cirq/stimcirq/_cirq_to_stim_test.py         | 14 +++++++++++++-
 glue/cirq/stimcirq/_measure_and_or_reset_gate.py |  6 ++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/glue/cirq/stimcirq/_cirq_to_stim_test.py b/glue/cirq/stimcirq/_cirq_to_stim_test.py
index 977d12d9..945dd67b 100644
--- a/glue/cirq/stimcirq/_cirq_to_stim_test.py
+++ b/glue/cirq/stimcirq/_cirq_to_stim_test.py
@@ -377,7 +377,7 @@ def test_on_tagged_loop():
             repetitions=3,
         ).with_tags('my_tag')
     )
-    
+
     stim_circuit = stimcirq.cirq_circuit_to_stim_circuit(c)
     assert stim.CircuitRepeatBlock in {type(instr) for instr in stim_circuit}
 
@@ -410,3 +410,15 @@ def test_custom_tagging():
         H[PAIR] 0
         TICK
     """)
+
+
+def test_round_trip_example_circuit():
+    stim_circuit = stim.Circuit.generated(
+        "surface_code:rotated_memory_x",
+        distance=3,
+        rounds=1,
+        after_clifford_depolarization=0.01,
+    )
+    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
diff --git a/glue/cirq/stimcirq/_measure_and_or_reset_gate.py b/glue/cirq/stimcirq/_measure_and_or_reset_gate.py
index 4e4aaf59..84002238 100644
--- a/glue/cirq/stimcirq/_measure_and_or_reset_gate.py
+++ b/glue/cirq/stimcirq/_measure_and_or_reset_gate.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Tuple
 
 import cirq
 import stim
@@ -91,7 +91,9 @@ def _stim_op_name(self) -> str:
             result += self.basis
         return result
 
-    def _stim_conversion_(self, *, edit_circuit: stim.Circuit, targets: List[int], tag: str, **kwargs):
+    def _stim_conversion_(self, *, edit_circuit: stim.Circuit, targets: List[int], tag: str, edit_measurement_key_lengths: List[Tuple[str, int]], **kwargs):
+        if self.measure:
+            edit_measurement_key_lengths.append((self.key, 1))
         if self.invert_measure:
             targets[0] = stim.target_inv(targets[0])
         if self.measure_flip_probability: