set.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245
  1. package collection
  2. import (
  3. "github.com/tal-tech/go-zero/core/lang"
  4. "github.com/tal-tech/go-zero/core/logx"
  5. )
  6. const (
  7. unmanaged = iota
  8. untyped
  9. intType
  10. int64Type
  11. uintType
  12. uint64Type
  13. stringType
  14. )
  15. // Set is not thread-safe, for concurrent use, make sure to use it with synchronization.
  16. type Set struct {
  17. data map[interface{}]lang.PlaceholderType
  18. tp int
  19. }
  20. // NewSet returns a managed Set, can only put the values with the same type.
  21. func NewSet() *Set {
  22. return &Set{
  23. data: make(map[interface{}]lang.PlaceholderType),
  24. tp: untyped,
  25. }
  26. }
  27. // NewUnmanagedSet returns a unmanaged Set, which can put values with different types.
  28. func NewUnmanagedSet() *Set {
  29. return &Set{
  30. data: make(map[interface{}]lang.PlaceholderType),
  31. tp: unmanaged,
  32. }
  33. }
  34. // Add adds i into s.
  35. func (s *Set) Add(i ...interface{}) {
  36. for _, each := range i {
  37. s.add(each)
  38. }
  39. }
  40. // AddInt adds int values ii into s.
  41. func (s *Set) AddInt(ii ...int) {
  42. for _, each := range ii {
  43. s.add(each)
  44. }
  45. }
  46. // AddInt64 adds int64 values ii into s.
  47. func (s *Set) AddInt64(ii ...int64) {
  48. for _, each := range ii {
  49. s.add(each)
  50. }
  51. }
  52. // AddUint adds uint values ii into s.
  53. func (s *Set) AddUint(ii ...uint) {
  54. for _, each := range ii {
  55. s.add(each)
  56. }
  57. }
  58. // AddUint64 adds uint64 values ii into s.
  59. func (s *Set) AddUint64(ii ...uint64) {
  60. for _, each := range ii {
  61. s.add(each)
  62. }
  63. }
  64. // AddStr adds string values ss into s.
  65. func (s *Set) AddStr(ss ...string) {
  66. for _, each := range ss {
  67. s.add(each)
  68. }
  69. }
  70. // Contains checks if i is in s.
  71. func (s *Set) Contains(i interface{}) bool {
  72. if len(s.data) == 0 {
  73. return false
  74. }
  75. s.validate(i)
  76. _, ok := s.data[i]
  77. return ok
  78. }
  79. // Keys returns the keys in s.
  80. func (s *Set) Keys() []interface{} {
  81. var keys []interface{}
  82. for key := range s.data {
  83. keys = append(keys, key)
  84. }
  85. return keys
  86. }
  87. // KeysInt returns the int keys in s.
  88. func (s *Set) KeysInt() []int {
  89. var keys []int
  90. for key := range s.data {
  91. if intKey, ok := key.(int); !ok {
  92. continue
  93. } else {
  94. keys = append(keys, intKey)
  95. }
  96. }
  97. return keys
  98. }
  99. // KeysInt64 returns int64 keys in s.
  100. func (s *Set) KeysInt64() []int64 {
  101. var keys []int64
  102. for key := range s.data {
  103. if intKey, ok := key.(int64); !ok {
  104. continue
  105. } else {
  106. keys = append(keys, intKey)
  107. }
  108. }
  109. return keys
  110. }
  111. // KeysUint returns uint keys in s.
  112. func (s *Set) KeysUint() []uint {
  113. var keys []uint
  114. for key := range s.data {
  115. if intKey, ok := key.(uint); !ok {
  116. continue
  117. } else {
  118. keys = append(keys, intKey)
  119. }
  120. }
  121. return keys
  122. }
  123. // KeysUint64 returns uint64 keys in s.
  124. func (s *Set) KeysUint64() []uint64 {
  125. var keys []uint64
  126. for key := range s.data {
  127. if intKey, ok := key.(uint64); !ok {
  128. continue
  129. } else {
  130. keys = append(keys, intKey)
  131. }
  132. }
  133. return keys
  134. }
  135. // KeysStr returns string keys in s.
  136. func (s *Set) KeysStr() []string {
  137. var keys []string
  138. for key := range s.data {
  139. if strKey, ok := key.(string); !ok {
  140. continue
  141. } else {
  142. keys = append(keys, strKey)
  143. }
  144. }
  145. return keys
  146. }
  147. // Remove removes i from s.
  148. func (s *Set) Remove(i interface{}) {
  149. s.validate(i)
  150. delete(s.data, i)
  151. }
  152. // Count returns the number of items in s.
  153. func (s *Set) Count() int {
  154. return len(s.data)
  155. }
  156. func (s *Set) add(i interface{}) {
  157. switch s.tp {
  158. case unmanaged:
  159. // do nothing
  160. case untyped:
  161. s.setType(i)
  162. default:
  163. s.validate(i)
  164. }
  165. s.data[i] = lang.Placeholder
  166. }
  167. func (s *Set) setType(i interface{}) {
  168. // s.tp can only be untyped here
  169. switch i.(type) {
  170. case int:
  171. s.tp = intType
  172. case int64:
  173. s.tp = int64Type
  174. case uint:
  175. s.tp = uintType
  176. case uint64:
  177. s.tp = uint64Type
  178. case string:
  179. s.tp = stringType
  180. }
  181. }
  182. func (s *Set) validate(i interface{}) {
  183. if s.tp == unmanaged {
  184. return
  185. }
  186. switch i.(type) {
  187. case int:
  188. if s.tp != intType {
  189. logx.Errorf("Error: element is int, but set contains elements with type %d", s.tp)
  190. }
  191. case int64:
  192. if s.tp != int64Type {
  193. logx.Errorf("Error: element is int64, but set contains elements with type %d", s.tp)
  194. }
  195. case uint:
  196. if s.tp != uintType {
  197. logx.Errorf("Error: element is uint, but set contains elements with type %d", s.tp)
  198. }
  199. case uint64:
  200. if s.tp != uint64Type {
  201. logx.Errorf("Error: element is uint64, but set contains elements with type %d", s.tp)
  202. }
  203. case string:
  204. if s.tp != stringType {
  205. logx.Errorf("Error: element is string, but set contains elements with type %d", s.tp)
  206. }
  207. }
  208. }