From f8a93608cb630efa989899313387082cff8967dc Mon Sep 17 00:00:00 2001
From: Fionn Malone <fmalone@google.com>
Date: Tue, 5 Sep 2023 05:40:30 +0000
Subject: [PATCH] Fix bug with simd_bits +=.

---
 src/stim/mem/simd_bits.test.cc       | 16 ++++++++++++++++
 src/stim/mem/simd_bits_range_ref.inl |  9 +++++----
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/stim/mem/simd_bits.test.cc b/src/stim/mem/simd_bits.test.cc
index 62b9eb9c9..dbfccc343 100644
--- 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
index 46fc514bd..e56837de6 100644
--- 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;
 }