gitee.com/lh-her-team/common@v1.5.1/sortedmap/sorted_map.go (about)

     1  package sortedmap
     2  
     3  import (
     4  	"sort"
     5  	"sync"
     6  	"sync/atomic"
     7  )
     8  
     9  type StringKeySortedMap struct {
    10  	keysLock   sync.RWMutex
    11  	keys       []string
    12  	m          sync.Map
    13  	needResort uint32
    14  }
    15  
    16  func NewStringKeySortedMap() *StringKeySortedMap {
    17  	return &StringKeySortedMap{
    18  		keys:       make([]string, 0),
    19  		m:          sync.Map{},
    20  		needResort: 0,
    21  	}
    22  }
    23  
    24  func NewStringKeySortedMapWithData(data map[string]string) *StringKeySortedMap {
    25  	if len(data) == 0 {
    26  		return NewStringKeySortedMap()
    27  	}
    28  	sortMap := &StringKeySortedMap{
    29  		keys:       make([]string, 0),
    30  		m:          sync.Map{},
    31  		needResort: 1,
    32  	}
    33  	for key, value := range data {
    34  		sortMap.m.Store(key, value)
    35  		sortMap.keys = append(sortMap.keys, key)
    36  	}
    37  	return sortMap
    38  }
    39  
    40  func NewStringKeySortedMapWithBytesData(data map[string][]byte) *StringKeySortedMap {
    41  	if len(data) == 0 {
    42  		return NewStringKeySortedMap()
    43  	}
    44  	sortMap := &StringKeySortedMap{
    45  		keys:       make([]string, 0),
    46  		m:          sync.Map{},
    47  		needResort: 1,
    48  	}
    49  	for key, value := range data {
    50  		sortMap.m.Store(key, value)
    51  		sortMap.keys = append(sortMap.keys, key)
    52  	}
    53  	return sortMap
    54  }
    55  
    56  func NewStringKeySortedMapWithInterfaceData(data map[string]interface{}) *StringKeySortedMap {
    57  	if len(data) == 0 {
    58  		return NewStringKeySortedMap()
    59  	}
    60  	sortMap := &StringKeySortedMap{
    61  		keys:       make([]string, 0),
    62  		m:          sync.Map{},
    63  		needResort: 1,
    64  	}
    65  	for key, value := range data {
    66  		sortMap.m.Store(key, value)
    67  		sortMap.keys = append(sortMap.keys, key)
    68  	}
    69  	return sortMap
    70  }
    71  
    72  func (m *StringKeySortedMap) Put(key string, val interface{}) {
    73  	_, exist := m.m.LoadOrStore(key, val)
    74  	if exist {
    75  		m.m.Store(key, val)
    76  	} else {
    77  		m.keysLock.Lock()
    78  		m.keys = append(m.keys, key)
    79  		m.keysLock.Unlock()
    80  		atomic.CompareAndSwapUint32(&m.needResort, 0, 1)
    81  	}
    82  }
    83  
    84  func (m *StringKeySortedMap) Get(key string) (val interface{}, ok bool) {
    85  	return m.m.Load(key)
    86  }
    87  
    88  func (m *StringKeySortedMap) Contains(key string) (ok bool) {
    89  	_, ok = m.m.Load(key)
    90  	return
    91  }
    92  
    93  func (m *StringKeySortedMap) sort() {
    94  	needResort := atomic.LoadUint32(&m.needResort)
    95  	if needResort == 1 {
    96  		m.keysLock.Lock()
    97  		if atomic.LoadUint32(&m.needResort) == 1 {
    98  			sort.Strings(m.keys)
    99  			atomic.StoreUint32(&m.needResort, 0)
   100  		}
   101  		m.keysLock.Unlock()
   102  	}
   103  }
   104  
   105  func (m *StringKeySortedMap) removeFromKeys(key string) {
   106  	m.keysLock.Lock()
   107  	if atomic.LoadUint32(&m.needResort) == 1 {
   108  		sort.Strings(m.keys)
   109  		atomic.StoreUint32(&m.needResort, 0)
   110  	}
   111  	length := len(m.keys)
   112  	idx := sort.SearchStrings(m.keys, key)
   113  	if idx != length {
   114  		m.keys = append(m.keys[:idx], m.keys[idx+1:]...)
   115  	}
   116  	m.keysLock.Unlock()
   117  }
   118  
   119  func (m *StringKeySortedMap) Remove(key string) (val interface{}, ok bool) {
   120  	val, ok = m.m.LoadAndDelete(key)
   121  	if ok {
   122  		m.removeFromKeys(key)
   123  	}
   124  	return
   125  }
   126  
   127  func (m *StringKeySortedMap) Range(f func(key string, val interface{}) (isContinue bool)) {
   128  	m.sort()
   129  	m.keysLock.RLock()
   130  	keys := m.keys
   131  	length := len(keys)
   132  	if length == 0 {
   133  		m.keysLock.RUnlock()
   134  		return
   135  	}
   136  	tempKeys := make([]string, length)
   137  	copy(tempKeys, keys)
   138  	m.keysLock.RUnlock()
   139  	for i := 0; i < length; i++ {
   140  		k := tempKeys[i]
   141  		v, ok := m.m.Load(k)
   142  		if ok && !f(k, v) {
   143  			break
   144  		}
   145  	}
   146  }
   147  
   148  func (m *StringKeySortedMap) Length() int {
   149  	m.keysLock.RLock()
   150  	defer m.keysLock.RUnlock()
   151  	return len(m.keys)
   152  }
   153  
   154  type IntKeySortedMap struct {
   155  	keysLock   sync.RWMutex
   156  	keys       []int
   157  	m          sync.Map
   158  	needResort uint32
   159  }
   160  
   161  func NewIntKeySortedMap() *IntKeySortedMap {
   162  	return &IntKeySortedMap{
   163  		keys:       make([]int, 0),
   164  		m:          sync.Map{},
   165  		needResort: 0,
   166  	}
   167  }
   168  
   169  func (m *IntKeySortedMap) Put(key int, val interface{}) {
   170  	_, exist := m.m.LoadOrStore(key, val)
   171  	if exist {
   172  		m.m.Store(key, val)
   173  	} else {
   174  		m.keysLock.Lock()
   175  		m.keys = append(m.keys, key)
   176  		m.keysLock.Unlock()
   177  		atomic.CompareAndSwapUint32(&m.needResort, 0, 1)
   178  	}
   179  }
   180  
   181  func (m *IntKeySortedMap) Get(key int) (val interface{}, ok bool) {
   182  	return m.m.Load(key)
   183  }
   184  
   185  func (m *IntKeySortedMap) Contains(key int) (ok bool) {
   186  	_, ok = m.m.Load(key)
   187  	return
   188  }
   189  
   190  func (m *IntKeySortedMap) sort() {
   191  	needResort := atomic.LoadUint32(&m.needResort)
   192  	if needResort == 1 {
   193  		m.keysLock.Lock()
   194  		if atomic.LoadUint32(&m.needResort) == 1 {
   195  			sort.Ints(m.keys)
   196  			atomic.StoreUint32(&m.needResort, 0)
   197  		}
   198  		m.keysLock.Unlock()
   199  	}
   200  }
   201  
   202  func (m *IntKeySortedMap) removeFromKeys(key int) {
   203  	m.keysLock.Lock()
   204  	if atomic.LoadUint32(&m.needResort) == 1 {
   205  		sort.Ints(m.keys)
   206  		atomic.StoreUint32(&m.needResort, 0)
   207  	}
   208  	length := len(m.keys)
   209  	idx := sort.SearchInts(m.keys, key)
   210  	if idx != length {
   211  		m.keys = append(m.keys[:idx], m.keys[idx+1:]...)
   212  	}
   213  	m.keysLock.Unlock()
   214  }
   215  
   216  func (m *IntKeySortedMap) Remove(key int) (val interface{}, ok bool) {
   217  	val, ok = m.m.LoadAndDelete(key)
   218  	if ok {
   219  		m.removeFromKeys(key)
   220  	}
   221  	return
   222  }
   223  
   224  func (m *IntKeySortedMap) Range(f func(val interface{}) (isContinue bool)) {
   225  	m.sort()
   226  	m.keysLock.RLock()
   227  	keys := m.keys
   228  	length := len(keys)
   229  	if length == 0 {
   230  		m.keysLock.RUnlock()
   231  		return
   232  	}
   233  	tempKeys := make([]int, length)
   234  	copy(tempKeys, keys)
   235  	m.keysLock.RUnlock()
   236  	for i := 0; i < length; i++ {
   237  		k := tempKeys[i]
   238  		v, ok := m.m.Load(k)
   239  		if ok && !f(v) {
   240  			break
   241  		}
   242  	}
   243  }
   244  
   245  func (m *IntKeySortedMap) Length() int {
   246  	m.keysLock.RLock()
   247  	defer m.keysLock.RUnlock()
   248  	return len(m.keys)
   249  }
   250  
   251  type Float64KeySortedMap struct {
   252  	keysLock   sync.RWMutex
   253  	keys       []float64
   254  	m          sync.Map
   255  	needResort uint32
   256  }
   257  
   258  func NewFloat64KeySortedMap() *Float64KeySortedMap {
   259  	return &Float64KeySortedMap{
   260  		keys:       make([]float64, 0),
   261  		m:          sync.Map{},
   262  		needResort: 0,
   263  	}
   264  }
   265  
   266  func (m *Float64KeySortedMap) Put(key float64, val interface{}) {
   267  	_, exist := m.m.LoadOrStore(key, val)
   268  	if exist {
   269  		m.m.Store(key, val)
   270  	} else {
   271  		m.keysLock.Lock()
   272  		m.keys = append(m.keys, key)
   273  		m.keysLock.Unlock()
   274  		atomic.CompareAndSwapUint32(&m.needResort, 0, 1)
   275  	}
   276  }
   277  
   278  func (m *Float64KeySortedMap) Get(key float64) (val interface{}, ok bool) {
   279  	return m.m.Load(key)
   280  }
   281  
   282  func (m *Float64KeySortedMap) Contains(key int) (ok bool) {
   283  	_, ok = m.m.Load(key)
   284  	return
   285  }
   286  
   287  func (m *Float64KeySortedMap) sort() {
   288  	needResort := atomic.LoadUint32(&m.needResort)
   289  	if needResort == 1 {
   290  		m.keysLock.Lock()
   291  		if atomic.LoadUint32(&m.needResort) == 1 {
   292  			sort.Float64s(m.keys)
   293  			atomic.StoreUint32(&m.needResort, 0)
   294  		}
   295  		m.keysLock.Unlock()
   296  	}
   297  }
   298  
   299  func (m *Float64KeySortedMap) removeFromKeys(key float64) {
   300  	m.keysLock.Lock()
   301  	if atomic.LoadUint32(&m.needResort) == 1 {
   302  		sort.Float64s(m.keys)
   303  		atomic.StoreUint32(&m.needResort, 0)
   304  	}
   305  	length := len(m.keys)
   306  	idx := sort.SearchFloat64s(m.keys, key)
   307  	if idx != length {
   308  		m.keys = append(m.keys[:idx], m.keys[idx+1:]...)
   309  	}
   310  	m.keysLock.Unlock()
   311  }
   312  
   313  func (m *Float64KeySortedMap) Remove(key float64) (val interface{}, ok bool) {
   314  	val, ok = m.m.LoadAndDelete(key)
   315  	if ok {
   316  		m.removeFromKeys(key)
   317  	}
   318  	return
   319  }
   320  
   321  func (m *Float64KeySortedMap) Range(f func(key float64, val interface{}) (isContinue bool)) {
   322  	m.sort()
   323  	m.keysLock.RLock()
   324  	keys := m.keys
   325  	length := len(keys)
   326  	if length == 0 {
   327  		m.keysLock.RUnlock()
   328  		return
   329  	}
   330  	tempKeys := make([]float64, length)
   331  	copy(tempKeys, keys)
   332  	m.keysLock.RUnlock()
   333  	for i := 0; i < length; i++ {
   334  		k := tempKeys[i]
   335  		v, ok := m.m.Load(k)
   336  		if ok && !f(k, v) {
   337  			break
   338  		}
   339  	}
   340  }
   341  
   342  func (m *Float64KeySortedMap) Length() int {
   343  	m.keysLock.RLock()
   344  	defer m.keysLock.RUnlock()
   345  	return len(m.keys)
   346  }