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  }