-
Notifications
You must be signed in to change notification settings - Fork 70
/
bitfield.go
147 lines (128 loc) · 3.7 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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
// Package bitfield provides support for manipulating bits in a []byte.
package bitfield
import (
"encoding/hex"
"errors"
)
func NumBytes(length uint32) int {
return int((uint64(length) + 7) / 8)
}
// Bitfield is described in BEP 3.
type Bitfield struct {
bytes []byte
length uint32
}
// New creates a new Bitfield of length bits.
func New(length uint32) *Bitfield {
return &Bitfield{
bytes: make([]byte, NumBytes(length)),
length: length,
}
}
// NewBytes returns a new Bitfield from bytes.
// Bytes in b are not copied. Unused bits in last byte are cleared.
// Panics if b is not big enough to hold "length" bits.
func NewBytes(b []byte, length uint32) (*Bitfield, error) {
div, mod := divMod32(length)
lastByteIncomplete := mod != 0
requiredBytes := div
if lastByteIncomplete {
requiredBytes++
}
if uint32(len(b)) != requiredBytes {
return nil, errors.New("invalid length")
}
if lastByteIncomplete {
b[len(b)-1] &= ^(0xff >> mod)
}
return &Bitfield{
bytes: b[:requiredBytes],
length: length,
}, nil
}
func (b *Bitfield) Copy() *Bitfield {
b2 := &Bitfield{
bytes: make([]byte, len(b.bytes)),
length: b.length,
}
copy(b2.bytes, b.bytes)
return b2
}
// Bytes returns bytes in b. If you modify the returned slice the bits in b are modified too.
func (b *Bitfield) Bytes() []byte { return b.bytes }
// Len returns the number of bits as given to New.
func (b *Bitfield) Len() uint32 { return b.length }
// Hex returns bytes as string. If not all the bits in last byte are used, they encode as not set.
func (b *Bitfield) Hex() string {
return hex.EncodeToString(b.bytes)
}
// Set bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b *Bitfield) Set(i uint32) {
b.checkIndex(i)
div, mod := divMod32(i)
b.bytes[div] |= 1 << (7 - mod)
}
// Clear bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b *Bitfield) Clear(i uint32) {
b.checkIndex(i)
div, mod := divMod32(i)
b.bytes[div] &= ^(1 << (7 - mod))
}
// Test bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b *Bitfield) Test(i uint32) bool {
b.checkIndex(i)
div, mod := divMod32(i)
return (b.bytes[div] & (1 << (7 - mod))) > 0
}
var countCache = [256]byte{
0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8,
}
// Count returns the count of set bits.
func (b *Bitfield) Count() uint32 {
var total uint32
for _, v := range b.bytes {
total += uint32(countCache[v])
}
return total
}
// All returns true if all bits are set, false otherwise.
func (b *Bitfield) All() bool {
return b.Count() == b.length
}
func (b *Bitfield) checkIndex(i uint32) {
if i >= b.Len() {
panic("index out of bound")
}
}
func (b *Bitfield) And(b2 *Bitfield) {
if b.length != b2.length {
panic("length mismatch")
}
for i := range b.bytes {
b.bytes[i] &= b2.bytes[i]
}
}
func (b *Bitfield) Or(b2 *Bitfield) {
if b.length != b2.length {
panic("length mismatch")
}
for i := range b.bytes {
b.bytes[i] |= b2.bytes[i]
}
}
func divMod32(a uint32) (uint32, uint32) { return a / 8, a % 8 }