Add stim::ReferenceSampleTree to support loop folded reference samp…

…ling (#772)

Get a compressed tree representation of the reference sample by using
`stim::ReferenceSampleTree::from_circuit_reference_sample(circuit)`.

The reason for a tree, instead of raw run length encoding, is to support
nested loops (e.g. physical surface code rounds repeating within
repeated logical operations).

Example test:

```
    CircuitGenParameters params(10000, 3, "rotated_memory_x");
    auto circuit = generate_surface_code_circuit(params).circuit;
    circuit.blocks[0].append_from_text("X 10 11 12 13");
    auto ref = ReferenceSampleTree::from_circuit_reference_sample(circuit);
    ASSERT_EQ(ref.str(), "1*(''+2*('00000000')+4999*('0110000000110000')+1*('000000000'))");
```

Example benchmark (25 milliseconds to do a distance 31 surface code with
a billion rounds):

```
    CircuitGenParameters params(1000000000, 31, "rotated_memory_x");
    auto circuit = generate_surface_code_circuit(params).circuit;
    simd_bits<MAX_BITWORD_WIDTH> ref(0);
    auto total = 0;
    benchmark_go([&]() {
        auto result = ReferenceSampleTree::from_circuit_reference_sample(circuit);
        total += result.empty();
    })
        .goal_millis(25);
```

Part of #768

file_lists/perf_files

@@ -18,4 +18,5 @@ src/stim/stabilizers/tableau.perf.cc
 src/stim/stabilizers/tableau_iter.perf.cc
 src/stim/util_bot/error_decomp.perf.cc
 src/stim/util_bot/probability_util.perf.cc
+src/stim/util_top/reference_sample_tree.perf.cc
 src/stim/util_top/stabilizers_to_tableau.perf.cc

file_lists/source_files_no_main

@@ -94,5 +94,6 @@ src/stim/util_top/circuit_vs_amplitudes.cc
 src/stim/util_top/export_crumble_url.cc
 src/stim/util_top/export_qasm.cc
 src/stim/util_top/export_quirk_url.cc
+src/stim/util_top/reference_sample_tree.cc
 src/stim/util_top/simplified_circuit.cc
 src/stim/util_top/transform_without_feedback.cc

file_lists/test_files

@@ -89,6 +89,7 @@ src/stim/util_top/export_crumble_url.test.cc
 src/stim/util_top/export_qasm.test.cc
 src/stim/util_top/export_quirk_url.test.cc
 src/stim/util_top/has_flow.test.cc
+src/stim/util_top/reference_sample_tree.test.cc
 src/stim/util_top/simplified_circuit.test.cc
 src/stim/util_top/stabilizers_to_tableau.test.cc
 src/stim/util_top/stabilizers_vs_amplitudes.test.cc

src/stim.h

@@ -115,6 +115,7 @@
 #include "stim/util_top/export_qasm.h"
 #include "stim/util_top/export_quirk_url.h"
 #include "stim/util_top/has_flow.h"
+#include "stim/util_top/reference_sample_tree.h"
 #include "stim/util_top/simplified_circuit.h"
 #include "stim/util_top/stabilizers_to_tableau.h"
 #include "stim/util_top/stabilizers_vs_amplitudes.h"

src/stim/io/measure_record.cc

@@ -53,3 +53,22 @@ void MeasureRecord::record_result(bool result) {
     storage.push_back(result);
     unwritten++;
 }
+
+void MeasureRecord::record_results(const std::vector<bool> &results) {
+    storage.insert(storage.end(), results.begin(), results.end());
+    unwritten += results.size();
+}
+
+void MeasureRecord::clear() {
+    unwritten = 0;
+    storage.clear();
+}
+
+void MeasureRecord::discard_results_past_max_lookback() {
+    if (storage.size() > max_lookback) {
+        storage.erase(storage.begin(), storage.end() - max_lookback);
+    }
+    if (unwritten > max_lookback) {
+        unwritten = max_lookback;
+    }
+}

src/stim/io/measure_record.h

@@ -47,8 +47,14 @@ struct MeasureRecord {
     /// Args:
     ///     lookback: How far back the measurement is. lookback=1 is the latest measurement, 2 the second latest, etc.
     bool lookback(size_t lookback) const;
+    /// Batch record.
+    void record_results(const std::vector<bool> &results);
     /// Appends a measurement to the record.
     void record_result(bool result);
+    /// Clear the record.
+    void clear();
+    /// Truncates the record to only include bits within the lookback limit.
+    void discard_results_past_max_lookback();
 };
 
 }  // namespace stim

src/stim/simulators/tableau_simulator.perf.cc

@@ -14,6 +14,7 @@
 
 #include "stim/simulators/tableau_simulator.h"
 
+#include "stim/gen/circuit_gen_params.h"
 #include "stim/perf.perf.h"
 
 using namespace stim;

src/stim/util_top/reference_sample_tree.cc

@@ -0,0 +1,207 @@
+#include "stim/util_top/reference_sample_tree.h"
+
+using namespace stim;
+
+bool ReferenceSampleTree::empty() const {
+    if (repetitions == 0) {
+        return true;
+    }
+    if (!prefix_bits.empty()) {
+        return false;
+    }
+    for (const auto &child: suffix_children) {
+        if (!child.empty()) {
+            return false;
+        }
+    }
+    return true;
+}
+
+void ReferenceSampleTree::flatten_and_simplify_into(std::vector<ReferenceSampleTree> &out) const {
+    if (repetitions == 0) {
+        return;
+    }
+
+    // Flatten children.
+    std::vector<ReferenceSampleTree> flattened;
+    if (!prefix_bits.empty()) {
+        flattened.push_back(ReferenceSampleTree{
+            .prefix_bits = prefix_bits,
+            .suffix_children = {},
+            .repetitions = 1,
+        });
+    }
+    for (const auto &child : suffix_children) {
+        child.flatten_and_simplify_into(flattened);
+    }
+
+    // Fuse children.
+    std::vector<ReferenceSampleTree> fused;
+    if (!flattened.empty()) {
+        fused.push_back(std::move(flattened[0]));
+    }
+    for (size_t k = 1; k < flattened.size(); k++) {
+        auto &dst = fused.back();
+        auto &src = flattened[k];
+
+        // Combine children with identical contents by adding their rep counts.
+        if (dst.prefix_bits == src.prefix_bits && dst.suffix_children == src.suffix_children) {
+            dst.repetitions += src.repetitions;
+
+        // Fuse children with unrepeated contents if they can be fused.
+        } else if (src.repetitions == 1 && dst.repetitions == 1 && dst.suffix_children.empty()) {
+            dst.suffix_children = std::move(src.suffix_children);
+            dst.prefix_bits.insert(dst.prefix_bits.end(), src.prefix_bits.begin(), src.prefix_bits.end());
+
+        } else {
+            fused.push_back(std::move(src));
+        }
+    }
+
+    if (repetitions == 1) {
+        // Un-nest all the children.
+        for (auto &e : fused) {
+            out.push_back(e);
+        }
+    } else if (fused.size() == 1) {
+        // Merge with single child.
+        fused[0].repetitions *= repetitions;
+        out.push_back(std::move(fused[0]));
+    } else if (fused.empty()) {
+        // Nothing to report.
+    } else if (fused[0].suffix_children.empty() && fused[0].repetitions == 1) {
+        // Take payload from first child.
+        ReferenceSampleTree result = std::move(fused[0]);
+        fused.erase(fused.begin());
+        result.repetitions = repetitions;
+        result.suffix_children = std::move(fused);
+        out.push_back(std::move(result));
+    } else {
+        out.push_back(ReferenceSampleTree{
+            .prefix_bits={},
+            .suffix_children=std::move(fused),
+            .repetitions=repetitions,
+        });
+    }
+}
+
+/// Finds how far back feedback operations ever look, within the loop.
+uint64_t stim::max_feedback_lookback_in_loop(const Circuit &loop) {
+    uint64_t furthest_lookback = 0;
+    for (const auto &inst : loop.operations) {
+        if (inst.gate_type == GateType::REPEAT) {
+            furthest_lookback = std::max(furthest_lookback, max_feedback_lookback_in_loop(inst.repeat_block_body(loop)));
+        } else {
+            auto f = GATE_DATA[inst.gate_type].flags;
+            if ((f & GateFlags::GATE_CAN_TARGET_BITS) && (f & GateFlags::GATE_TARGETS_PAIRS)) {
+                // Feedback-capable operation. Check for any measurement record targets.
+                for (auto t : inst.targets) {
+                    if (t.is_measurement_record_target()) {
+                        furthest_lookback = std::max(furthest_lookback, (uint64_t)-t.rec_offset());
+                    }
+                }
+            }
+        }
+    }
+    return furthest_lookback;
+}
+
+void ReferenceSampleTree::try_factorize(size_t period_factor) {
+    if (prefix_bits.size() != 0 || suffix_children.size() % period_factor != 0) {
+        return;
+    }
+
+    // Check if contents are periodic with the factor.
+    size_t h = suffix_children.size() / period_factor;
+    for (size_t k = h; k < suffix_children.size(); k++) {
+        if (suffix_children[k - h] != suffix_children[k]) {
+            return;
+        }
+    }
+
+    // Factorize.
+    suffix_children.resize(h);
+    repetitions *= period_factor;
+}
+
+ReferenceSampleTree ReferenceSampleTree::simplified() const {
+    std::vector<ReferenceSampleTree> flat;
+    flatten_and_simplify_into(flat);
+    if (flat.empty()) {
+        return ReferenceSampleTree();
+    } else if (flat.size() == 1) {
+        return std::move(flat[0]);
+    }
+
+    ReferenceSampleTree result;
+    result.repetitions = 1;
+
+     // Take payload from first child.
+    if (flat[0].repetitions == 1 && flat[0].suffix_children.empty()) {
+         result = std::move(flat[0]);
+         flat.erase(flat.begin());
+    }
+
+    result.suffix_children = std::move(flat);
+    return result;
+}
+
+size_t ReferenceSampleTree::size() const {
+    size_t result = prefix_bits.size();
+    for (const auto &child: suffix_children) {
+        result += child.size();
+    }
+    return result * repetitions;
+}
+
+void ReferenceSampleTree::decompress_into(std::vector<bool> &output) const {
+    for (uint64_t k = 0; k < repetitions; k++) {
+        output.insert(output.end(), prefix_bits.begin(), prefix_bits.end());
+        for (const auto &child: suffix_children) {
+            child.decompress_into(output);
+        }
+    }
+}
+
+ReferenceSampleTree ReferenceSampleTree::from_circuit_reference_sample(const Circuit &circuit) {
+    auto stats = circuit.compute_stats();
+    std::mt19937_64 irrelevant_rng{0};
+    CompressedReferenceSampleHelper<MAX_BITWORD_WIDTH> helper(
+        TableauSimulator<MAX_BITWORD_WIDTH>(
+            std::move(irrelevant_rng),
+            stats.num_qubits,
+            +1,
+            MeasureRecord(stats.max_lookback)));
+    return helper.do_loop_with_tortoise_hare_folding(circuit, 1).simplified();
+}
+
+std::string ReferenceSampleTree::str() const {
+    std::stringstream ss;
+    ss << *this;
+    return ss.str();
+}
+
+bool ReferenceSampleTree::operator==(const ReferenceSampleTree &other) const {
+    return repetitions == other.repetitions &&
+           prefix_bits == other.prefix_bits &&
+           suffix_children == other.suffix_children;
+}
+bool ReferenceSampleTree::operator!=(const ReferenceSampleTree &other) const {
+    return !(*this == other);
+}
+
+std::ostream &stim::operator<<(std::ostream &out, const ReferenceSampleTree &v) {
+    out << v.repetitions << "*";
+    out << "(";
+    out << "'";
+    for (auto b : v.prefix_bits) {
+        out << "01"[b];
+    }
+    out << "'";
+    for (const auto &child : v.suffix_children) {
+        out << "+";
+        out << child;
+    }
+    out << ")";
+    return out;
+}

src/stim/util_top/reference_sample_tree.h

@@ -0,0 +1,86 @@
+#ifndef _STIM_UTIL_TOP_REFERENCE_SAMPLE_TREE_H
+#define _STIM_UTIL_TOP_REFERENCE_SAMPLE_TREE_H
+
+#include "stim/simulators/tableau_simulator.h"
+
+namespace stim {
+
+/// A compressed tree representation of a reference sample.
+struct ReferenceSampleTree {
+    /// Raw bits to output before bits from the children.
+    std::vector<bool> prefix_bits;
+    /// Compressed representations of additional bits to output after the prefix.
+    std::vector<ReferenceSampleTree> suffix_children;
+    /// The number of times to repeatedly output the prefix+suffix bits.
+    size_t repetitions = 0;
+
+    /// Initializes a reference sample tree containing a reference sample for the given circuit.
+    static ReferenceSampleTree from_circuit_reference_sample(const Circuit &circuit);
+
+    /// Returns a tree with the same compressed contents, but a simpler tree structure.
+    ReferenceSampleTree simplified() const;
+
+    /// Checks if two trees are exactly the same, including structure (not just uncompressed contents).
+    bool operator==(const ReferenceSampleTree &other) const;
+    /// Checks if two trees are not exactly the same, including structure (not just uncompressed contents).
+    bool operator!=(const ReferenceSampleTree &other) const;
+    /// Returns a simple description of the tree's structure, like "5*('101'+6*('11'))".
+    std::string str() const;
+
+    /// Determines whether the tree contains any bits at all.
+    bool empty() const;
+    /// Computes the total size of the uncompressed bits represented by the tree.
+    size_t size() const;
+
+    /// Writes the contents of the tree into the given output vector.
+    void decompress_into(std::vector<bool> &output) const;
+
+    /// Folds redundant children into the repetition count, if they repeat this many times.
+    ///
+    /// For example, if the tree's children are [A, B, C, A, B, C] and the tree has no
+    /// prefix, then `try_factorize(2)` will reduce the children to [A, B, C] and double
+    /// the repetition count.
+    void try_factorize(size_t period_factor);
+
+   private:
+    /// Helper method for `simplified`.
+    void flatten_and_simplify_into(std::vector<ReferenceSampleTree> &out) const;
+};
+std::ostream &operator<<(std::ostream &out, const ReferenceSampleTree &v);
+
+/// Helper class for computing compressed reference samples.
+template <size_t W>
+struct CompressedReferenceSampleHelper {
+    TableauSimulator<W> sim;
+
+    CompressedReferenceSampleHelper(TableauSimulator<MAX_BITWORD_WIDTH> sim) : sim(sim) {
+    }
+
+    /// Processes a loop with no top-level folding.
+    ///
+    /// Loops containing within the body of this loop (or circuit body) may
+    /// still be compressed. Only the top-level loop is not folded.
+    ReferenceSampleTree do_loop_with_no_folding(
+        const Circuit &loop,
+        uint64_t reps);
+
+    /// Runs tortoise-and-hare analysis of the loop while simulating its
+    /// reference sample, in order to attempt to return a compressed
+    /// representation.
+    ReferenceSampleTree do_loop_with_tortoise_hare_folding(
+        const Circuit &loop,
+        uint64_t reps);
+
+    bool in_same_recent_state_as(
+        const CompressedReferenceSampleHelper<W> &other,
+        uint64_t max_record_lookback,
+        bool allow_false_negative) const;
+};
+
+uint64_t max_feedback_lookback_in_loop(const Circuit &loop);
+
+}  // namespace stim
+
+#include "stim/util_top/reference_sample_tree.inl"
+
+#endif