-
Notifications
You must be signed in to change notification settings - Fork 11
/
set.go
79 lines (67 loc) · 1.7 KB
/
set.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
// Copyright (C) 2023 Storj Labs, Inc.
// See LICENSE for copying information.
package location
import "math/bits"
const bitsPerBucket = 32
// Set implements a data-structure for fast lookups for country codes.
type Set [(countryCodeCount + bitsPerBucket - 1) / bitsPerBucket]uint32
// NewSet returns a set that has the specific countries set.
func NewSet(countries ...CountryCode) (r Set) {
for _, c := range countries {
r.Include(c)
}
return r
}
// NewFullSet returns a set that has every bit filled.
func NewFullSet() (r Set) {
for i := range r {
r[i] = 0xFFFFFFFF
}
return r
}
// Contains checks whether c exists in the set.
func (set *Set) Contains(c CountryCode) bool {
bucket, bit := c/bitsPerBucket, c%bitsPerBucket
if int(bucket) >= len(*set) {
return false
}
return set[bucket]&(1<<bit) != 0
}
// Include adds c to the set.
func (set *Set) Include(c CountryCode) {
bucket, bit := c/bitsPerBucket, c%bitsPerBucket
if int(bucket) >= len(*set) {
return
}
set[bucket] |= 1 << bit
}
// Remove removes c from the set.
func (set *Set) Remove(c CountryCode) {
bucket, bit := c/bitsPerBucket, c%bitsPerBucket
if int(bucket) >= len(*set) {
return
}
set[bucket] &^= 1 << bit
}
// Count returns the number of items in the set.
func (set *Set) Count() int {
total := 0
for _, v := range *set {
total += bits.OnesCount32(v)
}
return total
}
// With implements a fluid interface for constructing a set.
func (set Set) With(countries ...CountryCode) Set {
for _, c := range countries {
set.Include(c)
}
return set
}
// Without implements a fluid interface for constructing a set.
func (set Set) Without(countries ...CountryCode) Set {
for _, c := range countries {
set.Remove(c)
}
return set
}