diff --git a/glue/stimflow/src/stimflow/_layers/_layer_tag.py b/glue/stimflow/src/stimflow/_layers/_layer_tag.py index e858ca74..7963c830 100644 --- a/glue/stimflow/src/stimflow/_layers/_layer_tag.py +++ b/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] diff --git a/glue/stimflow/src/stimflow/_layers/_layer_tag_test.py b/glue/stimflow/src/stimflow/_layers/_layer_tag_test.py index decfa8b6..c0a79c4a 100644 --- a/glue/stimflow/src/stimflow/_layers/_layer_tag_test.py +++ b/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} \ No newline at end of file