set.go 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. // Copyright (c) Tailscale Inc & AUTHORS
  2. // SPDX-License-Identifier: BSD-3-Clause
  3. // Package set contains set types.
  4. package set
  5. import (
  6. "encoding/json"
  7. "maps"
  8. )
  9. // Set is a set of T.
  10. type Set[T comparable] map[T]struct{}
  11. // SetOf returns a new set constructed from the elements in slice.
  12. func SetOf[T comparable](slice []T) Set[T] {
  13. s := make(Set[T])
  14. s.AddSlice(slice)
  15. return s
  16. }
  17. // Clone returns a new set cloned from the elements in s.
  18. func (s Set[T]) Clone() Set[T] {
  19. return maps.Clone(s)
  20. }
  21. // Add adds e to s.
  22. func (s Set[T]) Add(e T) { s[e] = struct{}{} }
  23. // AddSlice adds each element of es to s.
  24. func (s Set[T]) AddSlice(es []T) {
  25. for _, e := range es {
  26. s.Add(e)
  27. }
  28. }
  29. // AddSet adds each element of es to s.
  30. func (s Set[T]) AddSet(es Set[T]) {
  31. for e := range es {
  32. s.Add(e)
  33. }
  34. }
  35. // Slice returns the elements of the set as a slice. The elements will not be
  36. // in any particular order.
  37. func (s Set[T]) Slice() []T {
  38. es := make([]T, 0, s.Len())
  39. for k := range s {
  40. es = append(es, k)
  41. }
  42. return es
  43. }
  44. // Delete removes e from the set.
  45. func (s Set[T]) Delete(e T) { delete(s, e) }
  46. // Contains reports whether s contains e.
  47. func (s Set[T]) Contains(e T) bool {
  48. _, ok := s[e]
  49. return ok
  50. }
  51. // Len reports the number of items in s.
  52. func (s Set[T]) Len() int { return len(s) }
  53. // Equal reports whether s is equal to other.
  54. func (s Set[T]) Equal(other Set[T]) bool {
  55. return maps.Equal(s, other)
  56. }
  57. func (s Set[T]) MarshalJSON() ([]byte, error) {
  58. return json.Marshal(s.Slice())
  59. }
  60. func (s *Set[T]) UnmarshalJSON(buf []byte) error {
  61. var ss []T
  62. if err := json.Unmarshal(buf, &ss); err != nil {
  63. return err
  64. }
  65. *s = SetOf(ss)
  66. return nil
  67. }