-
Notifications
You must be signed in to change notification settings - Fork 2
/
bitfield.go
119 lines (105 loc) · 2.58 KB
/
bitfield.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
package bitfield
// NOTE: Don't bother replacing the divisions/modulo with shifts/ands, go is smart.
import (
"fmt"
"math/bits"
)
// NewBitfield creates a new fixed-sized Bitfield (allocated up-front).
func NewBitfield(size int) (Bitfield, error) {
if size < 0 {
return nil, fmt.Errorf("bitfield size must be positive; got %d", size)
}
if size%8 != 0 {
return nil, fmt.Errorf("bitfield size must be a multiple of 8; got %d", size)
}
return make([]byte, size/8), nil
}
// FromBytes constructs a new bitfield from a serialized bitfield.
func FromBytes(size int, bits []byte) (Bitfield, error) {
bf, err := NewBitfield(size)
if err != nil {
return nil, err
}
start := len(bf) - len(bits)
if start < 0 {
return nil, fmt.Errorf("bitfield too small: got %d; need %d", size, len(bits)*8)
}
copy(bf[start:], bits)
return bf, nil
}
func (bf Bitfield) offset(i int) (uint, uint8) {
return uint(len(bf)) - (uint(i) / 8) - 1, uint8(i) % 8
}
// Bitfield is, well, a bitfield.
type Bitfield []byte
// Bytes returns the Bitfield as a byte string.
//
// This function *does not* copy.
func (bf Bitfield) Bytes() []byte {
for i, b := range bf {
if b != 0 {
return bf[i:]
}
}
return nil
}
// Bit returns the ith bit.
//
// Panics if the bit is out of bounds.
func (bf Bitfield) Bit(i int) bool {
idx, off := bf.offset(i)
return (bf[idx]>>off)&0x1 != 0
}
// SetBit sets the ith bit.
//
// Panics if the bit is out of bounds.
func (bf Bitfield) SetBit(i int) {
idx, off := bf.offset(i)
bf[idx] |= 1 << off
}
// UnsetBit unsets the ith bit.
//
// Panics if the bit is out of bounds.
func (bf Bitfield) UnsetBit(i int) {
idx, off := bf.offset(i)
bf[idx] &= 0xFF ^ (1 << off)
}
// SetBytes sets the bits to the given byte array.
//
// Panics if 'b' is larger than the bitfield.
func (bf Bitfield) SetBytes(b []byte) {
start := len(bf) - len(b)
if start < 0 {
panic("bitfield too small")
}
for i := range bf[:start] {
bf[i] = 0
}
copy(bf[start:], b)
}
// Ones returns the number of bits set.
func (bf Bitfield) Ones() int {
cnt := 0
for _, b := range bf {
cnt += bits.OnesCount8(b)
}
return cnt
}
// OnesBefore returns the number of bits set *before* this bit.
func (bf Bitfield) OnesBefore(i int) int {
idx, off := bf.offset(i)
cnt := bits.OnesCount8(bf[idx] << (8 - off))
for _, b := range bf[idx+1:] {
cnt += bits.OnesCount8(b)
}
return cnt
}
// OnesAfter returns the number of bits set *after* this bit.
func (bf Bitfield) OnesAfter(i int) int {
idx, off := bf.offset(i)
cnt := bits.OnesCount8(bf[idx] >> off)
for _, b := range bf[:idx] {
cnt += bits.OnesCount8(b)
}
return cnt
}