set.go

Documentation: edge-infra.dev/pkg/sds/lib/set

     1  //nolint:revive
     2  package set
     3  
     4  import (
     5  	"maps"
     6  	"slices"
     7  )
     8  
     9  type Set[K comparable] map[K]void
    10  type void struct{}
    11  
    12  var member = void{}
    13  
    14  func New[K comparable]() Set[K] {
    15  	return Set[K]{}
    16  }
    17  
    18  // Creates a set from each item in a slice.
    19  func FromSlice[S ~[]K, K comparable](slice S) Set[K] {
    20  	set := New[K]()
    21  	set.Add(slice...)
    22  	return set
    23  }
    24  
    25  // Creates a set from each key in a map.
    26  func FromMap[M ~map[K]V, K comparable, V any](m M) Set[K] {
    27  	return FromSlice(mapKeys(m))
    28  }
    29  
    30  // Returns a slice of the members in a set.
    31  func (set Set[K]) ToSlice() []K {
    32  	return mapKeys(set)
    33  }
    34  
    35  // Adds members to the set.
    36  func (set Set[K]) Add(items ...K) {
    37  	for _, item := range items {
    38  		set[item] = member
    39  	}
    40  }
    41  
    42  // Removes members from the set, if they exists.
    43  func (set Set[K]) Remove(items ...K) {
    44  	for _, item := range items {
    45  		delete(set, item)
    46  	}
    47  }
    48  
    49  // Returns true if the set has a given member.
    50  func (set Set[K]) HasMember(item K) bool {
    51  	_, ok := set[item]
    52  	return ok
    53  }
    54  
    55  // A u B, i.e. all members of A and B.
    56  func (left Set[K]) Union(right Set[K]) Set[K] {
    57  	union := Set[K]{}
    58  
    59  	for item := range left {
    60  		union.Add(item)
    61  	}
    62  
    63  	for item := range right {
    64  		union.Add(item)
    65  	}
    66  
    67  	return union
    68  }
    69  
    70  // A n B, i.e. the members in both A and B.
    71  func (left Set[K]) Intersection(right Set[K]) Set[K] {
    72  	intersection := Set[K]{}
    73  
    74  	for item := range left {
    75  		if right.HasMember(item) {
    76  			intersection.Add(item)
    77  		}
    78  	}
    79  
    80  	return intersection
    81  }
    82  
    83  // A - B, i.e. the members of A without any members also in B.
    84  func (left Set[K]) Difference(right Set[K]) Set[K] {
    85  	difference := left
    86  
    87  	for item := range right {
    88  		difference.Remove(item)
    89  	}
    90  
    91  	return difference
    92  }
    93  
    94  // (A u B) - (A n B), i.e. the members of A and B, but not in both.
    95  func (left Set[K]) SymmetricDifference(right Set[K]) Set[K] {
    96  	return left.Union(right).Difference(left.Intersection(right))
    97  }
    98  
    99  func mapKeys[M ~map[K]V, K comparable, V any](m M) []K {
   100  	return slices.Collect(maps.Keys(m))
   101  }
   102
View as plain text