github.com/sagernet/gvisor@v0.0.0-20240428053021-e691de28565f/pkg/bitmap/bitmap.go (about)

     1  // Copyright 2021 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package bitmap provides the implementation of bitmap.
    16  package bitmap
    17  
    18  import (
    19  	"fmt"
    20  	"math"
    21  	"math/bits"
    22  )
    23  
    24  // MaxBitEntryLimit defines the upper limit on how many bit entries are supported by this Bitmap
    25  // implementation.
    26  const MaxBitEntryLimit uint32 = math.MaxInt32
    27  
    28  // Bitmap implements an efficient bitmap.
    29  //
    30  // +stateify savable
    31  type Bitmap struct {
    32  	// numOnes is the number of ones in the bitmap.
    33  	numOnes uint32
    34  
    35  	// bitBlock holds the bits. The type of bitBlock is uint64 which means
    36  	// each number in bitBlock contains 64 entries.
    37  	bitBlock []uint64
    38  }
    39  
    40  // New create a new empty Bitmap.
    41  func New(size uint32) Bitmap {
    42  	b := Bitmap{}
    43  	bSize := (size + 63) / 64
    44  	b.bitBlock = make([]uint64, bSize)
    45  	return b
    46  }
    47  
    48  // IsEmpty verifies whether the Bitmap is empty.
    49  func (b *Bitmap) IsEmpty() bool {
    50  	return b.numOnes == 0
    51  }
    52  
    53  // Size returns the total number of bits in the bitmap.
    54  func (b *Bitmap) Size() int {
    55  	return len(b.bitBlock) * 64
    56  }
    57  
    58  // Grow grows the bitmap by at least toGrow bits.
    59  func (b *Bitmap) Grow(toGrow uint32) error {
    60  	newbitBlockSize := uint32(len(b.bitBlock)) + ((toGrow + 63) / 64)
    61  	if newbitBlockSize > MaxBitEntryLimit/8 {
    62  		return fmt.Errorf("requested bitmap size %d too large", newbitBlockSize*64)
    63  	}
    64  	bits := make([]uint64, (toGrow+63)/64)
    65  	b.bitBlock = append(b.bitBlock, bits...)
    66  	return nil
    67  }
    68  
    69  // Minimum return the smallest value in the Bitmap.
    70  func (b *Bitmap) Minimum() uint32 {
    71  	for i := 0; i < len(b.bitBlock); i++ {
    72  		if w := b.bitBlock[i]; w != 0 {
    73  			r := bits.TrailingZeros64(w)
    74  			return uint32(r + i*64)
    75  		}
    76  	}
    77  	return MaxBitEntryLimit
    78  }
    79  
    80  // FirstZero returns the first unset bit from the range [start, ).
    81  func (b *Bitmap) FirstZero(start uint32) (bit uint32, err error) {
    82  	i, nbit := int(start/64), start%64
    83  	n := len(b.bitBlock)
    84  	if i >= n {
    85  		return MaxBitEntryLimit, fmt.Errorf("given start of range exceeds bitmap size")
    86  	}
    87  	w := b.bitBlock[i] | ((1 << nbit) - 1)
    88  	for {
    89  		if w != ^uint64(0) {
    90  			r := bits.TrailingZeros64(^w)
    91  			return uint32(r + i*64), nil
    92  		}
    93  		i++
    94  		if i == n {
    95  			break
    96  		}
    97  		w = b.bitBlock[i]
    98  	}
    99  	return MaxBitEntryLimit, fmt.Errorf("bitmap has no unset bits")
   100  }
   101  
   102  // FirstOne returns the first set bit from the range [start, )
   103  func (b *Bitmap) FirstOne(start uint32) (bit uint32, err error) {
   104  	i, nbit := int(start/64), start%64
   105  	n := len(b.bitBlock)
   106  	if i >= n {
   107  		return MaxBitEntryLimit, fmt.Errorf("given start of range exceeds bitmap size")
   108  	}
   109  	w := b.bitBlock[i] & (math.MaxUint64 << nbit)
   110  	for {
   111  		if w != uint64(0) {
   112  			r := bits.TrailingZeros64(w)
   113  			return uint32(r + i*64), nil
   114  		}
   115  		i++
   116  		if i == n {
   117  			break
   118  		}
   119  		w = b.bitBlock[i]
   120  	}
   121  	return MaxBitEntryLimit, fmt.Errorf("bitmap has no set bits")
   122  }
   123  
   124  // Maximum return the largest value in the Bitmap.
   125  func (b *Bitmap) Maximum() uint32 {
   126  	for i := len(b.bitBlock) - 1; i >= 0; i-- {
   127  		if w := b.bitBlock[i]; w != 0 {
   128  			r := bits.LeadingZeros64(w)
   129  			return uint32(i*64 + 63 - r)
   130  		}
   131  	}
   132  	return uint32(0)
   133  }
   134  
   135  // Add add i to the Bitmap.
   136  func (b *Bitmap) Add(i uint32) {
   137  	blockNum, mask := i/64, uint64(1)<<(i%64)
   138  	// if blockNum is out of range, extend b.bitBlock
   139  	if x, y := int(blockNum), len(b.bitBlock); x >= y {
   140  		b.bitBlock = append(b.bitBlock, make([]uint64, x-y+1)...)
   141  	}
   142  	oldBlock := b.bitBlock[blockNum]
   143  	newBlock := oldBlock | mask
   144  	if oldBlock != newBlock {
   145  		b.bitBlock[blockNum] = newBlock
   146  		b.numOnes++
   147  	}
   148  }
   149  
   150  // Remove i from the Bitmap.
   151  func (b *Bitmap) Remove(i uint32) {
   152  	blockNum, mask := i/64, uint64(1)<<(i%64)
   153  	oldBlock := b.bitBlock[blockNum]
   154  	newBlock := oldBlock &^ mask
   155  	if oldBlock != newBlock {
   156  		b.bitBlock[blockNum] = newBlock
   157  		b.numOnes--
   158  	}
   159  }
   160  
   161  // Clone the Bitmap.
   162  func (b *Bitmap) Clone() Bitmap {
   163  	bitmap := Bitmap{b.numOnes, make([]uint64, len(b.bitBlock))}
   164  	copy(bitmap.bitBlock, b.bitBlock[:])
   165  	return bitmap
   166  }
   167  
   168  // countOnesForBlocks count all 1 bits within b.bitBlock of begin and that of end.
   169  // The begin block and end block are inclusive.
   170  func (b *Bitmap) countOnesForBlocks(begin, end uint32) uint64 {
   171  	ones := uint64(0)
   172  	beginBlock := begin / 64
   173  	endBlock := end / 64
   174  	for i := beginBlock; i <= endBlock; i++ {
   175  		ones += uint64(bits.OnesCount64(b.bitBlock[i]))
   176  	}
   177  	return ones
   178  }
   179  
   180  // countOnesForAllBlocks count all 1 bits in b.bitBlock.
   181  func (b *Bitmap) countOnesForAllBlocks() uint64 {
   182  	ones := uint64(0)
   183  	for i := 0; i < len(b.bitBlock); i++ {
   184  		ones += uint64(bits.OnesCount64(b.bitBlock[i]))
   185  	}
   186  	return ones
   187  }
   188  
   189  // flipRange flip the bits within range (begin and end). begin is inclusive and end is exclusive.
   190  func (b *Bitmap) flipRange(begin, end uint32) {
   191  	end--
   192  	beginBlock := begin / 64
   193  	endBlock := end / 64
   194  	if beginBlock == endBlock {
   195  		b.bitBlock[endBlock] ^= ((^uint64(0) << uint(begin%64)) & ((uint64(1) << (uint(end)%64 + 1)) - 1))
   196  	} else {
   197  		b.bitBlock[beginBlock] ^= ^(^uint64(0) << uint(begin%64))
   198  		for i := beginBlock; i < endBlock; i++ {
   199  			b.bitBlock[i] = ^b.bitBlock[i]
   200  		}
   201  		b.bitBlock[endBlock] ^= ((uint64(1) << (uint(end)%64 + 1)) - 1)
   202  	}
   203  }
   204  
   205  // clearRange clear the bits within range (begin and end). begin is inclusive and end is exclusive.
   206  func (b *Bitmap) clearRange(begin, end uint32) {
   207  	end--
   208  	beginBlock := begin / 64
   209  	endBlock := end / 64
   210  	if beginBlock == endBlock {
   211  		b.bitBlock[beginBlock] &= (((uint64(1) << uint(begin%64)) - 1) | ^((uint64(1) << (uint(end)%64 + 1)) - 1))
   212  	} else {
   213  		b.bitBlock[beginBlock] &= ((uint64(1) << uint(begin%64)) - 1)
   214  		for i := beginBlock + 1; i < endBlock; i++ {
   215  			b.bitBlock[i] &= ^b.bitBlock[i]
   216  		}
   217  		b.bitBlock[endBlock] &= ^((uint64(1) << (uint(end)%64 + 1)) - 1)
   218  	}
   219  }
   220  
   221  // ClearRange clear bits within range (begin and end) for the Bitmap. begin is inclusive and end is exclusive.
   222  func (b *Bitmap) ClearRange(begin, end uint32) {
   223  	blockRange := end/64 - begin/64
   224  	// When the number of cleared blocks is larger than half of the length of b.bitBlock,
   225  	// counting 1s for the entire bitmap has better performance.
   226  	if blockRange > uint32(len(b.bitBlock)/2) {
   227  		b.clearRange(begin, end)
   228  		b.numOnes = uint32(b.countOnesForAllBlocks())
   229  	} else {
   230  		oldRangeOnes := b.countOnesForBlocks(begin, end)
   231  		b.clearRange(begin, end)
   232  		newRangeOnes := b.countOnesForBlocks(begin, end)
   233  		b.numOnes += uint32(newRangeOnes - oldRangeOnes)
   234  	}
   235  }
   236  
   237  // FlipRange flip bits within range (begin and end) for the Bitmap. begin is inclusive and end is exclusive.
   238  func (b *Bitmap) FlipRange(begin, end uint32) {
   239  	blockRange := end/64 - begin/64
   240  	// When the number of flipped blocks is larger than half of the length of b.bitBlock,
   241  	// counting 1s for the entire bitmap has better performance.
   242  	if blockRange > uint32(len(b.bitBlock)/2) {
   243  		b.flipRange(begin, end)
   244  		b.numOnes = uint32(b.countOnesForAllBlocks())
   245  	} else {
   246  		oldRangeOnes := b.countOnesForBlocks(begin, end)
   247  		b.flipRange(begin, end)
   248  		newRangeOnes := b.countOnesForBlocks(begin, end)
   249  		b.numOnes += uint32(newRangeOnes - oldRangeOnes)
   250  	}
   251  }
   252  
   253  // ForEach calls `f` for each set bit in the range [start, end).
   254  //
   255  // If f returns false, ForEach stops the iteration.
   256  func (b *Bitmap) ForEach(start, end uint32, f func(idx uint32) bool) {
   257  	blockEnd := (end + 63) / 64
   258  	if blockEnd > uint32(len(b.bitBlock)) {
   259  		blockEnd = uint32(len(b.bitBlock))
   260  	}
   261  	// base is the start number of a bitBlock
   262  	base := start / 64 * 64
   263  	blockMask := ^((uint64(1) << (start % 64)) - 1)
   264  	for i := start / 64; i < blockEnd; i++ {
   265  		if i == end/64 {
   266  			blockMask &= (uint64(1) << (end % 64)) - 1
   267  		}
   268  		bitBlock := b.bitBlock[i] & blockMask
   269  		blockMask = ^uint64(0)
   270  		// Iterate through all the numbers held by this bit block.
   271  		for bitBlock != 0 {
   272  			// Extract the lowest set 1 bit.
   273  			j := bitBlock & -bitBlock
   274  			// Interpret the bit as the in32 number it represents and add it to result.
   275  			idx := base + uint32(bits.OnesCount64(j-1))
   276  			if !f(idx) {
   277  				return
   278  			}
   279  			bitBlock ^= j
   280  		}
   281  		base += 64
   282  	}
   283  }
   284  
   285  // ToSlice transform the Bitmap into slice. For example, a bitmap of [0, 1, 0, 1]
   286  // will return the slice [1, 3].
   287  func (b *Bitmap) ToSlice() []uint32 {
   288  	bitmapSlice := make([]uint32, 0, b.numOnes)
   289  	// base is the start number of a bitBlock
   290  	base := 0
   291  	for i := 0; i < len(b.bitBlock); i++ {
   292  		bitBlock := b.bitBlock[i]
   293  		// Iterate through all the numbers held by this bit block.
   294  		for bitBlock != 0 {
   295  			// Extract the lowest set 1 bit.
   296  			j := bitBlock & -bitBlock
   297  			// Interpret the bit as the in32 number it represents and add it to result.
   298  			bitmapSlice = append(bitmapSlice, uint32((base + int(bits.OnesCount64(j-1)))))
   299  			bitBlock ^= j
   300  		}
   301  		base += 64
   302  	}
   303  	return bitmapSlice
   304  }
   305  
   306  // GetNumOnes return the number of ones in the Bitmap.
   307  func (b *Bitmap) GetNumOnes() uint32 {
   308  	return b.numOnes
   309  }