Implement forEachSetBit

Author: s1ck

core/src/main/java/org/neo4j/gds/core/utils/paged/HugeAtomicPagedBitSet.java

@@ -24,6 +24,7 @@
 
 import java.util.concurrent.atomic.AtomicLongArray;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.LongConsumer;
 
 public final class HugeAtomicPagedBitSet {
 
@@ -149,6 +150,28 @@ public long cardinality() {
         return setBitCount;
     }
 
+    public void forEachSetBit(LongConsumer consumer) {
+        final Pages pages = this.pages.get();
+        final long pageCount = pages.length();
+        final long pageSize = this.pageSize;
+
+        long base = 0;
+
+        for (int pageIndex = 0; pageIndex < pageCount; pageIndex++) {
+            var page = pages.getPage(pageIndex);
+            for (int wordIndex = 0; wordIndex < pageSize; wordIndex++) {
+                long word = page.get(wordIndex);
+
+                while (word != 0) {
+                    long next = Long.numberOfTrailingZeros(word);
+                    consumer.accept(Long.SIZE * (base + wordIndex) + next);
+                    word = word ^ Long.lowestOneBit(word);
+                }
+            }
+            base += pageSize;
+        }
+    }
+
     public long capacity() {
         return pages.get().length() * (1L << pageShift);
     }

core/src/test/java/org/neo4j/gds/core/utils/paged/HugeAtomicPagedBitSetTest.java

@@ -26,8 +26,12 @@
 import org.neo4j.gds.core.concurrency.RunWithConcurrency;
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -120,6 +124,25 @@ void testCardinality(HugeAtomicPagedBitSet atomicBitSet) {
         assertThat(atomicBitSet.cardinality()).isEqualTo(3);
     }
 
+    @ParameterizedTest
+    @MethodSource("bitsets")
+    void testForEachSetBit(HugeAtomicPagedBitSet atomicBitSet) {
+        var rng = ThreadLocalRandom.current();
+        long bound = 42 * (1L << PAGE_SHIFT_BITS) + 42;
+
+        var expected = IntStream
+            .range(0, 10_000)
+            .mapToLong(__ -> rng.nextLong(0, bound))
+            .boxed()
+            .peek(atomicBitSet::set)
+            .collect(Collectors.toSet());
+
+        var actual = new HashSet<Long>();
+        atomicBitSet.forEachSetBit(actual::add);
+
+        assertThat(actual).isEqualTo(expected);
+    }
+
     @Test
     void writingAndGrowingShouldBeThreadSafe() {
         int concurrency = 8;