slice.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. // Copyright (c) Tailscale Inc & AUTHORS
  2. // SPDX-License-Identifier: BSD-3-Clause
  3. package set
  4. import (
  5. "golang.org/x/exp/slices"
  6. "tailscale.com/types/views"
  7. )
  8. // Slice is a set of elements tracked in a slice of unique elements.
  9. type Slice[T comparable] struct {
  10. slice []T
  11. set map[T]bool // nil until/unless slice is large enough
  12. }
  13. // Slice returns the a view of the underlying slice.
  14. // The elements are in order of insertion.
  15. // The returned value is only valid until ss is modified again.
  16. func (ss *Slice[T]) Slice() views.Slice[T] { return views.SliceOf(ss.slice) }
  17. // Contains reports whether v is in the set.
  18. // The amortized cost is O(1).
  19. func (ss *Slice[T]) Contains(v T) bool {
  20. if ss.set != nil {
  21. return ss.set[v]
  22. }
  23. return slices.Index(ss.slice, v) != -1
  24. }
  25. // Remove removes v from the set.
  26. // The cost is O(n).
  27. func (ss *Slice[T]) Remove(v T) {
  28. if ss.set != nil {
  29. if !ss.set[v] {
  30. return
  31. }
  32. delete(ss.set, v)
  33. }
  34. if ix := slices.Index(ss.slice, v); ix != -1 {
  35. ss.slice = append(ss.slice[:ix], ss.slice[ix+1:]...)
  36. }
  37. }
  38. // Add adds each element in vs to the set.
  39. // The amortized cost is O(1) per element.
  40. func (ss *Slice[T]) Add(vs ...T) {
  41. for _, v := range vs {
  42. if ss.Contains(v) {
  43. continue
  44. }
  45. ss.slice = append(ss.slice, v)
  46. if ss.set != nil {
  47. ss.set[v] = true
  48. } else if len(ss.slice) > 8 {
  49. ss.set = make(map[T]bool, len(ss.slice))
  50. for _, v := range ss.slice {
  51. ss.set[v] = true
  52. }
  53. }
  54. }
  55. }
  56. // AddSlice adds all elements in vs to the set.
  57. func (ss *Slice[T]) AddSlice(vs views.Slice[T]) {
  58. for i := 0; i < vs.Len(); i++ {
  59. ss.Add(vs.At(i))
  60. }
  61. }