smallbitset.go

package smallbitset

import (
	"fmt"
)

type BitSetI interface {
	Set(uint64)
	Clear(uint64)
	Test(uint64) bool
	IsSuperSet(BitSetI) bool
	Clone() BitSetI
	Count() int
	String() string
}

// SmallBitSet is backed by a uint64 - only use this for sets with a maximum cardinality of 64
type SmallBitSet uint64

func (b *SmallBitSet) Set(index uint) {
	if index > 63 {
		panic("cannot use SmallBitset with elements greater than 63")
	}
	*b |= 1 << index
}

func (b *SmallBitSet) Test(index uint) bool {
	return (*b)&(1<<index) != 0
}

func (b *SmallBitSet) IsSuperset(other SmallBitSet) bool {
	return *b&other == other
}

func (b *SmallBitSet) Clear(index uint) {
	*b &^= 1 << index
}

func (b *SmallBitSet) Equals(other SmallBitSet) bool {
	return *b == other
}

func (b *SmallBitSet) Clone() *SmallBitSet {
	return b
}

func (b *SmallBitSet) AsSlice() []uint {
	elements := make([]uint, 0, 64)
	for i := uint(0); i < 64; i++ {
		if b.Test(i) {
			elements = append(elements, i)
		}
	}
	return elements
}

func (b *SmallBitSet) String() string {
	return fmt.Sprintf("%v", b.AsSlice())
}

func FromSlice(elements []uint) SmallBitSet {
	b := SmallBitSet(0)
	for _, element := range elements {
		b.Set(element)
	}
	return b
}

smallbitset_test.go

package smallbitset

import (
	"testing"
)

func TestSetAndTest(t *testing.T) {
	var b SmallBitSet
	if b.Test(3) {
		t.Errorf("Expected bit 3 to be unset initially")
	}
	b.Set(3)
	if !b.Test(3) {
		t.Errorf("Expected bit 3 to be set after adding")
	}

	if b.Test(63) {
		t.Errorf("Expected bit 64 to be unset")
	}

	b.Set(63)
	if !b.Test(63) {
		t.Errorf("Expected bit 63 to be set after adding")
	}

	b.Set(63)
	if !b.Test(63) {
		t.Errorf("Expected bit 63 to be set after adding")
	}

	if !b.Test(3) {
		t.Errorf("Expected bit 3 to be set after adding")
	}
}

func TestSetOutOfBounds(t *testing.T) {
	defer func() {
		if r := recover(); r == nil {
			t.Errorf("Expected panic when adding element greater than 64")
		}
	}()

	b := SmallBitSet(0)
	b.Set(64)
}

func TestClear(t *testing.T) {
	var b SmallBitSet
	b.Set(2)
	if !b.Test(2) {
		t.Errorf("Expected bit 2 to be set after adding")
	}
	b.Clear(2)
	if b.Test(2) {
		t.Errorf("Expected bit 2 to be unset after removing")
	}
}

func TestIsSuperset(t *testing.T) {
	var b1, b2 SmallBitSet
	b1.Set(1)
	b1.Set(3)
	b2.Set(1)
	b2.Set(2)
	b2.Set(3)

	if b1.IsSuperset(b2) {
		t.Errorf("Did not expect b1 to be a superset of b2")
	}
	if !b2.IsSuperset(b1) {
		t.Errorf("Expected b2 to be a superset of b1")
	}
}

func TestEquals(t *testing.T) {
	var b1, b2 SmallBitSet
	b1.Set(1)
	b1.Set(4)
	b2.Set(4)
	b2.Set(1)
	if !b1.Equals(b2) {
		t.Errorf("Expected b1 and b2 to be equal")
	}
	b2.Set(5)
	if b1.Equals(b2) {
		t.Errorf("Expected b1 and b2 not to be equal after adding an extra element to b2")
	}
}

func TestString(t *testing.T) {
	var b SmallBitSet
	b.Set(3)
	b.Set(1)
	b.Set(5)
	expected := "[1 3 5]"
	result := b.String()
	if result != expected {
		t.Errorf("Expected string representation %q, got %q", expected, result)
	}
}

func TestBitsetFromSlice(t *testing.T) {
	elements := []uint{2, 4, 6}
	b := FromSlice(elements)
	for _, elem := range elements {
		if !b.Test(elem) {
			t.Errorf("Expected bitset to have element %d", elem)
		}
	}
	// Ensure that an element not in the slice is not set.
	if b.Test(1) {
		t.Errorf("Expected bitset not to have element 1")
	}
}