From a22a134b4f086abc28b461ceb36a2583e4536a42 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Fri, 14 Mar 2025 21:05:58 +0700
Subject: [PATCH] BitSet iterator

---
 gradle.properties                             |  2 +-
 .../kommons/collect/BitSetIterator.kt         | 82 +++++++++++++++++++
 .../kommons/test/BitSetIteratorTest.kt        | 50 +++++++++++
 3 files changed, 133 insertions(+), 1 deletion(-)
 create mode 100644 src/main/kotlin/ru/dbotthepony/kommons/collect/BitSetIterator.kt
 create mode 100644 src/test/kotlin/ru/dbotthepony/kommons/test/BitSetIteratorTest.kt

diff --git a/gradle.properties b/gradle.properties
index 32fcfca..aca9826 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -4,7 +4,7 @@ kotlin.code.style=official
 specifyKotlinAsDependency=false
 
 projectGroup=ru.dbotthepony.kommons
-projectVersion=3.5.2
+projectVersion=3.6.0
 
 guavaDepVersion=33.0.0
 gsonDepVersion=2.8.9
diff --git a/src/main/kotlin/ru/dbotthepony/kommons/collect/BitSetIterator.kt b/src/main/kotlin/ru/dbotthepony/kommons/collect/BitSetIterator.kt
new file mode 100644
index 0000000..60286a5
--- /dev/null
+++ b/src/main/kotlin/ru/dbotthepony/kommons/collect/BitSetIterator.kt
@@ -0,0 +1,82 @@
+package ru.dbotthepony.kommons.collect
+
+import java.util.BitSet
+
+sealed class BitSetIterator(protected val set: BitSet, first: Int, private val last: Int) : IntIterator() {
+	protected abstract fun findNext(from: Int): Int
+
+	private var hasNext = false
+	private var next = -1
+
+	init {
+		if (first < last) {
+			val nextSet = findNext(first)
+
+			if (nextSet < last && nextSet != -1) {
+				hasNext = true
+				next = nextSet
+			}
+		}
+	}
+
+	override fun nextInt(): Int {
+		if (!hasNext)
+			throw NoSuchElementException()
+
+		val nextSet = findNext(next + 1)
+
+		if (nextSet < last && nextSet != -1) {
+			hasNext = true
+			val value = next
+			next = nextSet
+			return value
+		} else {
+			hasNext = false
+			return next
+		}
+	}
+
+	override fun hasNext(): Boolean {
+		return hasNext
+	}
+
+	class Set(set: BitSet, first: Int, last: Int) : BitSetIterator(set, first, last) {
+		override fun findNext(from: Int): Int {
+			return set.nextSetBit(from)
+		}
+	}
+
+	class Clear(set: BitSet, first: Int, last: Int) : BitSetIterator(set, first, last) {
+		override fun findNext(from: Int): Int {
+			return set.nextClearBit(from)
+		}
+	}
+}
+
+/**
+ * Iterates set bits in this [BitSet], from [first] (inclusive) up to [last] (exclusive)
+ */
+fun BitSet.iterateSetBits(first: Int, last: Int): IntIterator {
+	return BitSetIterator.Set(this, first, last)
+}
+
+/**
+ * Iterates clear bits in this [BitSet], from [first] (inclusive) up to [last] (exclusive)
+ */
+fun BitSet.iterateClearBits(first: Int, last: Int): IntIterator {
+	return BitSetIterator.Clear(this, first, last)
+}
+
+/**
+ * Iterates set bits in this [BitSet], from 0 (inclusive) up to [last] (exclusive)
+ */
+fun BitSet.iterateSetBits(last: Int): IntIterator {
+	return BitSetIterator.Set(this, 0, last)
+}
+
+/**
+ * Iterates clear bits in this [BitSet], from 0 (inclusive) up to [last] (exclusive)
+ */
+fun BitSet.iterateClearBits(last: Int): IntIterator {
+	return BitSetIterator.Clear(this, 0, last)
+}
diff --git a/src/test/kotlin/ru/dbotthepony/kommons/test/BitSetIteratorTest.kt b/src/test/kotlin/ru/dbotthepony/kommons/test/BitSetIteratorTest.kt
new file mode 100644
index 0000000..8967faa
--- /dev/null
+++ b/src/test/kotlin/ru/dbotthepony/kommons/test/BitSetIteratorTest.kt
@@ -0,0 +1,50 @@
+package ru.dbotthepony.kommons.test
+
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import ru.dbotthepony.kommons.collect.iterateClearBits
+import ru.dbotthepony.kommons.collect.iterateSetBits
+import java.util.BitSet
+
+object BitSetIteratorTest {
+	@Test
+	@DisplayName("Bit set iterator")
+	fun test() {
+		val bitSet = BitSet()
+		bitSet[1] = true
+		bitSet[4] = true
+		bitSet[6] = true
+		bitSet[7] = true
+		bitSet[8] = true
+		bitSet[10] = true
+		bitSet[14] = true
+
+		val expectedSet = intArrayOf(1, 4, 6, 7, 8, 10, 14)
+		val expectedUnset = intArrayOf(0, 2, 3, 5, 9, 11, 12, 13)
+
+		var i = 0
+
+		for (bit in bitSet.iterateSetBits(15))
+			Assertions.assertEquals(expectedSet[i++], bit)
+
+		i = 0
+
+		for (bit in bitSet.iterateClearBits(15))
+			Assertions.assertEquals(expectedUnset[i++], bit)
+
+		i = 1
+
+		for (bit in bitSet.iterateSetBits(2, 15))
+			Assertions.assertEquals(expectedSet[i++], bit)
+
+		i = 2
+
+		for (bit in bitSet.iterateClearBits(3, 15))
+			Assertions.assertEquals(expectedUnset[i++], bit)
+
+		Assertions.assertFalse(bitSet.iterateSetBits(15, 100).hasNext())
+		Assertions.assertTrue(bitSet.iterateClearBits(1).hasNext())
+		Assertions.assertFalse(bitSet.iterateClearBits(1, 2).hasNext())
+	}
+}