fixed bitset iterator

drexlerd · May 21, 2024 · cada725 · cada725
1 parent 759b74c
commit cada725
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 75 deletions.
diff --git a/include/flatmemory/details/types/bitset.hpp b/include/flatmemory/details/types/bitset.hpp
@@ -507,6 +507,13 @@ class Operator<Bitset<Block>>
                 // Advance position
                 ++m_pos;
                 ++m_bit_index;
+                if (m_bit_index == 64)
+                {
+                    ++m_block_index;
+                    m_bit_index = 0;
+                    m_cur_block = m_blocks[m_block_index];
+                }
+
                 if (m_cur_block)
                 {
                     // If there are set bits in the current value
@@ -529,16 +536,6 @@ class Operator<Bitset<Block>>
             } while (m_pos < m_end_pos);
         }
 
-        size_t find_end_pos() const
-        {
-            assert(m_num_blocks > 0);
-            // Find the last block that differs from the default block
-            int32_t last_relevant_block_index = static_cast<int64_t>(m_num_blocks) - 1;
-            for (; (last_relevant_block_index >= 0) && (m_blocks[last_relevant_block_index] == block_zeroes); --last_relevant_block_index) {}
-            // Point to the last position of the next block (it must not exist)
-            return (last_relevant_block_index + 1) * block_size - 1;
-        }
-
     public:
         using difference_type = int;
         using value_type = size_t;
@@ -551,19 +548,30 @@ class Operator<Bitset<Block>>
             m_blocks(blocks),
             m_num_blocks(num_blocks),
             m_block_index(0),
-            m_bit_index(-1)  // set bit_index to -1 for advance step
-            ,
+            m_bit_index(-1),
             m_cur_block(num_blocks > 0 ? blocks[m_block_index] : block_zeroes),
-            m_end_pos(find_end_pos()),
-            m_pos(begin ? -1 : m_end_pos)  // set to -1 for advance step
+            m_end_pos((m_num_blocks + 1) * block_size - 1),
+            m_pos(begin ? -1 : m_end_pos)
         {
             if (default_bit_value)
             {
                 throw std::runtime_error("Cannot iterate over infinite set.");
             }
             if (begin && m_pos != m_end_pos)
             {
-                next_set_bit();
+                // Operator::block_msb_one
+                if ((num_blocks > 0) && ((m_blocks[0] & 1) == 0))
+                {
+                    // The first bit is not set, so we need to find it
+                    next_set_bit();
+                }
+                else
+                {
+                    // The first bit is set, advance by 1 manually
+                    m_pos = 0;
+                    m_bit_index = 0;
+                    m_cur_block >>= 1;
+                }
             }
         }
 

diff --git a/tests/unit/types/bitset.cpp b/tests/unit/types/bitset.cpp
@@ -262,6 +262,79 @@ TEST(FlatmemoryTests, TypesBitsetGetTest)
     EXPECT_FALSE(const_view.get(64));
 }
 
+/**
+ * Iterators
+ */
+
+TEST(FlatmemoryTests, TypesBitsetIteratorTest)
+{
+    size_t num_bits = 200;
+    auto builder = Builder<Bitset<uint64_t>>(num_bits, false);
+    builder.set(0);
+    builder.set(2);
+    builder.set(4);
+    builder.set(99);
+
+    auto it = builder.begin();
+    auto end = builder.end();
+    EXPECT_EQ(*it, 0);
+    ++it;
+    EXPECT_EQ(*it, 2);
+    ++it;
+    EXPECT_EQ(*it, 4);
+    ++it;
+    EXPECT_EQ(*it, 99);
+    ++it;
+    EXPECT_EQ(it, end);
+
+    builder.finish();
+
+    auto view = View<Bitset<uint64_t>>(builder.buffer().data());
+    it = view.begin();
+    end = view.end();
+    EXPECT_EQ(*it, 0);
+    ++it;
+    EXPECT_EQ(*it, 2);
+    ++it;
+    EXPECT_EQ(*it, 4);
+    ++it;
+    EXPECT_EQ(*it, 99);
+    ++it;
+    EXPECT_EQ(it, end);
+
+    auto const_view = ConstView<Bitset<uint64_t>>(builder.buffer().data());
+    it = const_view.begin();
+    end = const_view.end();
+    EXPECT_EQ(*it, 0);
+    ++it;
+    EXPECT_EQ(*it, 2);
+    ++it;
+    EXPECT_EQ(*it, 4);
+    ++it;
+    EXPECT_EQ(*it, 99);
+    ++it;
+    EXPECT_EQ(it, end);
+}
+
+TEST(FlatmemoryTests, TypesBitsetIterator2Test)
+{
+    size_t num_bits = 0;
+    auto builder = Builder<Bitset<uint64_t>>(num_bits, false);
+
+    auto it = builder.begin();
+    auto end = builder.end();
+    EXPECT_EQ(it, end);
+}
+
+TEST(FlatmemoryTests, TypesBitsetIterator3Test)
+{
+    // Set the last bit within the size of the bitset.
+    auto builder = Builder<Bitset<uint64_t>>(63, false);
+    builder.set(63);
+    builder.set(127);
+    EXPECT_EQ(builder.count(), 2);
+}
+
 /**
  * TODO: rework all the remaining parts
  */
@@ -465,64 +538,4 @@ TEST(FlatmemoryTests, TypesBitsetEqualityTest)
     EXPECT_NE(const_view2.hash(), const_view3.hash());
 }
 
-TEST(FlatmemoryTests, TypesBitsetIteratorTest)
-{
-    size_t num_bits = 200;
-    auto builder = Builder<Bitset<uint64_t>>(num_bits, false);
-    builder.set(0);
-    builder.set(2);
-    builder.set(4);
-    builder.set(99);
-
-    auto it = builder.begin();
-    auto end = builder.end();
-    EXPECT_EQ(*it, 0);
-    ++it;
-    EXPECT_EQ(*it, 2);
-    ++it;
-    EXPECT_EQ(*it, 4);
-    ++it;
-    EXPECT_EQ(*it, 99);
-    ++it;
-    EXPECT_EQ(it, end);
-
-    builder.finish();
-
-    auto view = View<Bitset<uint64_t>>(builder.buffer().data());
-    it = view.begin();
-    end = view.end();
-    EXPECT_EQ(*it, 0);
-    ++it;
-    EXPECT_EQ(*it, 2);
-    ++it;
-    EXPECT_EQ(*it, 4);
-    ++it;
-    EXPECT_EQ(*it, 99);
-    ++it;
-    EXPECT_EQ(it, end);
-
-    auto const_view = ConstView<Bitset<uint64_t>>(builder.buffer().data());
-    it = const_view.begin();
-    end = const_view.end();
-    EXPECT_EQ(*it, 0);
-    ++it;
-    EXPECT_EQ(*it, 2);
-    ++it;
-    EXPECT_EQ(*it, 4);
-    ++it;
-    EXPECT_EQ(*it, 99);
-    ++it;
-    EXPECT_EQ(it, end);
-}
-
-TEST(FlatmemoryTests, TypesBitsetIterator2Test)
-{
-    size_t num_bits = 0;
-    auto builder = Builder<Bitset<uint64_t>>(num_bits, false);
-
-    auto it = builder.begin();
-    auto end = builder.end();
-    EXPECT_EQ(it, end);
-}
-
 }