Fix bug with simd_bits +=.

quantumlib · Sep 5, 2023 · f8a9360 · f8a9360
1 parent c135a61
commit f8a9360
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/src/stim/mem/simd_bits.test.cc b/src/stim/mem/simd_bits.test.cc
@@ -228,6 +228,22 @@ TEST_EACH_WORD_SIZE_W(simd_bits, add_assignment, {
     m0 += m1;
     ASSERT_EQ(m0[0], 0);
     ASSERT_EQ(m0[64], 1);
+    // Test carrying across multiple (>=2) words.
+    size_t num_bits = 193;
+    simd_bits<W> add(num_bits);
+    simd_bits<W> one(num_bits);
+    for (size_t word = 0; word < add.num_u64_padded() - 1; word++) {
+        for (size_t k = 0; k < 64; k++) {
+            add[word * 64 + k] = 1;
+        }
+    }
+    one[0] = 1;
+    add += one;
+    // These should all overflow and carries should propagate to the last word.
+    for (size_t k = 0; k < num_bits - 1; k++) {
+        ASSERT_EQ(add[k], 0);
+    }
+    ASSERT_EQ(add[num_bits - 1], 1);
 })
 
 TEST_EACH_WORD_SIZE_W(simd_bits, right_shift_assignment, {

diff --git a/src/stim/mem/simd_bits_range_ref.inl b/src/stim/mem/simd_bits_range_ref.inl
@@ -57,11 +57,12 @@ simd_bits_range_ref<W> simd_bits_range_ref<W>::operator=(const simd_bits_range_r
 template <size_t W>
 simd_bits_range_ref<W> simd_bits_range_ref<W>::operator+=(const simd_bits_range_ref<W> other) {
     size_t num_u64 = num_u64_padded();
-    for (size_t w = 0; w < num_u64 - 1; w++) {
-        u64[w] += other.u64[w];
-        u64[w + 1] += (u64[w] < other.u64[w]);
+    uint64_t carry{0};
+    for (size_t w = 0; w < num_u64; w++) {
+        uint64_t val_before = u64[w];
+        u64[w] += other.u64[w] + carry;
+        carry = u64[w] < val_before;
     }
-    u64[num_u64 - 1] += other.u64[num_u64 - 1];
     return *this;
 }