From 40b0e7d9e267a6c5d5feb2d1c955948828793a61 Mon Sep 17 00:00:00 2001
From: Nick Gasson <nick@nickg.me.uk>
Date: Sat, 23 Mar 2024 09:48:14 +0000
Subject: [PATCH] Fix some boundary cases in mask_iter

---
 src/mask.c       | 23 ++++++++++++++---------
 test/test_misc.c | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/src/mask.c b/src/mask.c
index 266c73c4..19b513f2 100644
--- a/src/mask.c
+++ b/src/mask.c
@@ -303,18 +303,23 @@ bool mask_iter(bit_mask_t *m, int *bit)
    if (*bit + 1 < 0 || *bit + 1 >= m->size)
       return false;
    else if (m->size > 64) {
-      const int first = (*bit + 1) / 64;
-      const uint64_t remain = m->ptr[first] & ~mask_for_range(0, *bit % 64);
-      const int fs = __builtin_ffsll(remain);
-      if (fs > 0) {
-         *bit = fs - 1;
-         return true;
+      int word = (*bit + 1) / 64;
+
+      if ((*bit + 1) % 64 > 0) {
+         const uint64_t remain = m->ptr[word] & ~mask_for_range(0, *bit % 64);
+         const int fs = __builtin_ffsll(remain);
+         if (fs > 0) {
+            *bit = word*64 + fs - 1;
+            return true;
+         }
+
+         word++;
       }
 
-      for (int i = first + 1; i < (m->size + 63) / 64; i++) {
-         const int fs = __builtin_ffsll(m->ptr[i]);
+      for (; word < (m->size + 63) / 64; word++) {
+         const int fs = __builtin_ffsll(m->ptr[word]);
          if (fs > 0) {
-            *bit = i*64 + fs - 1;
+            *bit = word*64 + fs - 1;
             return true;
          }
       }
diff --git a/test/test_misc.c b/test/test_misc.c
index 2df1740a..403b12e2 100644
--- a/test/test_misc.c
+++ b/test/test_misc.c
@@ -578,6 +578,40 @@ START_TEST(test_mask_iter)
    fail_if(mask_iter(&m, &bit));
 
    mask_free(&m);
+
+   mask_init(&m, 85);
+
+   mask_set(&m, 64);
+   mask_set(&m, 70);
+
+   bit = -1;
+   fail_unless(mask_iter(&m, &bit));
+   fail_unless(bit == 64);
+   fail_unless(mask_iter(&m, &bit));
+   fail_unless(bit == 70);
+   fail_if(mask_iter(&m, &bit));
+
+   mask_free(&m);
+
+   mask_init(&m, 124);
+
+   mask_set(&m, 56);
+   mask_set(&m, 63);
+   mask_set(&m, 69);
+   mask_set(&m, 123);
+
+   bit = -1;
+   fail_unless(mask_iter(&m, &bit));
+   fail_unless(bit == 56);
+   fail_unless(mask_iter(&m, &bit));
+   fail_unless(bit == 63);
+   fail_unless(mask_iter(&m, &bit));
+   fail_unless(bit == 69);
+   fail_unless(mask_iter(&m, &bit));
+   fail_unless(bit == 123);
+   fail_if(mask_iter(&m, &bit));
+
+   mask_free(&m);
 }
 END_TEST
 
-- 
2.39.2