-
Notifications
You must be signed in to change notification settings - Fork 0
/
buckets.go
97 lines (85 loc) · 2.94 KB
/
buckets.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
package boom
// Buckets is a fast, space-efficient array of buckets where each bucket can
// store up to a configured maximum value.
type Buckets struct {
data []byte
bucketSize uint8
max uint8
count uint
}
// NewBuckets creates a new Buckets with the provided number of buckets where
// each bucket is the specified number of bits.
func NewBuckets(count uint, bucketSize uint8) *Buckets {
return &Buckets{
count: count,
data: make([]byte, (count*uint(bucketSize)+7)/8),
bucketSize: bucketSize,
max: (1 << bucketSize) - 1,
}
}
// MaxBucketValue returns the maximum value that can be stored in a bucket.
func (b *Buckets) MaxBucketValue() uint8 {
return b.max
}
// Count returns the number of buckets.
func (b *Buckets) Count() uint {
return b.count
}
// Increment will increment the value in the specified bucket by the provided
// delta. A bucket can be decremented by providing a negative delta. The value
// is clamped to zero and the maximum bucket value. Returns itself to allow for
// chaining.
func (b *Buckets) Increment(bucket uint, delta int32) *Buckets {
val := int32(b.getBits(bucket*uint(b.bucketSize), uint(b.bucketSize))) + delta
if val > int32(b.max) {
val = int32(b.max)
} else if val < 0 {
val = 0
}
b.setBits(uint32(bucket)*uint32(b.bucketSize), uint32(b.bucketSize), uint32(val))
return b
}
// Set will set the bucket value. The value is clamped to zero and the maximum
// bucket value. Returns itself to allow for chaining.
func (b *Buckets) Set(bucket uint, value uint8) *Buckets {
if value > b.max {
value = b.max
}
b.setBits(uint32(bucket)*uint32(b.bucketSize), uint32(b.bucketSize), uint32(value))
return b
}
// Get returns the value in the specified bucket.
func (b *Buckets) Get(bucket uint) uint32 {
return b.getBits(bucket*uint(b.bucketSize), uint(b.bucketSize))
}
// Reset restores the Buckets to the original state. Returns itself to allow
// for chaining.
func (b *Buckets) Reset() *Buckets {
b.data = make([]byte, (b.count*uint(b.bucketSize)+7)/8)
return b
}
// getBits returns the bits at the specified offset and length.
func (b *Buckets) getBits(offset, length uint) uint32 {
byteIndex := offset / 8
byteOffset := offset % 8
if byteOffset+length > 8 {
rem := 8 - byteOffset
return b.getBits(offset, rem) | (b.getBits(offset+rem, length-rem) << rem)
}
bitMask := uint32((1 << length) - 1)
return (uint32(b.data[byteIndex]) & (bitMask << byteOffset)) >> byteOffset
}
// setBits sets bits at the specified offset and length.
func (b *Buckets) setBits(offset, length, bits uint32) {
byteIndex := offset / 8
byteOffset := offset % 8
if byteOffset+length > 8 {
rem := 8 - byteOffset
b.setBits(offset, rem, bits)
b.setBits(offset+rem, length-rem, bits>>rem)
return
}
bitMask := uint32((1 << length) - 1)
b.data[byteIndex] = byte(uint32(b.data[byteIndex]) & ^(bitMask << byteOffset))
b.data[byteIndex] = byte(uint32(b.data[byteIndex]) | ((bits & bitMask) << byteOffset))
}