From e0d99ebe9f0c29b9ad5559efff9c9ab4952c89da Mon Sep 17 00:00:00 2001
From: Craig Gidney <craiggidney@google.com>
Date: Tue, 2 Jul 2024 20:49:18 -0700
Subject: [PATCH] Fix PauliString.before reset not working (#792)

- Also switch `circuit.insert` to doing auto-fusion of operations
---
 doc/python_api_reference_vDev.md              |  10 +-
 doc/stim.pyi                                  |  10 +-
 glue/python/src/stim/__init__.pyi             |  10 +-
 src/stim/circuit/circuit.cc                   |  32 +++-
 src/stim/circuit/circuit.h                    |   1 +
 src/stim/circuit/circuit.pybind.cc            |  10 +-
 src/stim/circuit/circuit.test.cc              | 139 ++++++++++++++++++
 .../stabilizers/pauli_string_pybind_test.py   |  27 ++++
 src/stim/stabilizers/pauli_string_ref.h       |   1 +
 src/stim/stabilizers/pauli_string_ref.inl     |  39 ++++-
 10 files changed, 247 insertions(+), 32 deletions(-)

diff --git a/doc/python_api_reference_vDev.md b/doc/python_api_reference_vDev.md
index d05fb9dc1..8846c0ac8 100644
--- a/doc/python_api_reference_vDev.md
+++ b/doc/python_api_reference_vDev.md
@@ -2236,9 +2236,8 @@ def insert(
 ) -> None:
     """Inserts an operation at the given index, pushing existing operations forward.
 
-    Note that, unlike when appending operations or parsing stim circuit files,
-    inserted operations aren't automatically fused into the preceding operation.
-    This is to avoid creating complicated situations where it's difficult to reason
+    Beware that inserted operations are automatically fused with the preceding
+    and following operations, if possible. This can make it complex to reason
     about how the indices of operations change in response to insertions.
 
     Args:
@@ -2249,7 +2248,7 @@ def insert(
             indices relative to the end of the circuit instead of the start.
 
             Instructions before the index are not shifted. Instructions that
-            were at or after the index are shifted forwards.
+            were at or after the index are shifted forwards as needed.
         operation: The object to insert. This can be a single
             stim.CircuitInstruction or an entire stim.Circuit.
 
@@ -2273,8 +2272,7 @@ def insert(
         stim.Circuit('''
             H 0
             Y 3 4 5
-            S 1
-            S 999
+            S 1 999
             CX 0 1
             CZ 2 3
             X 2
diff --git a/doc/stim.pyi b/doc/stim.pyi
index 8d603dd77..9e15ce302 100644
--- a/doc/stim.pyi
+++ b/doc/stim.pyi
@@ -1607,9 +1607,8 @@ class Circuit:
     ) -> None:
         """Inserts an operation at the given index, pushing existing operations forward.
 
-        Note that, unlike when appending operations or parsing stim circuit files,
-        inserted operations aren't automatically fused into the preceding operation.
-        This is to avoid creating complicated situations where it's difficult to reason
+        Beware that inserted operations are automatically fused with the preceding
+        and following operations, if possible. This can make it complex to reason
         about how the indices of operations change in response to insertions.
 
         Args:
@@ -1620,7 +1619,7 @@ class Circuit:
                 indices relative to the end of the circuit instead of the start.
 
                 Instructions before the index are not shifted. Instructions that
-                were at or after the index are shifted forwards.
+                were at or after the index are shifted forwards as needed.
             operation: The object to insert. This can be a single
                 stim.CircuitInstruction or an entire stim.Circuit.
 
@@ -1644,8 +1643,7 @@ class Circuit:
             stim.Circuit('''
                 H 0
                 Y 3 4 5
-                S 1
-                S 999
+                S 1 999
                 CX 0 1
                 CZ 2 3
                 X 2
diff --git a/glue/python/src/stim/__init__.pyi b/glue/python/src/stim/__init__.pyi
index 8d603dd77..9e15ce302 100644
--- a/glue/python/src/stim/__init__.pyi
+++ b/glue/python/src/stim/__init__.pyi
@@ -1607,9 +1607,8 @@ class Circuit:
     ) -> None:
         """Inserts an operation at the given index, pushing existing operations forward.
 
-        Note that, unlike when appending operations or parsing stim circuit files,
-        inserted operations aren't automatically fused into the preceding operation.
-        This is to avoid creating complicated situations where it's difficult to reason
+        Beware that inserted operations are automatically fused with the preceding
+        and following operations, if possible. This can make it complex to reason
         about how the indices of operations change in response to insertions.
 
         Args:
@@ -1620,7 +1619,7 @@ class Circuit:
                 indices relative to the end of the circuit instead of the start.
 
                 Instructions before the index are not shifted. Instructions that
-                were at or after the index are shifted forwards.
+                were at or after the index are shifted forwards as needed.
             operation: The object to insert. This can be a single
                 stim.CircuitInstruction or an entire stim.Circuit.
 
@@ -1644,8 +1643,7 @@ class Circuit:
             stim.Circuit('''
                 H 0
                 Y 3 4 5
-                S 1
-                S 999
+                S 1 999
                 CX 0 1
                 CZ 2 3
                 X 2
diff --git a/src/stim/circuit/circuit.cc b/src/stim/circuit/circuit.cc
index 615a20e38..86b762f74 100644
--- a/src/stim/circuit/circuit.cc
+++ b/src/stim/circuit/circuit.cc
@@ -225,13 +225,21 @@ void circuit_read_single_operation(Circuit &circuit, char lead_char, SOURCE read
 }
 
 void Circuit::try_fuse_last_two_ops() {
-    auto &ops = operations;
-    size_t n = ops.size();
-    if (n > 1 && ops[n - 2].can_fuse(ops[n - 1])) {
-        fuse_data(ops[n - 2].targets, ops[n - 1].targets, target_buf);
-        ops.pop_back();
+    if (operations.size() >= 2) {
+        try_fuse_after(operations.size() - 2);
     }
 }
+
+void Circuit::try_fuse_after(size_t index) {
+    if (index + 1 >= operations.size()) {
+        return;
+    }
+    if (operations[index].can_fuse(operations[index + 1])) {
+        fuse_data(operations[index].targets, operations[index + 1].targets, target_buf);
+        operations.erase(operations.begin() + index + 1);
+    }
+}
+
 template <typename SOURCE>
 void circuit_read_operations(Circuit &circuit, SOURCE read_char, READ_CONDITION read_condition) {
     auto &ops = circuit.operations;
@@ -358,6 +366,12 @@ void Circuit::safe_insert(size_t index, const CircuitInstruction &instruction) {
     copy.args = arg_buf.take_copy(copy.args);
     copy.targets = target_buf.take_copy(copy.targets);
     operations.insert(operations.begin() + index, copy);
+
+    // Fuse at boundaries.
+    try_fuse_after(index);
+    if (index > 0) {
+        try_fuse_after(index - 1);
+    }
 }
 
 void Circuit::safe_insert(size_t index, const Circuit &circuit) {
@@ -381,6 +395,14 @@ void Circuit::safe_insert(size_t index, const Circuit &circuit) {
             operations[k].args = arg_buf.take_copy(operations[k].args);
         }
     }
+
+    // Fuse at boundaries.
+    if (!circuit.operations.empty()) {
+        try_fuse_after(index + circuit.operations.size() - 1);
+        if (index > 0) {
+            try_fuse_after(index - 1);
+        }
+    }
 }
 
 void Circuit::safe_insert_repeat_block(size_t index, uint64_t repeat_count, const Circuit &block) {
diff --git a/src/stim/circuit/circuit.h b/src/stim/circuit/circuit.h
index e8670e83e..f8268dea8 100644
--- a/src/stim/circuit/circuit.h
+++ b/src/stim/circuit/circuit.h
@@ -280,6 +280,7 @@ struct Circuit {
     std::string describe_instruction_location(size_t instruction_offset) const;
 
     void try_fuse_last_two_ops();
+    void try_fuse_after(size_t index);
 };
 
 void vec_pad_add_mul(std::vector<double> &target, SpanRef<const double> offset, uint64_t mul = 1);
diff --git a/src/stim/circuit/circuit.pybind.cc b/src/stim/circuit/circuit.pybind.cc
index 608b60bff..f364cc4e1 100644
--- a/src/stim/circuit/circuit.pybind.cc
+++ b/src/stim/circuit/circuit.pybind.cc
@@ -1134,9 +1134,8 @@ void stim_pybind::pybind_circuit_methods(pybind11::module &, pybind11::class_<Ci
             Inserts an operation at the given index, pushing existing operations forward.
             @signature def insert(self, index: int, operation: Union[stim.CircuitInstruction, stim.Circuit]) -> None:
 
-            Note that, unlike when appending operations or parsing stim circuit files,
-            inserted operations aren't automatically fused into the preceding operation.
-            This is to avoid creating complicated situations where it's difficult to reason
+            Beware that inserted operations are automatically fused with the preceding
+            and following operations, if possible. This can make it complex to reason
             about how the indices of operations change in response to insertions.
 
             Args:
@@ -1147,7 +1146,7 @@ void stim_pybind::pybind_circuit_methods(pybind11::module &, pybind11::class_<Ci
                     indices relative to the end of the circuit instead of the start.
 
                     Instructions before the index are not shifted. Instructions that
-                    were at or after the index are shifted forwards.
+                    were at or after the index are shifted forwards as needed.
                 operation: The object to insert. This can be a single
                     stim.CircuitInstruction or an entire stim.Circuit.
 
@@ -1171,8 +1170,7 @@ void stim_pybind::pybind_circuit_methods(pybind11::module &, pybind11::class_<Ci
                 stim.Circuit('''
                     H 0
                     Y 3 4 5
-                    S 1
-                    S 999
+                    S 1 999
                     CX 0 1
                     CZ 2 3
                     X 2
diff --git a/src/stim/circuit/circuit.test.cc b/src/stim/circuit/circuit.test.cc
index cdde5e379..a35928882 100644
--- a/src/stim/circuit/circuit.test.cc
+++ b/src/stim/circuit/circuit.test.cc
@@ -1769,3 +1769,142 @@ TEST(circuit, generate_test_circuit_with_all_operations) {
     }
     ASSERT_EQ(seen.size(), NUM_DEFINED_GATES);
 }
+
+TEST(circuit, insert_circuit) {
+    Circuit c(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(2, Circuit(R"CIRCUIT(
+        H 1
+        X 3
+        S 2
+    )CIRCUIT"));
+    ASSERT_EQ(c.operations.size(), 5);
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0 1
+        X 3
+        S 2 0
+        CX 0 1
+    )CIRCUIT"));
+
+    c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(0, Circuit(R"CIRCUIT(
+        H 1
+        X 3
+        S 2
+    )CIRCUIT"));
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        H 1
+        X 3
+        S 2
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT"));
+
+    c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(4, Circuit(R"CIRCUIT(
+        H 1
+        X 3
+        S 2
+    )CIRCUIT"));
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+        H 1
+        X 3
+        S 2
+    )CIRCUIT"));
+}
+
+TEST(circuit, insert_instruction) {
+    Circuit c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(2, Circuit("H 1").operations[0]);
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0 1
+        S 0
+        CX 0 1
+    )CIRCUIT"));
+
+    c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(2, Circuit("S 1").operations[0]);
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 1 0
+        CX 0 1
+    )CIRCUIT"));
+
+    c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(2, Circuit("X 1").operations[0]);
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        X 1
+        S 0
+        CX 0 1
+    )CIRCUIT"));
+
+    c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(0, Circuit("X 1").operations[0]);
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        X 1
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT"));
+
+    c = Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+    )CIRCUIT");
+    c.safe_insert(4, Circuit("X 1").operations[0]);
+    ASSERT_EQ(c, Circuit(R"CIRCUIT(
+        CX 0 1
+        H 0
+        S 0
+        CX 0 1
+        X 1
+    )CIRCUIT"));
+}
diff --git a/src/stim/stabilizers/pauli_string_pybind_test.py b/src/stim/stabilizers/pauli_string_pybind_test.py
index 5058e6a55..4d7b5af1d 100644
--- a/src/stim/stabilizers/pauli_string_pybind_test.py
+++ b/src/stim/stabilizers/pauli_string_pybind_test.py
@@ -928,3 +928,30 @@ def test_pauli_indices():
     assert stim.PauliString("_XYZ").pauli_indices("_") == [0]
     with pytest.raises(ValueError, match="Invalid character"):
         assert stim.PauliString("_XYZ").pauli_indices("k")
+
+
+def test_before_reset():
+    assert stim.PauliString("Z").before(stim.Circuit("R 0")) == stim.PauliString("_")
+    assert stim.PauliString("Z").before(stim.Circuit("MR 0")) == stim.PauliString("_")
+    assert stim.PauliString("Z").before(stim.Circuit("M 0")) == stim.PauliString("Z")
+
+    assert stim.PauliString("X").before(stim.Circuit("RX 0")) == stim.PauliString("_")
+    assert stim.PauliString("X").before(stim.Circuit("MRX 0")) == stim.PauliString("_")
+    assert stim.PauliString("X").before(stim.Circuit("MX 0")) == stim.PauliString("X")
+
+    assert stim.PauliString("Y").before(stim.Circuit("RY 0")) == stim.PauliString("_")
+    assert stim.PauliString("Y").before(stim.Circuit("MRY 0")) == stim.PauliString("_")
+    assert stim.PauliString("Y").before(stim.Circuit("MY 0")) == stim.PauliString("Y")
+
+    with pytest.raises(ValueError):
+        stim.PauliString("Z").before(stim.Circuit("RX 0"))
+    with pytest.raises(ValueError):
+        stim.PauliString("Z").before(stim.Circuit("RY 0"))
+    with pytest.raises(ValueError):
+        stim.PauliString("Z").before(stim.Circuit("MRX 0"))
+    with pytest.raises(ValueError):
+        stim.PauliString("Z").before(stim.Circuit("MRY 0"))
+    with pytest.raises(ValueError):
+        stim.PauliString("Z").before(stim.Circuit("MX 0"))
+    with pytest.raises(ValueError):
+        stim.PauliString("Z").before(stim.Circuit("MY 0"))
diff --git a/src/stim/stabilizers/pauli_string_ref.h b/src/stim/stabilizers/pauli_string_ref.h
index d93ca517f..b87599b39 100644
--- a/src/stim/stabilizers/pauli_string_ref.h
+++ b/src/stim/stabilizers/pauli_string_ref.h
@@ -159,6 +159,7 @@ struct PauliStringRef {
     void check_avoids_MPP(const CircuitInstruction &inst);
     void check_avoids_reset(const CircuitInstruction &inst);
     void check_avoids_measurement(const CircuitInstruction &inst);
+    void undo_reset_xyz(const CircuitInstruction &inst);
 
     void do_single_cx(const CircuitInstruction &inst, uint32_t c, uint32_t t);
     void do_single_cy(const CircuitInstruction &inst, uint32_t c, uint32_t t);
diff --git a/src/stim/stabilizers/pauli_string_ref.inl b/src/stim/stabilizers/pauli_string_ref.inl
index 5b49756a1..73c5e5461 100644
--- a/src/stim/stabilizers/pauli_string_ref.inl
+++ b/src/stim/stabilizers/pauli_string_ref.inl
@@ -172,6 +172,39 @@ void PauliStringRef<W>::do_circuit(const Circuit &circuit) {
     });
 }
 
+template <size_t W>
+void PauliStringRef<W>::undo_reset_xyz(const CircuitInstruction &inst) {
+    bool x_dep, z_dep;
+    if (inst.gate_type == GateType::R || inst.gate_type == GateType::MR) {
+        x_dep = true;
+        z_dep = false;
+    } else if (inst.gate_type == GateType::RX || inst.gate_type == GateType::MRX) {
+        x_dep = false;
+        z_dep = true;
+    } else if (inst.gate_type == GateType::RY || inst.gate_type == GateType::MRY) {
+        x_dep = true;
+        z_dep = true;
+    } else {
+        throw std::invalid_argument("Unrecognized measurement type: " + inst.str());
+    }
+    for (const auto &t : inst.targets) {
+        assert(t.is_qubit_target());
+        auto q = t.qubit_value();
+        if (q < num_qubits && ((xs[q] & x_dep) ^ (zs[q] & z_dep))) {
+            std::stringstream ss;
+            ss << "The pauli observable '" << *this;
+            ss << "' doesn't have a well specified value before '" << inst;
+            ss << "' because it anticommutes with the reset.";
+            throw std::invalid_argument(ss.str());
+        }
+    }
+    for (const auto &t : inst.targets) {
+        auto q = t.qubit_value();
+        xs[q] = false;
+        zs[q] = false;
+    }
+}
+
 template <size_t W>
 void PauliStringRef<W>::check_avoids_measurement(const CircuitInstruction &inst) {
     bool x_dep, z_dep;
@@ -193,7 +226,7 @@ void PauliStringRef<W>::check_avoids_measurement(const CircuitInstruction &inst)
         if (q < num_qubits && ((xs[q] & x_dep) ^ (zs[q] & z_dep))) {
             std::stringstream ss;
             ss << "The pauli observable '" << *this;
-            ss << "' doesn't have a well specified value after '" << inst;
+            ss << "' doesn't have a well specified value across '" << inst;
             ss << "' because it anticommutes with the measurement.";
             throw std::invalid_argument(ss.str());
         }
@@ -238,7 +271,7 @@ void PauliStringRef<W>::check_avoids_MPP(const CircuitInstruction &inst) {
         if (anticommutes) {
             std::stringstream ss;
             ss << "The pauli observable '" << *this;
-            ss << "' doesn't have a well specified value after '" << inst;
+            ss << "' doesn't have a well specified value across '" << inst;
             ss << "' because it anticommutes with the measurement.";
             throw std::invalid_argument(ss.str());
         }
@@ -552,7 +585,7 @@ void PauliStringRef<W>::undo_instruction(const CircuitInstruction &inst) {
         case GateType::MR:
         case GateType::MRX:
         case GateType::MRY:
-            check_avoids_reset(inst);
+            undo_reset_xyz(inst);
             break;
 
         case GateType::M: