diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 3dd8c5d8..1381cbee 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -555,7 +555,11 @@ def strip(circuit: stim.Circuit) -> stim.Circuit: for instr in circuit: if isinstance(instr, stim.CircuitRepeatBlock): stripped = strip(instr.body_copy()) - result.append(stim.CircuitRepeatBlock(instr.repeat_count, stripped)) + result.append( + stim.CircuitRepeatBlock( + instr.repeat_count, stripped, tag=instr.tag + ) + ) continue if instr.name in ["OBSERVABLE_INCLUDE", "DETECTOR"]: continue @@ -881,7 +885,11 @@ def fix_tags(circuit: stim.Circuit) -> stim.Circuit: for instr in circuit: if isinstance(instr, stim.CircuitRepeatBlock): fixed = fix_tags(instr.body_copy()) - result.append(stim.CircuitRepeatBlock(instr.repeat_count, fixed)) + result.append( + stim.CircuitRepeatBlock( + instr.repeat_count, fixed, tag=instr.tag + ) + ) continue name = instr.name diff --git a/src/tsim/utils/clifford.py b/src/tsim/utils/clifford.py index 6d5fc02e..cb283754 100644 --- a/src/tsim/utils/clifford.py +++ b/src/tsim/utils/clifford.py @@ -157,7 +157,9 @@ def expand_clifford_rotations(source: stim.Circuit) -> stim.Circuit: if isinstance(instr, stim.CircuitRepeatBlock): out.append( stim.CircuitRepeatBlock( - instr.repeat_count, expand_clifford_rotations(instr.body_copy()) + instr.repeat_count, + expand_clifford_rotations(instr.body_copy()), + tag=instr.tag, ) ) continue diff --git a/src/tsim/utils/diagram.py b/src/tsim/utils/diagram.py index 166c7bba..0e3240ba 100644 --- a/src/tsim/utils/diagram.py +++ b/src/tsim/utils/diagram.py @@ -383,7 +383,9 @@ def _replace_tagged_gates( if isinstance(instr, stim.CircuitRepeatBlock): modified_body = _replace_tagged_gates(instr.body_copy(), replace_dict) modified_circ.append( - stim.CircuitRepeatBlock(instr.repeat_count, modified_body) + stim.CircuitRepeatBlock( + instr.repeat_count, modified_body, tag=instr.tag + ) ) continue diff --git a/test/unit/test_circuit.py b/test/unit/test_circuit.py index 46644c7d..43dacf3b 100644 --- a/test/unit/test_circuit.py +++ b/test/unit/test_circuit.py @@ -754,6 +754,20 @@ def test_stim_circuit_repeat_block_preserves_non_clifford(): assert block.repeat_count == 2 +def test_stim_circuit_preserves_repeat_block_tags(): + """Tags on REPEAT blocks survive the round-trip through stim_circuit.""" + c = Circuit( + "REPEAT[outer] 3 {\n H 0\n REPEAT[inner] 2 {\n CX 0 1\n }\n}" + ) + stim_c = c.stim_circuit + outer = stim_c[0] + assert isinstance(outer, stim.CircuitRepeatBlock) + assert outer.tag == "outer" + inner = outer.body_copy()[-1] + assert isinstance(inner, stim.CircuitRepeatBlock) + assert inner.tag == "inner" + + def test_get_graph(): """Test get_graph returns a ZX graph.""" c = Circuit("H 0\nCNOT 0 1")