From 26bb560eb75624fefa2a753f9deb420945cc237b Mon Sep 17 00:00:00 2001
From: Greg Burd <greg@burd.me>
Date: Mon, 8 Apr 2024 22:01:30 -0400
Subject: [PATCH] more rank() tests; found/fixed bug

---
 flake.nix       |  1 +
 src/sparsemap.c |  9 +++++----
 tests/common.c  | 27 +++++++++++++++++++++++++++
 tests/test.c    | 41 +++++++++++++++++++++++++++++++++++++----
 4 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/flake.nix b/flake.nix
index cdd6e23..e270e61 100644
--- a/flake.nix
+++ b/flake.nix
@@ -44,6 +44,7 @@
                 pkg-config
                 python3
                 ripgrep
+                valgrind
               ];
               buildInputs = with pkgs; [
                 libbacktrace
diff --git a/src/sparsemap.c b/src/sparsemap.c
index c5c6968..f02a81f 100644
--- a/src/sparsemap.c
+++ b/src/sparsemap.c
@@ -499,10 +499,11 @@ __sm_chunk_map_rank(__sm_chunk_t *map, size_t first, size_t last, size_t *after)
               *after = 0;
             }
           }
-          for (size_t k = ks; k < last && k < sizeof(sm_bitvec_t); k++) {
-            if (w & ((sm_bitvec_t)1 << k)) {
-              ret++;
-            }
+          uint64_t mask = ((uint64_t)1 << (last + 1)) - 1 - (((uint64_t)1 << ks) - 1);
+          uint64_t masked = w & mask;
+          while (masked) {
+            ret += masked & 1;
+            masked >>= 1;
           }
           return (ret);
         }
diff --git a/tests/common.c b/tests/common.c
index 57a7f39..01015c3 100644
--- a/tests/common.c
+++ b/tests/common.c
@@ -264,3 +264,30 @@ bitmap_from_uint64(sparsemap_t *map, uint64_t number) {
         sparsemap_set(map, i, bit);
     }
 }
+
+uint32_t
+rank_uint64(uint64_t number, int n, int p)
+{
+    if (p < n || p > 63) {
+        return 0;
+    }
+
+    /* Create a mask for the range between n and p.
+       This works by shifting 1 to the left (p+1) times, subtracting 1 to have
+       a sequence of p 1's, then shifting n times to the left to position it
+       starting at n. Finally, subtracting (1 << n) - 1 removes the bits below
+       n from the mask. */
+    uint64_t mask = ((uint64_t)1 << (p + 1)) - 1 - (((uint64_t)1 << n) - 1);
+
+    /* Apply the mask and count the set bits in the result. */
+    uint64_t maskedNumber = number & mask;
+
+    /* Count the bits set in maskedNumber. */
+    uint32_t count = 0;
+    while (maskedNumber) {
+      count += maskedNumber & 1;
+      maskedNumber >>= 1;
+    }
+
+    return count;
+}
diff --git a/tests/test.c b/tests/test.c
index 28f4e1b..b127881 100644
--- a/tests/test.c
+++ b/tests/test.c
@@ -439,7 +439,7 @@ test_api_select_setup(const MunitParameter params[], void *user_data)
   sparsemap_t *map = (sparsemap_t *)test_api_setup(params, user_data);
 
   sparsemap_init(map, buf, 1024, 0);
-  populate_map(map, 1024, 3 * 1024);
+  bitmap_from_uint64(map, ((uint64_t)0xfeedface << 32) | 0xbadc0ffee);
 
   return (void *)map;
 }
@@ -458,7 +458,11 @@ test_api_select(const MunitParameter params[], void *data)
 
   assert_ptr_not_null(map);
 
-  size_t size = sparsemap_select(map, 1);
+  /* NOTE: select() is 0-based, to get the bit position of the 1st logical bit set
+     call select(map, 0), to get the 18th, select(map, 17), etc. */
+  assert_true(sparsemap_select(map, 0) == 1);
+  assert_true(sparsemap_select(map, 4) == 6);
+  assert_true(sparsemap_select(map, 17) == 26);
 
   return MUNIT_OK;
 }
@@ -470,7 +474,6 @@ test_api_rank_setup(const MunitParameter params[], void *user_data)
   sparsemap_t *map = (sparsemap_t *)test_api_setup(params, user_data);
 
   sparsemap_init(map, buf, 1024, 0);
-  populate_map(map, 1024, 3 * 1024);
 
   return (void *)map;
 }
@@ -484,12 +487,42 @@ test_api_rank_tear_down(void *fixture)
 static MunitResult
 test_api_rank(const MunitParameter params[], void *data)
 {
+  int rank;
   sparsemap_t *map = (sparsemap_t *)data;
   (void)params;
 
   assert_ptr_not_null(map);
 
-  size_t size = sparsemap_rank(map, 0, 1);
+  for (int i = 0; i < 10; i++) {
+    sparsemap_set(map, i, true);
+  }
+  for (int i = 0; i < 10; i++) {
+    assert_true(sparsemap_is_set(map, i));
+  }
+  for (int i = 10; i < 1000; i++) {
+    assert_true(!sparsemap_is_set(map, i));
+  }
+  /* rank() is also 0-based, for consistency (and confusion sake); consider the
+     range as [start, end] of [0, 9] counts the bits set in the first 10
+     positions (starting from the LSB) in the index. */
+  int r1 = sparsemap_rank(map, 0, 9);
+  int r2 = rank_uint64((uint64_t)-1, 0, 9);
+  assert_true(r1 == r2);
+  assert_true(sparsemap_rank(map, 0, 9) == 10);
+
+  for (int i = 0; i < 10; i++) {
+    for (int j = i; j < 10; j++) {
+      r1 = sparsemap_rank(map, i, j);
+      r2 = rank_uint64((uint64_t)-1, i, j);
+      assert_true(r1 == r2);
+    }
+  }
+
+  sparsemap_clear(map);
+  uint64_t bits = ((uint64_t)0xfeedface << 32) | 0xbadc0ffee;
+  bitmap_from_uint64(map, bits);
+  rank = sparsemap_rank(map, 3, 18);
+  assert_true(rank == rank_uint64(bits, 3, 18));
 
   return MUNIT_OK;
 }