go-collect

set.go

package collect

// Set represents a generic set of comparable elements.
type Set[T comparable] map[T]struct{}

// Has checks whether the set holds a given value.
func (s Set[T]) Has(v T) bool {
	_, ok := s[v]
	return ok
}

// Add adds values to the set. You may add as many as you wish.
func (s Set[T]) Add(vs ...T) {
	for _, v := range vs {
		s[v] = struct{}{}
	}
}

// Remove removes values from the set. You may remove as many as you wish.
func (s Set[T]) Remove(vs ...T) {
	for _, v := range vs {
		delete(s, v)
	}
}

// Clear reinitializes the set.
func (s *Set[T]) Clear() {
	(*s) = NewSet[T]()
}

// Items returns a slice of all elements currently in the set.
// Note that unless you actually want a collection, you should call Iter or
// Filter instead. Items calls Iter internally.
func (s Set[T]) Items() []T {
	out := make([]T, 0, len(s))
	for v := range s.Iter() {
		out = append(out, v)
	}
	return out
}

func fromChan[T comparable](cs ...chan T) Set[T] {
	out := NewSet[T]()
	for _, c := range cs {
		for v := range c {
			out.Add(v)
		}
	}
	return out
}

// Difference returns a new set that holds the elements present in the receiver, but not s2.
func (s Set[T]) Difference(s2 Set[T]) Set[T] {
	return fromChan(s.Filter(func(v T)bool {
		return !s2.Has(v)
	}))
}

// Intersection returns a new set that holds elements present in both sets.
func (s Set[T]) Intersection(s2 Set[T]) Set[T] {
	return fromChan(s.Filter(func(v T)bool {
		return s2.Has(v)
	}))
}

// Union returns a new set that holds the elements of both sets.
func (s Set[T]) Union(s2 Set[T]) Set[T] {
	return fromChan(s.Iter(), s2.Iter())
}

// Size returns how many elements are in the set.
func (s Set[T]) Size() int {
	return len(s)
}

// NewSet initializes a set of a given type.
// Note that if you pass any initial values, the type will be inferred.
func NewSet[T comparable](vs ...T) Set[T] {
	var out Set[T] = make(map[T]struct{})
	out.Add(vs...)
	return out
}

func (s Set[T]) iter(c chan T, filter func(T)bool) {
	for k := range s {
		if filter(k) {
			c <- k
		}
	}
	close(c)
}

// Iter allows you to iterate over the set.
// It is efficient even for very large sets.
func (s Set[T]) Iter() chan T {
	out := make(chan T)
	go s.iter(out, func(_ T)bool { return true })
	return out
}

// Filter is like Iter, but you only get elements that pass the filter function.
// It is efficient even for very large sets.
func (s Set[T]) Filter(filter func(T)bool) chan T {
	out := make(chan T)
	go s.iter(out, filter)
	return out
}