github.com/MetalBlockchain/metalgo@v1.11.9/utils/bag/unique_bag.go (about) 1 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved. 2 // See the file LICENSE for licensing terms. 3 4 package bag 5 6 import ( 7 "fmt" 8 "strings" 9 10 "golang.org/x/exp/maps" 11 12 "github.com/MetalBlockchain/metalgo/utils" 13 "github.com/MetalBlockchain/metalgo/utils/set" 14 ) 15 16 // Maps a key to a bitset. 17 type UniqueBag[T comparable] map[T]set.Bits64 18 19 func (b *UniqueBag[T]) init() { 20 if *b == nil { 21 *b = make(map[T]set.Bits64, minBagSize) 22 } 23 } 24 25 // Adds [n] to the bitset associated with each key in [keys]. 26 func (b *UniqueBag[T]) Add(n uint, keys ...T) { 27 var bs set.Bits64 28 bs.Add(n) 29 30 for _, key := range keys { 31 b.UnionSet(key, bs) 32 } 33 } 34 35 // Unions [set] with the bitset associated with [key]. 36 func (b *UniqueBag[T]) UnionSet(key T, set set.Bits64) { 37 b.init() 38 39 previousSet := (*b)[key] 40 previousSet.Union(set) 41 (*b)[key] = previousSet 42 } 43 44 // Removes each element of [set] from the bitset associated with [key]. 45 func (b *UniqueBag[T]) DifferenceSet(key T, set set.Bits64) { 46 b.init() 47 48 previousSet := (*b)[key] 49 previousSet.Difference(set) 50 (*b)[key] = previousSet 51 } 52 53 // For each key/bitset pair in [diff], removes each element of the bitset 54 // from the bitset associated with the key in [b]. 55 // Keys in [diff] that are not in [b] are ignored. 56 // Bitset elements in [diff] that are not in the bitset associated with 57 // the key in [b] are ignored. 58 func (b *UniqueBag[T]) Difference(diff *UniqueBag[T]) { 59 b.init() 60 61 for key, previousSet := range *b { 62 if previousSetDiff, exists := (*diff)[key]; exists { 63 previousSet.Difference(previousSetDiff) 64 } 65 (*b)[key] = previousSet 66 } 67 } 68 69 // Returns the bitset associated with [key]. 70 func (b *UniqueBag[T]) GetSet(key T) set.Bits64 { 71 return (*b)[key] 72 } 73 74 // Removes the bitset associated with [key]. 75 func (b *UniqueBag[T]) RemoveSet(key T) { 76 delete(*b, key) 77 } 78 79 // Returns the keys. 80 func (b *UniqueBag[T]) List() []T { 81 return maps.Keys(*b) 82 } 83 84 // Returns a bag with the given [threshold] where each key is 85 // in the bag once for each element in the key's bitset. 86 func (b *UniqueBag[T]) Bag(threshold int) Bag[T] { 87 bag := Bag[T]{ 88 counts: make(map[T]int, len(*b)), 89 } 90 bag.SetThreshold(threshold) 91 for key, bs := range *b { 92 bag.AddCount(key, bs.Len()) 93 } 94 return bag 95 } 96 97 func (b *UniqueBag[T]) PrefixedString(prefix string) string { 98 sb := strings.Builder{} 99 100 sb.WriteString(fmt.Sprintf("UniqueBag[%T]: (Size = %d)", utils.Zero[T](), len(*b))) 101 for key, set := range *b { 102 sb.WriteString(fmt.Sprintf("\n%s %v: %s", prefix, key, set)) 103 } 104 105 return sb.String() 106 } 107 108 func (b *UniqueBag[_]) String() string { 109 return b.PrefixedString("") 110 } 111 112 // Removes all key --> bitset pairs. 113 func (b *UniqueBag[_]) Clear() { 114 clear(*b) 115 }