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

…ling

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,152 @@
+#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;
+    }
+
+    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.
+    auto &f = flattened;
+    for (size_t k = 0; k < f.size(); k++) {
+        while (k + 1 < f.size() && f[k].repetitions == 1 && f[k].suffix_children.empty() && f[k + 1].repetitions == 1) {
+            f[k].suffix_children = std::move(f[k + 1].suffix_children);
+            f[k].prefix_bits.insert(f[k].prefix_bits.end(), f[k + 1].prefix_bits.begin(), f[k + 1].prefix_bits.end());
+            f.erase(f.begin() + k + 1);
+        }
+    }
+
+    if (repetitions == 1) {
+        // Un-nest all the children.
+        for (auto &e : flattened) {
+            out.push_back(e);
+        }
+    } else if (flattened.size() == 1) {
+        // Merge with single child.
+        flattened[0].repetitions *= repetitions;
+        out.push_back(std::move(flattened[0]));
+    } else if (flattened.empty()) {
+        // Nothing to report.
+    } else if (flattened[0].suffix_children.empty() && flattened[0].repetitions == 1) {
+        // Take payload from first child.
+        auto result = std::move(flattened[0]);
+        flattened.erase(flattened.begin());
+        result.repetitions = repetitions;
+        result.suffix_children = std::move(flattened);
+        out.push_back(std::move(result));
+    } else {
+        out.push_back(ReferenceSampleTree{
+            .prefix_bits={},
+            .suffix_children=std::move(flattened),
+            .repetitions=repetitions,
+        });
+    }
+}
+
+ReferenceSampleTree ReferenceSampleTree::simplified() const {
+    std::vector<ReferenceSampleTree> flat;
+    flatten_and_simplify_into(flat);
+    if (flat.empty()) {
+        return ReferenceSampleTree();
+    } else if (flat.size() == 1) {
+        return flat[0];
+    }
+
+    ReferenceSampleTree result;
+    result.repetitions = 1;
+    if (flat[0].repetitions == 1 && flat[0].suffix_children.empty()) {
+         // Take payload from first child.
+         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,69 @@
+#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 {
+    /// Bits to repeatedly output before outputting bits for 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 and 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;
+
+    /// Determines whether the tree contains any bits at all.
+    bool empty() const;
+
+    bool operator==(const ReferenceSampleTree &other) const;
+    bool operator!=(const ReferenceSampleTree &other) const;
+    std::string str() 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;
+
+   private:
+    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);
+};
+
+}  // namespace stim
+
+#include "stim/util_top/reference_sample_tree.inl"
+
+#endif