/
bitmask.go
212 lines (185 loc) · 4.78 KB
/
bitmask.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
/*
Copyright 2019 The Kubernetes Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package bitmask
import (
"fmt"
"math/bits"
"strconv"
)
// BitMask interface allows hint providers to create BitMasks for TopologyHints
type BitMask interface {
Add(bits ...int) error
Remove(bits ...int) error
And(masks ...BitMask)
Or(masks ...BitMask)
Clear()
Fill()
IsEqual(mask BitMask) bool
IsEmpty() bool
IsSet(bit int) bool
AnySet(bits []int) bool
IsNarrowerThan(mask BitMask) bool
String() string
Count() int
GetBits() []int
}
type bitMask uint64
// NewEmptyBitMask creates a new, empty BitMask
func NewEmptyBitMask() BitMask {
s := bitMask(0)
return &s
}
// NewBitMask creates a new BitMask
func NewBitMask(bits ...int) (BitMask, error) {
s := bitMask(0)
err := (&s).Add(bits...)
if err != nil {
return nil, err
}
return &s, nil
}
// Add adds the bits with topology affinity to the BitMask
func (s *bitMask) Add(bits ...int) error {
mask := *s
for _, i := range bits {
if i < 0 || i >= 64 {
return fmt.Errorf("bit number must be in range 0-63")
}
mask |= 1 << uint64(i)
}
*s = mask
return nil
}
// Remove removes specified bits from BitMask
func (s *bitMask) Remove(bits ...int) error {
mask := *s
for _, i := range bits {
if i < 0 || i >= 64 {
return fmt.Errorf("bit number must be in range 0-63")
}
mask &^= 1 << uint64(i)
}
*s = mask
return nil
}
// And performs and operation on all bits in masks
func (s *bitMask) And(masks ...BitMask) {
for _, m := range masks {
*s &= *m.(*bitMask)
}
}
// Or performs or operation on all bits in masks
func (s *bitMask) Or(masks ...BitMask) {
for _, m := range masks {
*s |= *m.(*bitMask)
}
}
// Clear resets all bits in mask to zero
func (s *bitMask) Clear() {
*s = 0
}
// Fill sets all bits in mask to one
func (s *bitMask) Fill() {
*s = bitMask(^uint64(0))
}
// IsEmpty checks mask to see if all bits are zero
func (s *bitMask) IsEmpty() bool {
return *s == 0
}
// IsSet checks bit in mask to see if bit is set to one
func (s *bitMask) IsSet(bit int) bool {
if bit < 0 || bit >= 64 {
return false
}
return (*s & (1 << uint64(bit))) > 0
}
// AnySet checks bit in mask to see if any provided bit is set to one
func (s *bitMask) AnySet(bits []int) bool {
for _, b := range bits {
if s.IsSet(b) {
return true
}
}
return false
}
// IsEqual checks if masks are equal
func (s *bitMask) IsEqual(mask BitMask) bool {
return *s == *mask.(*bitMask)
}
// IsNarrowerThan checks if one mask is narrower than another.
//
// A mask is said to be "narrower" than another if it has lets bits set. If the
// same number of bits are set in both masks, then the mask with more
// lower-numbered bits set wins out.
func (s *bitMask) IsNarrowerThan(mask BitMask) bool {
if s.Count() == mask.Count() {
if *s < *mask.(*bitMask) {
return true
}
}
return s.Count() < mask.Count()
}
// String converts mask to string
func (s *bitMask) String() string {
grouping := 2
for shift := 64 - grouping; shift > 0; shift -= grouping {
if *s > (1 << uint(shift)) {
return fmt.Sprintf("%0"+strconv.Itoa(shift+grouping)+"b", *s)
}
}
return fmt.Sprintf("%0"+strconv.Itoa(grouping)+"b", *s)
}
// Count counts number of bits in mask set to one
func (s *bitMask) Count() int {
return bits.OnesCount64(uint64(*s))
}
// Getbits returns each bit number with bits set to one
func (s *bitMask) GetBits() []int {
var bits []int
for i := uint64(0); i < 64; i++ {
if (*s & (1 << i)) > 0 {
bits = append(bits, int(i))
}
}
return bits
}
// And is a package level implementation of 'and' between first and masks
func And(first BitMask, masks ...BitMask) BitMask {
s := *first.(*bitMask)
s.And(masks...)
return &s
}
// Or is a package level implementation of 'or' between first and masks
func Or(first BitMask, masks ...BitMask) BitMask {
s := *first.(*bitMask)
s.Or(masks...)
return &s
}
// IterateBitMasks iterates all possible masks from a list of bits,
// issuing a callback on each mask.
func IterateBitMasks(bits []int, callback func(BitMask)) {
var iterate func(bits, accum []int, size int)
iterate = func(bits, accum []int, size int) {
if len(accum) == size {
mask, _ := NewBitMask(accum...)
callback(mask)
return
}
for i := range bits {
iterate(bits[i+1:], append(accum, bits[i]), size)
}
}
for i := 1; i <= len(bits); i++ {
iterate(bits, []int{}, i)
}
}