ReferenceSampleTree Random Access (#983)

* Adding `operator[]` to `ReferenceSampleTree`.
* Allows quick access to any reference sample bit within the tree at a
given absolute index.
* Updating `src/stim/util_top/reference_sample_tree.test.cc` to check
bit positions, as applicable.

Author: qec-pconner

src/stim/util_top/reference_sample_tree.cc

@@ -1,5 +1,10 @@
 #include "stim/util_top/reference_sample_tree.h"
 
+#if defined(_WIN32)
+#include <intrin.h>
+#pragma intrinsic(_umul128)
+#endif
+
 using namespace stim;
 
 bool ReferenceSampleTree::empty() const {
@@ -149,12 +154,23 @@ ReferenceSampleTree ReferenceSampleTree::simplified() const {
     return result;
 }
 
-size_t ReferenceSampleTree::size() const {
-    size_t result = prefix_bits.size();
+uint64_t ReferenceSampleTree::size() const {
+    uint64_t result = prefix_bits.size();
     for (const auto &child : suffix_children) {
         result += child.size();
     }
-    return result * repetitions;
+#if defined(__GNUC__) || defined(__clang__)
+    bool overflow = __builtin_mul_overflow(result, repetitions, &result);
+    assert(!overflow);
+#elif defined(_WIN64)
+    uint64_t overflow;
+    result = _umul128(result, repetitions, &overflow);
+    assert(overflow == 0);
+#else
+    assert((repetitions == 0) || (result <= (UINT64_MAX / repetitions)));
+    result *= repetitions;
+#endif
+    return result;
 }
 
 void ReferenceSampleTree::decompress_into(std::vector<bool> &output) const {
@@ -189,6 +205,57 @@ bool ReferenceSampleTree::operator!=(const ReferenceSampleTree &other) const {
     return !(*this == other);
 }
 
+bool ReferenceSampleTree::operator[](uint64_t index) const {
+    uint64_t current_absolute_index = 0;
+    bool result;
+    bool value_found = try_get_bit_value(index, current_absolute_index, result);
+    assert(value_found);
+    return result;
+}
+
+bool ReferenceSampleTree::try_get_bit_value(
+    uint64_t desired_absolute_index, uint64_t &current_absolute_index, bool &bit_value) const {
+    // Run through once to allow shallow accesses (without unnecessary full iteration of the tree).
+    const uint64_t current_relative_starting_index = current_absolute_index;
+    {
+        const uint64_t relative_index = desired_absolute_index - current_absolute_index;
+        if (relative_index < prefix_bits.size()) {
+            bit_value = prefix_bits[relative_index];
+            return true;
+        }
+        current_absolute_index += prefix_bits.size();
+        for (const ReferenceSampleTree &child : suffix_children) {
+            if (child.try_get_bit_value(desired_absolute_index, current_absolute_index, bit_value)) {
+                return true;
+            }
+        }
+    }
+    // After the first full iteration, extrapolate the repetition size and skip to the proper iteration.
+    const uint64_t single_iteration_size = current_absolute_index - current_relative_starting_index;
+    const uint64_t skip_to_iteration_count =
+        (desired_absolute_index - current_relative_starting_index) / single_iteration_size;
+    if (skip_to_iteration_count < repetitions) {
+        // If the desired index is in this part of the tree, skip forward to the appropriate iteration.
+        current_absolute_index += single_iteration_size * (skip_to_iteration_count - 1);
+        // Do the final iteration to find the value.
+        const uint64_t relative_index = desired_absolute_index - current_absolute_index;
+        if (relative_index < prefix_bits.size()) {
+            bit_value = prefix_bits[relative_index];
+            return true;
+        }
+        current_absolute_index += prefix_bits.size();
+        for (const ReferenceSampleTree &child : suffix_children) {
+            if (child.try_get_bit_value(desired_absolute_index, current_absolute_index, bit_value)) {
+                return true;
+            }
+        }
+    } else {
+        // Advance past this node for parent to continue.
+        current_absolute_index += single_iteration_size * (repetitions - 1);
+    }
+    return false;
+}
+
 std::ostream &stim::operator<<(std::ostream &out, const ReferenceSampleTree &v) {
     out << v.repetitions << "*";
     out << "(";

src/stim/util_top/reference_sample_tree.h

@@ -24,13 +24,15 @@ struct ReferenceSampleTree {
     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 the bit value for a given absolute index.
+    bool operator[](uint64_t index) 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;
+    uint64_t size() const;
 
     /// Writes the contents of the tree into the given output vector.
     void decompress_into(std::vector<bool> &output) const;
@@ -45,6 +47,8 @@ struct ReferenceSampleTree {
    private:
     /// Helper method for `simplified`.
     void flatten_and_simplify_into(std::vector<ReferenceSampleTree> &out) const;
+    /// Helper method for `operator[]`.
+    bool try_get_bit_value(uint64_t desired_absolute_index, uint64_t & current_absolute_index, bool & bit_value) const;
 };
 std::ostream &operator<<(std::ostream &out, const ReferenceSampleTree &v);
 

src/stim/util_top/reference_sample_tree.test.cc

@@ -15,6 +15,9 @@ void expect_tree_matches_normal_reference_sample_of(const ReferenceSampleTree &t
     }
     auto expected = TableauSimulator<MAX_BITWORD_WIDTH>::reference_sample_circuit(circuit);
     EXPECT_EQ(actual, expected);
+    for (size_t index = 0; index < decompressed.size(); ++index) {
+        ASSERT_EQ(tree[index], decompressed[index]) << "index: " << index;
+    }
 }
 
 TEST(ReferenceSampleTree, equality) {
@@ -102,20 +105,24 @@ TEST(ReferenceSampleTree, simplified) {
 
 TEST(ReferenceSampleTree, decompress_into) {
     std::vector<bool> result;
-    ReferenceSampleTree{
+    ReferenceSampleTree tree_under_test{
         .prefix_bits = {1, 1, 0, 1},
         .suffix_children = {ReferenceSampleTree{
             .prefix_bits = {1},
             .suffix_children = {},
             .repetitions = 5,
         }},
         .repetitions = 2,
+    };
+    tree_under_test.decompress_into(result);
+    std::vector<bool> expected = std::vector<bool>{1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1};
+    ASSERT_EQ(result, expected);
+    for (size_t index = 0; index < expected.size(); ++index) {
+        ASSERT_EQ(tree_under_test[index], expected[index]) << "index: " << index;
     }
-        .decompress_into(result);
-    ASSERT_EQ(result, (std::vector<bool>{1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1}));
 
     result.clear();
-    ReferenceSampleTree{
+    tree_under_test = ReferenceSampleTree{
         .prefix_bits = {1, 1, 0, 1},
         .suffix_children =
             {
@@ -131,10 +138,14 @@ TEST(ReferenceSampleTree, decompress_into) {
                 },
             },
         .repetitions = 1,
+    };
+    tree_under_test.decompress_into(result);
+    expected =
+        std::vector<bool>{1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0};
+    ASSERT_EQ(result, expected);
+    for (size_t index = 0; index < expected.size(); ++index) {
+        ASSERT_EQ(tree_under_test[index], expected[index]) << "index: " << index;
     }
-        .decompress_into(result);
-    ASSERT_EQ(result, (std::vector<bool>{1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0,
-                                         1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0}));
 }
 
 TEST(ReferenceSampleTree, simple_circuit) {
@@ -292,3 +303,54 @@ TEST(ReferenceSampleTree, surface_code_with_pauli_vs_normal_reference_sample) {
     ASSERT_EQ(ref.size(), circuit.count_measurements());
     expect_tree_matches_normal_reference_sample_of(ref, circuit);
 }
+
+TEST(ReferenceSampleTree, random_access_large_tree) {
+    ReferenceSampleTree tree_under_test{
+        .prefix_bits = {1, 1, 0, 1},
+        .suffix_children =
+            {ReferenceSampleTree{
+                 .prefix_bits = {1, 0, 1},
+                 .suffix_children = {},
+                 .repetitions = 60'000'000,
+             },
+             ReferenceSampleTree{
+                 .prefix_bits = {0, 0, 0, 0, 0, 0, 1},
+                 .suffix_children = {ReferenceSampleTree{
+                     .prefix_bits = {1, 1, 1, 0, 0, 1},
+                     .suffix_children = {},
+                     .repetitions = 42,
+                 }},
+                 .repetitions = 2'000'000'000,
+             },
+             ReferenceSampleTree{
+                 .prefix_bits = {0, 0, 0, 0, 1},
+                 .suffix_children = {},
+                 .repetitions = 999'000'000,
+             }},
+        .repetitions = 1'234'000,
+    };
+
+    std::vector<bool> expected_beginning = std::vector<bool>{1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1};
+    for (size_t index = 0; index < expected_beginning.size(); ++index) {
+        ASSERT_EQ(tree_under_test[index], expected_beginning[index]) << "index: " << index;
+    }
+
+    uint64_t whole_tree_size = tree_under_test.size();
+    ASSERT_EQ(
+        whole_tree_size,
+        1'234'000ULL * (4 + (60'000'000ULL * 3) + (2'000'000'000ULL * (7 + (42 * 6))) + (999'000'000ULL * 5)));
+
+    std::vector<bool> expected_ending =
+        std::vector<bool>{0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1};
+    for (uint64_t index = 0; index < expected_ending.size(); ++index) {
+        ASSERT_EQ(tree_under_test[whole_tree_size - expected_ending.size() + index], expected_ending[index])
+            << "index: " << index;
+    }
+    // Same thing on previous outer iteration.
+    uint64_t one_outer_iter_before_ending = whole_tree_size - (whole_tree_size / 1'234'000ULL);
+    for (uint64_t index = 0; index < expected_ending.size(); ++index) {
+        ASSERT_EQ(
+            tree_under_test[one_outer_iter_before_ending - expected_ending.size() + index], expected_ending[index])
+            << "index: " << index;
+    }
+}