diff --git a/nomad/structs/bitmap.go b/nomad/structs/bitmap.go new file mode 100644 index 000000000..22a4aedb8 --- /dev/null +++ b/nomad/structs/bitmap.go @@ -0,0 +1,39 @@ +package structs + +import "fmt" + +// Bitmap is a simple uncompressed bitmap +type Bitmap []byte + +// NewBitmap returns a bitmap with up to size indexes +func NewBitmap(size int) (Bitmap, error) { + if size <= 0 { + return nil, fmt.Errorf("bitmap must be positive size") + } + if size&7 != 0 { + return nil, fmt.Errorf("bitmap must be byte aligned") + } + b := make([]byte, size>>3) + return Bitmap(b), nil +} + +// Set is used to set the given index of the bitmap +func (b Bitmap) Set(idx uint) { + bucket := idx >> 3 + mask := byte(1 << (idx & 7)) + b[bucket] |= mask +} + +// Check is used to check the given index of the bitmap +func (b Bitmap) Check(idx uint) bool { + bucket := idx >> 3 + mask := byte(1 << (idx & 7)) + return (b[bucket] & mask) != 0 +} + +// Clear is used to efficiently clear the bitmap +func (b Bitmap) Clear() { + for i := range b { + b[i] = 0 + } +} diff --git a/nomad/structs/bitmap_test.go b/nomad/structs/bitmap_test.go new file mode 100644 index 000000000..076d6c6d4 --- /dev/null +++ b/nomad/structs/bitmap_test.go @@ -0,0 +1,58 @@ +package structs + +import "testing" + +func TestBitmap(t *testing.T) { + // Check invalid sizes + _, err := NewBitmap(0) + if err == nil { + t.Fatalf("bad") + } + _, err = NewBitmap(7) + if err == nil { + t.Fatalf("bad") + } + + // Create a normal bitmap + b, err := NewBitmap(256) + if err != nil { + t.Fatalf("err: %v", err) + } + + // Set a few bits + b.Set(0) + b.Set(255) + + // Verify the bytes + if b[0] == 0 { + t.Fatalf("bad") + } + if !b.Check(0) { + t.Fatalf("bad") + } + + // Verify the bytes + if b[len(b)-1] == 0 { + t.Fatalf("bad") + } + if !b.Check(255) { + t.Fatalf("bad") + } + + // All other bits should be unset + for i := 1; i < 255; i++ { + if b.Check(uint(i)) { + t.Fatalf("bad") + } + } + + // Clear + b.Clear() + + // All bits should be unset + for i := 0; i < 256; i++ { + if b.Check(uint(i)) { + t.Fatalf("bad") + } + } +}