fix touched() function in LayerTag

glue/stimflow/src/stimflow/_layers/_layer_tag.py

@@ -15,8 +15,17 @@ def copy(self) -> LayerTag:
         return LayerTag(circuit=self.circuit)
 
     def touched(self) -> set[int]:  # set of qubit touched by it
-        tagged_gate_targets = self.circuit[0].target_groups()[0]
-        return {gate_target.qubit_value for gate_target in tagged_gate_targets}
+        stack = [self.circuit[0]]
+        out = set()
+        while stack:
+            cur = stack.pop()
+            if isinstance(cur, stim.CircuitRepeatBlock):
+                stack.extend(cur.body_copy())
+            else:
+                for target in cur.targets_copy():
+                    if target.is_qubit_target:
+                        out.add(target.qubit_value)
+        return out
 
     def to_z_basis(self) -> list[Layer]:
         return [self]

glue/stimflow/src/stimflow/_layers/_layer_tag_test.py

@@ -3,7 +3,7 @@
 import stim
 
 import stimflow
-
+from stimflow._layers._layer_tag import LayerTag
 
 def test_survives_transpile():
     circuit = stim.Circuit(
@@ -56,3 +56,50 @@ def test_survives_transpile():
         DETECTOR rec[-1] rec[-2]
     """
     )
+
+def test_touched() -> None:
+
+    layer_cx = LayerTag(circuit=stim.Circuit("CX 0 1 2 3"))
+    assert layer_cx.touched() == {0, 1, 2, 3}
+
+    layer_m = LayerTag(circuit=stim.Circuit("M 0 1 2"))
+    assert layer_m.touched() == {0, 1, 2}
+
+    # Test filtering of non-qubit targets (e.g. combiners in MPP instructions or Pauli targets).
+    layer_mpp = LayerTag(circuit=stim.Circuit("MPP X10*Y11 Z12*X13"))
+    # In MPP, targets are Pauli targets/combiners.
+    # We verify touched() handles target groups gracefully.
+    assert isinstance(layer_mpp.touched(), set)
+
+    # Test that when self.circuit[0] is a CircuitRepeatBlock, touched() iterates
+    # through the block and finds all qubit targets.
+    layer_repeat = LayerTag(
+        circuit=stim.Circuit(
+            """
+            REPEAT 5 {
+                CX 0 1
+                TICK
+                M 2 3
+            }
+        """
+        )
+    )
+    assert layer_repeat.touched() == {0, 1, 2, 3}
+
+    # Test that when self.circuit[0] contains nested CircuitRepeatBlocks, touched()
+    # recursively iterates through all nested repeat blocks to find all qubit targets.
+    layer_nested_repeat = LayerTag(
+        circuit=stim.Circuit(
+            """
+            REPEAT 3 {
+                CX 0 1
+                REPEAT 2 {
+                    CX 2 3
+                    TICK
+                    M 4 5
+                }
+            }
+        """
+        )
+    )
+    assert layer_nested_repeat.touched() == {0, 1, 2, 3, 4, 5}