forked from ava-labs/avalanchego
-
Notifications
You must be signed in to change notification settings - Fork 4
/
unique_bag.go
115 lines (92 loc) · 2.61 KB
/
unique_bag.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
// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package bag
import (
"fmt"
"strings"
"golang.org/x/exp/maps"
"github.com/MetalBlockchain/metalgo/utils"
"github.com/MetalBlockchain/metalgo/utils/set"
)
// Maps a key to a bitset.
type UniqueBag[T comparable] map[T]set.Bits64
func (b *UniqueBag[T]) init() {
if *b == nil {
*b = make(map[T]set.Bits64, minBagSize)
}
}
// Adds [n] to the bitset associated with each key in [keys].
func (b *UniqueBag[T]) Add(n uint, keys ...T) {
var bs set.Bits64
bs.Add(n)
for _, key := range keys {
b.UnionSet(key, bs)
}
}
// Unions [set] with the bitset associated with [key].
func (b *UniqueBag[T]) UnionSet(key T, set set.Bits64) {
b.init()
previousSet := (*b)[key]
previousSet.Union(set)
(*b)[key] = previousSet
}
// Removes each element of [set] from the bitset associated with [key].
func (b *UniqueBag[T]) DifferenceSet(key T, set set.Bits64) {
b.init()
previousSet := (*b)[key]
previousSet.Difference(set)
(*b)[key] = previousSet
}
// For each key/bitset pair in [diff], removes each element of the bitset
// from the bitset associated with the key in [b].
// Keys in [diff] that are not in [b] are ignored.
// Bitset elements in [diff] that are not in the bitset associated with
// the key in [b] are ignored.
func (b *UniqueBag[T]) Difference(diff *UniqueBag[T]) {
b.init()
for key, previousSet := range *b {
if previousSetDiff, exists := (*diff)[key]; exists {
previousSet.Difference(previousSetDiff)
}
(*b)[key] = previousSet
}
}
// Returns the bitset associated with [key].
func (b *UniqueBag[T]) GetSet(key T) set.Bits64 {
return (*b)[key]
}
// Removes the bitset associated with [key].
func (b *UniqueBag[T]) RemoveSet(key T) {
delete(*b, key)
}
// Returns the keys.
func (b *UniqueBag[T]) List() []T {
return maps.Keys(*b)
}
// Returns a bag with the given [threshold] where each key is
// in the bag once for each element in the key's bitset.
func (b *UniqueBag[T]) Bag(threshold int) Bag[T] {
bag := Bag[T]{
counts: make(map[T]int, len(*b)),
}
bag.SetThreshold(threshold)
for key, bs := range *b {
bag.AddCount(key, bs.Len())
}
return bag
}
func (b *UniqueBag[T]) PrefixedString(prefix string) string {
sb := strings.Builder{}
sb.WriteString(fmt.Sprintf("UniqueBag[%T]: (Size = %d)", utils.Zero[T](), len(*b)))
for key, set := range *b {
sb.WriteString(fmt.Sprintf("\n%s %v: %s", prefix, key, set))
}
return sb.String()
}
func (b *UniqueBag[_]) String() string {
return b.PrefixedString("")
}
// Removes all key --> bitset pairs.
func (b *UniqueBag[_]) Clear() {
maps.Clear(*b)
}