octopus/pkg/tools/set.go
2023-03-22 22:45:17 +08:00

102 lines
1.6 KiB
Go

package tools
import (
"sync"
)
type Set[T comparable] struct {
mu sync.RWMutex
set map[T]struct{}
}
func NewSetWithCapacity[T comparable](c int) *Set[T] {
set := make(map[T]struct{}, c)
return &Set[T]{set: set}
}
func NewSet[T comparable](elems ...T) *Set[T] {
set := make(map[T]struct{}, len(elems))
for _, elem := range elems {
set[elem] = struct{}{}
}
return &Set[T]{set: set}
}
func (s *Set[T]) Add(ele ...T) *Set[T] {
s.mu.Lock()
defer s.mu.Unlock()
for _, e := range ele {
s.set[e] = struct{}{}
}
return s
}
func (s *Set[T]) Del(ele ...T) *Set[T] {
s.mu.Lock()
defer s.mu.Unlock()
for _, e := range ele {
delete(s.set, e)
}
return s
}
func (s *Set[T]) Has(ele T) bool {
s.mu.RLock()
defer s.mu.RUnlock()
_, exists := s.set[ele]
return exists
}
// 取交集.
func (s *Set[T]) Intersect(other *Set[T]) *Set[T] {
s.mu.RLock()
defer s.mu.RUnlock()
ret := NewSet[T]()
for ele := range s.set {
if other.Has(ele) {
ret.Add(ele)
}
}
return ret
}
func (s *Set[T]) Len() int {
return len(s.set)
}
func (s *Set[T]) Set() map[T]struct{} {
return s.set
}
func (s *Set[T]) Values() []T {
s.mu.RLock()
defer s.mu.RUnlock()
ret := make([]T, 0, len(s.set))
for ele := range s.set {
ret = append(ret, ele)
}
return ret
}
func (s *Set[T]) Compare(other *Set[T]) (less, more *Set[T]) {
less = NewSet[T]()
more = NewSet[T]()
// less: s - other
s.mu.RLock()
other.mu.RLock()
defer s.mu.RUnlock()
defer other.mu.RUnlock()
for ele := range other.set {
if !s.Has(ele) {
less.Add(ele)
}
}
// more: other - s
for ele := range s.set {
if !other.Has(ele) {
more.Add(ele)
}
}
return less, more
}