From 51f76c69c4d368f4da743fd67e322df28cc41ff2 Mon Sep 17 00:00:00 2001
From: Pierre Curto <pierre.curto@gmail.com>
Date: Wed, 6 Sep 2023 10:36:21 +0200
Subject: [PATCH] lib/std/collections: add Bitset

Signed-off-by: Pierre Curto <pierre.curto@gmail.com>
---
 lib/std/collections/bitset.c3          | 145 +++++++++++++++++++++++++
 test/unit/stdlib/collections/bitset.c3 |  92 ++++++++++++++++
 2 files changed, 237 insertions(+)
 create mode 100644 lib/std/collections/bitset.c3
 create mode 100644 test/unit/stdlib/collections/bitset.c3

diff --git a/lib/std/collections/bitset.c3 b/lib/std/collections/bitset.c3
new file mode 100644
index 000000000..e46c536e1
--- /dev/null
+++ b/lib/std/collections/bitset.c3
@@ -0,0 +1,145 @@
+/**
+ * @require SIZE > 0
+ **/
+module std::collections::bitset(<SIZE>);
+
+def Type = uint;
+
+const BITS = Type.sizeof * 8;
+const SZ = (SIZE + BITS - 1) / BITS;
+
+struct BitSet
+{
+	Type[SZ] data;
+}
+
+fn usz BitSet.cardinality(&self)
+{
+	usz n;
+	foreach (x : self.data)
+	{
+		n += x.popcount();
+	}
+	return n;
+}
+
+/**
+ * @require i < SIZE
+ **/
+fn void BitSet.set(&self, usz i)
+{
+	usz q = i / BITS;
+	usz r = i % BITS;
+	self.data[q] |= 1 << r;
+}
+
+/**
+ * @require i < SIZE
+ **/
+fn void BitSet.unset(&self, usz i)
+{
+	usz q = i / BITS;
+	usz r = i % BITS;
+	self.data[q] &= ~(1 << r);
+}
+
+/**
+ * @require i < SIZE
+ **/
+fn bool BitSet.get(&self, usz i) @operator([]) @inline
+{
+	usz q = i / BITS;
+	usz r = i % BITS;
+	return self.data[q] & (1 << r) != 0;
+}
+
+fn usz BitSet.len(&self) @operator(len) @inline
+{
+	return SZ * BITS;
+}
+
+/**
+ * @require i < SIZE
+ **/
+fn void BitSet.set_bool(&self, usz i, bool value) @operator([]=) @inline
+{
+	if (value) return self.set(i);
+	self.unset(i);
+}
+
+/**
+ * @require Type.kindof == UNSIGNED_INT
+ **/
+module std::collections::growablebitset(<Type>);
+import std::collections::list;
+
+const BITS = Type.sizeof * 8;
+
+def GrowableBitSetList = List(<Type>);
+
+struct GrowableBitSet
+{
+	GrowableBitSetList data;
+}
+
+fn void GrowableBitSet.init(&self, usz initial_capacity = 1, Allocator* using = mem::heap())
+{
+	self.data.init(initial_capacity, using);
+}
+
+fn void GrowableBitSet.tinit(&self)
+{
+	self.init(.using = mem::temp());
+}
+
+fn usz GrowableBitSet.cardinality(&self)
+{
+	usz n;
+	foreach (x : self.data)
+	{
+		n += x.popcount();
+	}
+	return n;
+}
+
+fn void GrowableBitSet.set(&self, usz i)
+{
+	usz q = i / BITS;
+	usz r = i % BITS;
+	if (q >= self.data.len())
+	{
+		usz n = q + 1;
+		self.data.reserve(n);
+		self.data.size = n;
+	}
+	self.data.set(q, self.data[q] | (1 << r));
+}
+
+fn void GrowableBitSet.unset(&self, usz i)
+{
+	usz q = i / BITS;
+	usz r = i % BITS;
+	if (q >= self.data.len()) return;
+	self.data.set(q, self.data[q] &~ (1 << r));
+}
+
+fn bool GrowableBitSet.get(&self, usz i) @operator([]) @inline
+{
+	usz q = i / BITS;
+	usz r = i % BITS;
+	if (q >= self.data.len()) return false;
+	return self.data[q] & (1 << r) != 0;
+}
+
+fn usz GrowableBitSet.len(&self) @operator(len)
+{
+	usz n = self.data.len() * BITS;
+	if (n > 0) n -= (usz)self.data[n - 1].ctz() - 1;
+	return n;
+}
+
+fn void GrowableBitSet.set_bool(&self, usz i, bool value) @operator([]=) @inline
+{
+	if (value) return self.set(i);
+	self.unset(i);
+}
\ No newline at end of file
diff --git a/test/unit/stdlib/collections/bitset.c3 b/test/unit/stdlib/collections/bitset.c3
new file mode 100644
index 000000000..a6db770e2
--- /dev/null
+++ b/test/unit/stdlib/collections/bitset.c3
@@ -0,0 +1,92 @@
+module bitset_test @test;
+import std::collections::bitset;
+import std::collections::growablebitset;
+import std::collections::list;
+import std::io;
+
+def List = List(<usz>);
+
+def BitSet = BitSet(<2048>);
+
+fn void! set_get()
+{
+	BitSet bs;
+	assert(bs.cardinality() == 0);
+
+	assert(!bs.get(0));
+	bs.set(0);
+	assert(bs.get(0));
+	assert(bs.cardinality() == 1);
+
+	assert(!bs.get(2000));
+	bs[2000] = true;
+	assert(bs.get(2000));
+	assert(bs.cardinality() == 2);
+
+	List found;
+	found.tinit();
+	foreach (i, x : bs)
+	{
+		switch (i)
+		{
+			case 0:
+			case 2000:
+				assert(x);
+				found.push(i);
+			default:
+				assert(!x);
+		}
+	}
+	assert(found.array_view() == usz[]{0, 2000});
+
+	bs.unset(0);
+	assert(!bs.get(0));
+	bs[2000] = false;
+	assert(!bs.get(2000));
+	assert(bs.cardinality() == 0);
+}
+
+def GrowableBitSet = GrowableBitSet(<char>);
+
+fn void! growable_set_get()
+{
+	GrowableBitSet bs;
+	bs.tinit();
+	assert(bs.cardinality() == 0);
+
+	assert(!bs.get(0));
+	bs.set(0);
+	assert(bs.get(0));
+	assert(bs.cardinality() == 1);
+	assert(bs.len() == 1);
+
+	assert(!bs.get(2000));
+	bs[2000] = true;
+	assert(bs.get(2000));
+	assert(bs.cardinality() == 2);
+
+	assert(bs.data.len() == 251);
+	assert(bs.len() == 2001);
+
+	List found;
+	found.tinit();
+	foreach (i, x : bs)
+	{
+		switch (i)
+		{
+			case 0:
+			case 2000:
+				assert(x);
+				found.push(i);
+			default:
+				assert(!x);
+		}
+	}
+	assert(found.array_view() == usz[]{0, 2000});
+
+	bs.unset(0);
+	assert(!bs.get(0));
+	bs[2000] = false;
+	assert(!bs.get(2000));
+	assert(bs.cardinality() == 0);
+}
\ No newline at end of file