setmatrix.go 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. package common
  2. import (
  3. "sync"
  4. mapset "github.com/deckarep/golang-set"
  5. )
  6. // SetMatrix is a map of Sets
  7. type SetMatrix interface {
  8. // Get returns the members of the set for a specific key as a slice.
  9. Get(key string) ([]interface{}, bool)
  10. // Contains is used to verify is an element is in a set for a specific key
  11. // returns true if the element is in the set
  12. // returns true if there is a set for the key
  13. Contains(key string, value interface{}) (bool, bool)
  14. // Insert inserts the mapping between the IP and the endpoint identifier
  15. // returns true if the mapping was not present, false otherwise
  16. // returns also the number of endpoints associated to the IP
  17. Insert(key string, value interface{}) (bool, int)
  18. // Remove removes the mapping between the IP and the endpoint identifier
  19. // returns true if the mapping was deleted, false otherwise
  20. // returns also the number of endpoints associated to the IP
  21. Remove(key string, value interface{}) (bool, int)
  22. // Cardinality returns the number of elements in the set of a specfic key
  23. // returns false if the key is not in the map
  24. Cardinality(key string) (int, bool)
  25. // String returns the string version of the set, empty otherwise
  26. // returns false if the key is not in the map
  27. String(key string) (string, bool)
  28. }
  29. type setMatrix struct {
  30. matrix map[string]mapset.Set
  31. sync.Mutex
  32. }
  33. // NewSetMatrix creates a new set matrix object
  34. func NewSetMatrix() SetMatrix {
  35. s := &setMatrix{}
  36. s.init()
  37. return s
  38. }
  39. func (s *setMatrix) init() {
  40. s.matrix = make(map[string]mapset.Set)
  41. }
  42. func (s *setMatrix) Get(key string) ([]interface{}, bool) {
  43. s.Lock()
  44. defer s.Unlock()
  45. set, ok := s.matrix[key]
  46. if !ok {
  47. return nil, ok
  48. }
  49. return set.ToSlice(), ok
  50. }
  51. func (s *setMatrix) Contains(key string, value interface{}) (bool, bool) {
  52. s.Lock()
  53. defer s.Unlock()
  54. set, ok := s.matrix[key]
  55. if !ok {
  56. return false, ok
  57. }
  58. return set.Contains(value), ok
  59. }
  60. func (s *setMatrix) Insert(key string, value interface{}) (bool, int) {
  61. s.Lock()
  62. defer s.Unlock()
  63. set, ok := s.matrix[key]
  64. if !ok {
  65. s.matrix[key] = mapset.NewSet()
  66. s.matrix[key].Add(value)
  67. return true, 1
  68. }
  69. return set.Add(value), set.Cardinality()
  70. }
  71. func (s *setMatrix) Remove(key string, value interface{}) (bool, int) {
  72. s.Lock()
  73. defer s.Unlock()
  74. set, ok := s.matrix[key]
  75. if !ok {
  76. return false, 0
  77. }
  78. var removed bool
  79. if set.Contains(value) {
  80. set.Remove(value)
  81. removed = true
  82. // If the set is empty remove it from the matrix
  83. if set.Cardinality() == 0 {
  84. delete(s.matrix, key)
  85. }
  86. }
  87. return removed, set.Cardinality()
  88. }
  89. func (s *setMatrix) Cardinality(key string) (int, bool) {
  90. s.Lock()
  91. defer s.Unlock()
  92. set, ok := s.matrix[key]
  93. if !ok {
  94. return 0, ok
  95. }
  96. return set.Cardinality(), ok
  97. }
  98. func (s *setMatrix) String(key string) (string, bool) {
  99. s.Lock()
  100. defer s.Unlock()
  101. set, ok := s.matrix[key]
  102. if !ok {
  103. return "", ok
  104. }
  105. return set.String(), ok
  106. }