add bitmap

provides the implementation of bitmap.

Signed-off-by: Howard Zhang <howard.zhang@arm.com>

3 files changed, 524 insertions, 0 deletions

diff --git a/pkg/bitmap/BUILD b/pkg/bitmap/BUILD
new file mode 100644
index 000000000..1af60002a
--- /dev/null
+++ b/pkg/bitmap/BUILD
@@ -0,0 +1,18 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "bitmap",
+    srcs = ["bitmap.go"],
+    visibility = ["//:sandbox"],
+    deps = ["@org_golang_x_sys//unix:go_default_library"],
+)
+
+go_test(
+    name = "bitmap_test",
+    size = "small",
+    srcs = ["bitmap_test.go"],
+    library = ":bitmap",
+    deps = ["@org_golang_x_sys//unix:go_default_library"],
+)
diff --git a/pkg/bitmap/bitmap.go b/pkg/bitmap/bitmap.go
new file mode 100644
index 000000000..40cf88103
--- /dev/null
+++ b/pkg/bitmap/bitmap.go
@@ -0,0 +1,212 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package bitmap provides the implementation of bitmap.
+package bitmap
+
+import (
+	"math"
+	"math/bits"
+)
+
+// Bitmap implements an efficient bitmap.
+//
+// +stateify savable
+type Bitmap struct {
+	// numOnes is the number of ones in the bitmap.
+	numOnes  uint32
+
+	// bitBlock holds the bits. The type of bitBlock is uint64 which means
+	// each number in bitBlock contains 64 entries.
+	bitBlock []uint64
+}
+
+// BitmapWithSize create a new empty Bitmap.
+func BitmapWithSize(size uint32) Bitmap {
+	b := Bitmap{}
+	bSize := (size + 63) / 64
+	b.bitBlock = make([]uint64, bSize)
+	return b
+}
+
+// IsEmpty verifies whether the Bitmap is empty.
+func (b *Bitmap) IsEmpty() bool {
+	return b.numOnes == 0
+}
+
+// Minimum return the smallest value in the Bitmap.
+func (b *Bitmap) Minimum() uint32 {
+	for i := 0; i < len(b.bitBlock); i++ {
+		if w := b.bitBlock[i]; w != 0 {
+			r := bits.TrailingZeros64(w)
+			return uint32(r + i*64)
+		}
+	}
+	return math.MaxInt32
+}
+
+// Maximum return the largest value in the Bitmap.
+func (b *Bitmap) Maximum() uint32 {
+	for i := len(b.bitBlock) - 1; i >= 0; i-- {
+		if w := b.bitBlock[i]; w != 0 {
+			r := bits.LeadingZeros64(w)
+			return uint32(i*64 + 63 - r)
+		}
+	}
+	return uint32(0)
+}
+
+// Add add i to the Bitmap.
+func (b *Bitmap) Add(i uint32) {
+	blockNum, mask := i/64, uint64(1)<<(i%64)
+	// if blockNum is out of range, extend b.bitBlock
+	if x, y := int(blockNum), len(b.bitBlock); x >= y {
+		b.bitBlock = append(b.bitBlock, make([]uint64, x-y+1)...)
+	}
+	oldBlock := b.bitBlock[blockNum]
+	newBlock := oldBlock | mask
+	if oldBlock != newBlock {
+		b.bitBlock[blockNum] = newBlock
+		b.numOnes++
+	}
+}
+
+// Remove i from the Bitmap.
+func (b *Bitmap) Remove(i uint32) {
+	blockNum, mask := i/64, uint64(1)<<(i%64)
+	oldBlock := b.bitBlock[blockNum]
+	newBlock := oldBlock &^ mask
+	if oldBlock != newBlock {
+		b.bitBlock[blockNum] = newBlock
+		b.numOnes--
+	}
+}
+
+// Clone the Bitmap.
+func (b *Bitmap) Clone() Bitmap {
+	bitmap := Bitmap{b.numOnes, make([]uint64, len(b.bitBlock))}
+	copy(bitmap.bitBlock, b.bitBlock[:])
+	return bitmap
+}
+
+// countOnesForBlocks count all 1 bits within b.bitBlock of begin and that of end.
+// The begin block and end block are inclusive.
+func (b *Bitmap) countOnesForBlocks(begin, end uint32) uint64 {
+	ones := uint64(0)
+	beginBlock := begin / 64
+	endBlock := end / 64
+	for i := beginBlock; i <= endBlock; i++ {
+		ones += uint64(bits.OnesCount64(b.bitBlock[i]))
+	}
+	return ones
+}
+
+// countOnesForAllBlocks count all 1 bits in b.bitBlock.
+func (b *Bitmap) countOnesForAllBlocks() uint64 {
+	ones := uint64(0)
+	for i := 0; i < len(b.bitBlock); i++ {
+		ones += uint64(bits.OnesCount64(b.bitBlock[i]))
+	}
+	return ones
+}
+
+// flipRange flip the bits within range (begin and end). begin is inclusive and end is exclusive.
+func (b *Bitmap) flipRange(begin, end uint32) {
+	end--
+	beginBlock := begin / 64
+	endBlock := end / 64
+	if beginBlock == endBlock {
+		b.bitBlock[endBlock] ^= ((^uint64(0) << uint(begin%64)) & ((uint64(1) << (uint(end) % 64 + 1))-1))
+	} else {
+		b.bitBlock[beginBlock] ^= ^(^uint64(0) << uint(begin%64))
+		for i := beginBlock; i < endBlock; i++ {
+			b.bitBlock[i] = ^b.bitBlock[i]
+		}
+		b.bitBlock[endBlock] ^= ((uint64(1) << (uint(end) % 64 + 1))-1)
+	}
+}
+
+// clearRange clear the bits within range (begin and end). begin is inclusive and end is exclusive.
+func (b *Bitmap) clearRange(begin, end uint32) {
+	end--
+	beginBlock := begin / 64
+	endBlock := end / 64
+	if beginBlock == endBlock {
+		b.bitBlock[beginBlock] &= (((uint64(1) << uint(begin%64)) - 1) | ^((uint64(1) << (uint(end) % 64 + 1)) - 1))
+	} else {
+		b.bitBlock[beginBlock] &= ((uint64(1) << uint(begin%64)) - 1)
+		for i := beginBlock + 1; i < endBlock; i++ {
+			b.bitBlock[i] &= ^b.bitBlock[i]
+		}
+		b.bitBlock[endBlock] &= ^((uint64(1) << (uint(end) % 64 + 1)) - 1)
+	}
+}
+
+// ClearRange clear bits within range (begin and end) for the Bitmap. begin is inclusive and end is exclusive.
+func (b *Bitmap) ClearRange(begin, end uint32) {
+	blockRange := end/64 - begin/64
+	// When the number of cleared blocks is larger than half of the length of b.bitBlock,
+	// counting 1s for the entire bitmap has better performance.
+	if blockRange > uint32(len(b.bitBlock)/2) {
+		b.clearRange(begin, end)
+		b.numOnes = uint32(b.countOnesForAllBlocks())
+	} else {
+		oldRangeOnes := b.countOnesForBlocks(begin, end)
+		b.clearRange(begin, end)
+		newRangeOnes := b.countOnesForBlocks(begin, end)
+		b.numOnes += uint32(newRangeOnes - oldRangeOnes)
+	}
+}
+
+// FlipRange flip bits within range (begin and end) for the Bitmap. begin is inclusive and end is exclusive.
+func (b *Bitmap) FlipRange(begin, end uint32) {
+	blockRange := end/64 - begin/64
+	// When the number of flipped blocks is larger than half of the length of b.bitBlock,
+	// counting 1s for the entire bitmap has better performance.
+	if blockRange > uint32(len(b.bitBlock)/2) {
+		b.flipRange(begin, end)
+		b.numOnes = uint32(b.countOnesForAllBlocks())
+	} else {
+		oldRangeOnes := b.countOnesForBlocks(begin, end)
+		b.flipRange(begin, end)
+		newRangeOnes := b.countOnesForBlocks(begin, end)
+		b.numOnes += uint32(newRangeOnes - oldRangeOnes)
+	}
+}
+
+// ToSlice transform the Bitmap into slice. For example, a bitmap of [0, 1, 0, 1]
+// will return the slice [1, 3].
+func (b *Bitmap) ToSlice() []uint32 {
+	bitmapSlice := make([]uint32, 0, b.numOnes)
+	// base is the start number of a bitBlock
+	base := 0
+	for i := 0; i < len(b.bitBlock); i++ {
+		bitBlock := b.bitBlock[i]
+		// Iterate through all the numbers held by this bit block.
+		for bitBlock != 0 {
+			// Extract the lowest set 1 bit.
+			j := bitBlock & -bitBlock
+			// Interpret the bit as the in32 number it represents and add it to result.
+			bitmapSlice = append(bitmapSlice, uint32((base + int(bits.OnesCount64(j - 1)))))
+			bitBlock ^= j
+		}
+		base += 64
+	}
+	return bitmapSlice
+}
+
+// GetNumOnes return the the number of ones in the Bitmap.
+func (b *Bitmap) GetNumOnes() uint32 {
+	return b.numOnes
+}
diff --git a/pkg/bitmap/bitmap_test.go b/pkg/bitmap/bitmap_test.go
new file mode 100644
index 000000000..a7c337b55
--- /dev/null
+++ b/pkg/bitmap/bitmap_test.go
@@ -0,0 +1,294 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package bitmap
+
+import (
+	"reflect"
+	"testing"
+)
+
+// generateFilledSlice generates a slice fill with numbers.
+func generateFilledSlice(min, max, length int) []uint32 {
+	if max == min {
+		return []uint32{uint32(min)}
+	}
+	if length > (max - min) {
+		return nil
+	}
+	randSlice := make([]uint32, length)
+	if length != 0 {
+		rangeNum := uint32((max - min)/length)
+		randSlice[0], randSlice[length-1] = uint32(min), uint32(max)
+		for i := 1 ; i < length - 1; i++ {
+			randSlice[i] = randSlice[i-1] + rangeNum
+		}
+	}
+	return randSlice
+}
+
+// generateFilledBitmap generates a Bitmap filled with fillNum of numbers,
+// and returns the slice and bitmap.
+func generateFilledBitmap(min, max, fillNum int) ([]uint32, Bitmap) {
+	bitmap := BitmapWithSize(uint32(max))
+	randSlice := generateFilledSlice(min, max, fillNum)
+	for i := 0; i < fillNum; i++ {
+		bitmap.Add(randSlice[i])
+	}
+	return randSlice, bitmap
+}
+
+func TestNewBitmap(t *testing.T) {
+	tests :=  []struct {
+		name string
+		size int
+		expectSize int
+	}{
+		{"length 1", 1, 1},
+		{"length 1024", 1024, 16},
+		{"length 1025", 1025, 17},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T){
+			if bitmap := BitmapWithSize(uint32(tt.size)); len(bitmap.bitBlock) != tt.expectSize {
+				t.Errorf("BitmapWithSize created bitmap with %v, bitBlock size: %d, wanted: %d", tt.name, len(bitmap.bitBlock), tt.expectSize)
+			}
+		})
+	}
+}
+
+func TestAdd(t *testing.T) {
+	tests := []struct {
+		name string
+		bitmapSize int
+		addNum int
+	}{
+		{"Add with null bitmap.bitBlock", 0, 10},
+		{"Add without extending bitBlock",64, 10},
+		{"Add without extending bitblock with margin number", 63, 64},
+		{"Add with extended one block",1024, 1025},
+		{"Add with extended more then one block", 1024, 2048},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			bitmap := BitmapWithSize(uint32(tt.bitmapSize))
+			bitmap.Add(uint32(tt.addNum))
+			bitmapSlice := bitmap.ToSlice()
+			if bitmapSlice[0] != uint32(tt.addNum) {
+				t.Errorf("%v, get number: %d, wanted: %d.", tt.name, bitmapSlice[0], tt.addNum)
+			}
+		})
+	}
+}
+
+func TestRemove(t *testing.T) {
+	bitmap := BitmapWithSize(uint32(1024))
+	firstSlice := generateFilledSlice(0, 511, 50)
+	secondSlice := generateFilledSlice(512, 1024, 50)
+	for i := 0; i < 50; i++ {
+		bitmap.Add(firstSlice[i])
+		bitmap.Add(secondSlice[i])
+	}
+
+	for i := 0; i < 50; i++ {
+		bitmap.Remove(firstSlice[i])
+	}
+	bitmapSlice := bitmap.ToSlice()
+	if !reflect.DeepEqual(bitmapSlice, secondSlice) {
+		 t.Errorf("After Remove() firstSlice, remained slice: %v, wanted: %v", bitmapSlice, secondSlice)
+	}
+
+	for i := 0; i < 50; i++ {
+		bitmap.Remove(secondSlice[i])
+	}
+	bitmapSlice = bitmap.ToSlice()
+	emptySlice := make([]uint32,0)
+	if !reflect.DeepEqual(bitmapSlice, emptySlice) {
+		t.Errorf("After Remove secondSlice, remained slice: %v, wanted: %v", bitmapSlice, emptySlice)
+	}
+
+}
+
+// Verify flip bits within one bitBlock, one bit and bits cross multi bitBlocks.
+func TestFlipRange(t *testing.T) {
+	tests := []struct {
+		name string
+		flipRangeMin int
+		flipRangeMax int
+		filledSliceLen int
+	}{
+		{"Flip one number in bitmap", 77, 77, 1},
+		{"Flip numbers within one bitBlocks", 5, 60, 20},
+		{"Flip numbers that cross multi bitBlocks", 20, 1000, 300},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			fillSlice, bitmap := generateFilledBitmap(tt.flipRangeMin, tt.flipRangeMax, tt.filledSliceLen)
+			flipFillSlice := make([]uint32,0)
+			for i, j := tt.flipRangeMin, 0; i <= tt.flipRangeMax; i++ {
+				if uint32(i) != fillSlice[j] {
+					flipFillSlice = append(flipFillSlice, uint32(i))
+				} else {
+					j++
+				}
+			}
+
+			bitmap.FlipRange(uint32(tt.flipRangeMin), uint32(tt.flipRangeMax + 1))
+			flipBitmapSlice := bitmap.ToSlice()
+			if !reflect.DeepEqual(flipFillSlice, flipBitmapSlice) {
+				t.Errorf("%v, flipped slice: %v, wanted: %v", tt.name, flipBitmapSlice, flipFillSlice)
+			}
+		})
+	}
+}
+
+// Verify clear bits within one bitBlock, one bit and bits cross multi bitBlocks.
+func TestClearRange(t *testing.T) {
+	tests := []struct {
+		name string
+		clearRangeMin int
+		clearRangeMax int
+		bitmapSize int
+	}{
+		{"ClearRange clear one number", 5, 5, 64},
+		{"ClearRange clear numbers within one bitBlock", 4, 61, 64},
+		{"ClearRange clear numbers cross multi bitBlocks", 20, 254, 512},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			bitmap := BitmapWithSize(uint32(tt.bitmapSize))
+			bitmap.FlipRange(uint32(0), uint32(tt.bitmapSize))
+			bitmap.ClearRange(uint32(tt.clearRangeMin), uint32(tt.clearRangeMax + 1))
+			clearedBitmapSlice := bitmap.ToSlice()
+			clearedSlice := make([]uint32,0)
+			for i := 0; i < tt.bitmapSize; i++ {
+				if i < tt.clearRangeMin || i > tt.clearRangeMax {
+					clearedSlice = append(clearedSlice, uint32(i))
+				}
+			}
+			if !reflect.DeepEqual(clearedSlice, clearedBitmapSlice) {
+				t.Errorf("%v, cleared slice: %v, wanted: %v", tt.name, clearedBitmapSlice, clearedSlice)
+			}
+		})
+
+	}
+}
+
+func TestMinimum(t *testing.T) {
+	randSlice, bitmap := generateFilledBitmap(0, 1024, 200)
+	min := bitmap.Minimum()
+	if min != randSlice[0] {
+		t.Errorf("Minimum() returns: %v, wanted: %v", min, randSlice[0])
+	}
+
+	bitmap.ClearRange(uint32(0), uint32(200))
+	min = bitmap.Minimum()
+	bitmapSlice := bitmap.ToSlice()
+	if min != bitmapSlice[0] {
+		t.Errorf("After ClearRange, Minimum() returns: %v, wanted: %v", min, bitmapSlice[0])
+	}
+
+	bitmap.FlipRange(uint32(2), uint32(11))
+	min = bitmap.Minimum()
+	bitmapSlice = bitmap.ToSlice()
+	if min != bitmapSlice[0] {
+		t.Errorf("After Flip, Minimum() returns: %v, wanted: %v", min, bitmapSlice[0])
+	}
+}
+
+func TestMaximum(t *testing.T) {
+	randSlice, bitmap := generateFilledBitmap(0, 1024, 200)
+	max := bitmap.Maximum()
+	if max != randSlice[len(randSlice) - 1] {
+		t.Errorf("Maximum() returns: %v, wanted: %v", max, randSlice[len(randSlice) - 1])
+	}
+
+	bitmap.ClearRange(uint32(1000), uint32(1025))
+	max = bitmap.Maximum()
+	bitmapSlice := bitmap.ToSlice()
+	if max != bitmapSlice[len(bitmapSlice) - 1] {
+		t.Errorf("After ClearRange, Maximum() returns: %v, wanted: %v", max, bitmapSlice[len(bitmapSlice) - 1])
+	}
+
+	bitmap.FlipRange(uint32(1001), uint32(1021))
+	max = bitmap.Maximum()
+	bitmapSlice = bitmap.ToSlice()
+	if max != bitmapSlice[len(bitmapSlice) - 1] {
+		t.Errorf("After Flip, Maximum() returns: %v, wanted: %v", max, bitmapSlice[len(bitmapSlice) - 1])
+	}
+}
+
+func TestBitmapNumOnes(t *testing.T) {
+	randSlice, bitmap := generateFilledBitmap(0, 1024, 200)
+	bitmapOnes := bitmap.GetNumOnes()
+	if bitmapOnes != uint32(200) {
+		t.Errorf("GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 200)
+	}
+	// Remove 10 numbers.
+	for i := 5; i < 15; i++ {
+		bitmap.Remove(randSlice[i])
+	}
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(190) {
+		t.Errorf("After Remove 10 number, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 190)
+	}
+	// Remove the 10 number again, the length supposed not change.
+	for i := 5; i < 15; i++ {
+		bitmap.Remove(randSlice[i])
+	}
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(190) {
+		t.Errorf("After Remove the 10 number again, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 190)
+	}
+
+	// Add 10 number.
+	for i := 1080; i < 1090; i++{
+		bitmap.Add(uint32(i))
+	}
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(200) {
+		t.Errorf("After Add 10 number, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 200)
+	}
+
+	// Add the 10 number again, the length supposed not change.
+	for i := 1080; i < 1090; i++{
+		bitmap.Add(uint32(i))
+	}
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(200) {
+		t.Errorf("After Add the 10 number again, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 200)
+	}
+
+	// Flip 10 bits from 0 to 1.
+	bitmap.FlipRange(uint32(1025), uint32(1035))
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(210) {
+		t.Errorf("After Flip, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 210)
+	}
+
+	// ClearRange numbers range from [0, 1025).
+	bitmap.ClearRange(uint32(0), uint32(1025))
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(20) {
+		t.Errorf("After ClearRange, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 20)
+	}
+}