From d22d72b963b9d80cd8ec2762e58f63a5d9b457df Mon Sep 17 00:00:00 2001
From: Karl Gaissmaier <karl.gaissmaier@uni-ulm.de>
Date: Thu, 29 Aug 2024 23:44:32 +0200
Subject: [PATCH] add inlineable bitset.MustNew

---
 bitset.go      | 13 +++++++++++++
 bitset_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/bitset.go b/bitset.go
index cdf91c3..5ad1745 100644
--- a/bitset.go
+++ b/bitset.go
@@ -164,6 +164,19 @@ func New(length uint) (bset *BitSet) {
 	return bset
 }
 
+// MustNew creates a new BitSet with the given length bits.
+// It panics if length exceeds the possible capacity or by a lack of memory.
+func MustNew(length uint) (bset *BitSet) {
+	if length >= Cap() {
+		panic("You are exceeding the capacity")
+	}
+
+	return &BitSet{
+		length,
+		make([]uint64, wordsNeeded(length)), // may panic on lack of memory
+	}
+}
+
 // Cap returns the total possible capacity, or number of bits
 // that can be stored in the BitSet theoretically. Under 32-bit system,
 // it is 4294967295 and under 64-bit system, it is 18446744073709551615.
diff --git a/bitset_test.go b/bitset_test.go
index ca21d8b..d9d012f 100644
--- a/bitset_test.go
+++ b/bitset_test.go
@@ -420,6 +420,54 @@ func TestNullCount(t *testing.T) {
 	}
 }
 
+func TestMustNew(t *testing.T) {
+	testCases := []struct {
+		length uint
+		nwords int
+	}{
+		{
+			length: 0,
+			nwords: 0,
+		},
+		{
+			length: 1,
+			nwords: 1,
+		},
+		{
+			length: 64,
+			nwords: 1,
+		},
+		{
+			length: 65,
+			nwords: 2,
+		},
+		{
+			length: 512,
+			nwords: 8,
+		},
+		{
+			length: 513,
+			nwords: 9,
+		},
+	}
+
+	for _, tc := range testCases {
+		b := MustNew(tc.length)
+		if len(b.set) != tc.nwords {
+			t.Errorf("length = %d, len(b.set) got: %d, want: %d", tc.length, len(b.set), tc.nwords)
+		}
+	}
+}
+
+func TestPanicMustNew(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Error("length too big should have caused a panic")
+		}
+	}()
+	MustNew(Cap())
+}
+
 func TestPanicDifferenceBNil(t *testing.T) {
 	var b *BitSet
 	var compare = New(10)