github.com/MetalBlockchain/metalgo@v1.11.9/utils/set/bits.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package set
     5  
     6  import (
     7  	"encoding/hex"
     8  	"math/big"
     9  	"math/bits"
    10  )
    11  
    12  // Bits is a bit-set backed by a big.Int
    13  // Holds values ranging from [0, INT_MAX] (arch-dependent)
    14  // Trying to use negative values will result in a panic.
    15  // This implementation is NOT thread-safe.
    16  type Bits struct {
    17  	bits *big.Int
    18  }
    19  
    20  // NewBits returns a new instance of Bits with [bits] set to 1.
    21  //
    22  // Invariants:
    23  // 1. Negative bits will cause a panic.
    24  // 2. Duplicate bits are allowed but will cause a no-op.
    25  func NewBits(bits ...int) Bits {
    26  	b := Bits{new(big.Int)}
    27  	for _, bit := range bits {
    28  		b.Add(bit)
    29  	}
    30  	return b
    31  }
    32  
    33  // Add sets the [i]'th bit to 1
    34  func (b Bits) Add(i int) {
    35  	b.bits.SetBit(b.bits, i, 1)
    36  }
    37  
    38  // Union performs the set union with another set.
    39  // This adds all elements in [other] to [b]
    40  func (b Bits) Union(other Bits) {
    41  	b.bits.Or(b.bits, other.bits)
    42  }
    43  
    44  // Intersection performs the set intersection with another set
    45  // This sets [b] to include only elements in both [b] and [other]
    46  func (b Bits) Intersection(other Bits) {
    47  	b.bits.And(b.bits, other.bits)
    48  }
    49  
    50  // Difference removes all the elements in [other] from this set
    51  func (b Bits) Difference(other Bits) {
    52  	b.bits.AndNot(b.bits, other.bits)
    53  }
    54  
    55  // Remove sets the [i]'th bit to 0
    56  func (b Bits) Remove(i int) {
    57  	b.bits.SetBit(b.bits, i, 0)
    58  }
    59  
    60  // Clear empties out the bitset
    61  func (b Bits) Clear() {
    62  	b.bits.SetUint64(0)
    63  }
    64  
    65  // Contains returns true if the [i]'th bit is 1, and false otherwise
    66  func (b Bits) Contains(i int) bool {
    67  	return b.bits.Bit(i) == 1
    68  }
    69  
    70  // BitLen returns the bit length of this bitset
    71  func (b Bits) BitLen() int {
    72  	return b.bits.BitLen()
    73  }
    74  
    75  // Len returns the amount of 1's in the bitset
    76  //
    77  // This is typically referred to as the "Hamming Weight"
    78  // of a set of bits.
    79  func (b Bits) Len() int {
    80  	result := 0
    81  	for _, word := range b.bits.Bits() {
    82  		result += bits.OnesCount(uint(word))
    83  	}
    84  	return result
    85  }
    86  
    87  // Returns the byte representation of this bitset
    88  func (b Bits) Bytes() []byte {
    89  	return b.bits.Bytes()
    90  }
    91  
    92  // Inverse of Bits.Bytes()
    93  func BitsFromBytes(bytes []byte) Bits {
    94  	return Bits{
    95  		bits: new(big.Int).SetBytes(bytes),
    96  	}
    97  }
    98  
    99  // String returns the hex representation of this bitset
   100  func (b Bits) String() string {
   101  	return hex.EncodeToString(b.bits.Bytes())
   102  }