github.com/coreservice-io/utils@v0.3.0/randmap/rand_map.go (about)

     1  package randmap
     2  
     3  import (
     4  	"math/rand"
     5  	"sync"
     6  )
     7  
     8  type element struct {
     9  	content     interface{}
    10  	slice_index int
    11  }
    12  
    13  type RandMap struct {
    14  	m sync.RWMutex
    15  	// Where the objects you care about are stored.
    16  	container map[any]element
    17  	// A slice of the map keys used in the map above. We put them in a slice
    18  	// so that we can get a random key by choosing a random index.
    19  	keys []any
    20  	// We store the index of each key, so that when we remove an item, we can
    21  	// quickly remove it from the slice above.
    22  	sliceKeyIndex map[any]int
    23  	//len of keys
    24  	counter int
    25  }
    26  
    27  func NewRandMap() *RandMap {
    28  	return &RandMap{
    29  		container:     make(map[any]element),
    30  		sliceKeyIndex: make(map[any]int),
    31  		counter:       0,
    32  	}
    33  }
    34  
    35  func (s *RandMap) Count() int {
    36  	return s.counter
    37  }
    38  
    39  func (s *RandMap) Set(key any, item interface{}) {
    40  	s.m.Lock()
    41  	defer s.m.Unlock()
    42  
    43  	if old_ele, ok := s.container[key]; ok {
    44  		//old exist already
    45  		s.container[key] = element{item, old_ele.slice_index}
    46  	} else {
    47  		// add map key to slice of map keys
    48  		s.keys = append(s.keys, key)
    49  		// store object in map
    50  		s.container[key] = element{item, s.counter}
    51  		s.sliceKeyIndex[key] = s.counter
    52  		s.counter++
    53  	}
    54  }
    55  
    56  func (s *RandMap) Get(key any) interface{} {
    57  	s.m.RLock()
    58  	defer s.m.RUnlock()
    59  
    60  	if ele, ok := s.container[key]; ok {
    61  		return ele.content
    62  	} else {
    63  		return nil
    64  	}
    65  
    66  }
    67  
    68  func (s *RandMap) Remove(key any) {
    69  	s.m.Lock()
    70  	defer s.m.Unlock()
    71  	// get index in key slice for key
    72  	index, exists := s.sliceKeyIndex[key]
    73  	if !exists {
    74  		// item does not exist
    75  		return
    76  	}
    77  	delete(s.sliceKeyIndex, key)
    78  
    79  	counter_prev := s.counter - 1
    80  
    81  	// remove key from slice of keys
    82  	s.keys[index] = s.keys[counter_prev]
    83  	s.keys = s.keys[:counter_prev]
    84  
    85  	// we just swapped the last element to another position.
    86  	// so we need to update it's index (if it was not in last position)
    87  	if counter_prev != index { //not the last index
    88  		otherKey := s.keys[index]
    89  		s.sliceKeyIndex[otherKey] = index
    90  	}
    91  
    92  	// remove object from map
    93  	delete(s.container, key)
    94  
    95  	s.counter--
    96  }
    97  
    98  func (s *RandMap) Random() interface{} {
    99  
   100  	if s.counter <= 0 {
   101  		return nil
   102  	}
   103  
   104  	s.m.RLock()
   105  	defer s.m.RUnlock()
   106  
   107  	randomIndex := rand.Intn(s.counter)
   108  	key := s.keys[randomIndex]
   109  
   110  	if ele, ok := s.container[key]; ok {
   111  		return ele.content
   112  	} else {
   113  		return nil
   114  	}
   115  }
   116  
   117  func (s *RandMap) PopRandom() interface{} {
   118  
   119  	if s.counter <= 0 {
   120  		return nil
   121  	}
   122  
   123  	s.m.RLock()
   124  	randomIndex := rand.Intn(s.counter)
   125  	key := s.keys[randomIndex]
   126  
   127  	item := s.container[key]
   128  	s.m.RUnlock()
   129  
   130  	s.Remove(key)
   131  
   132  	return item.content
   133  }
   134  
   135  func (s *RandMap) Loop(callback func(key,value interface{}) bool) {
   136  	s.m.Lock()
   137  	defer s.m.Unlock()
   138  
   139  	for k,v:=range s.container{
   140  		if !callback(k,v.content){
   141  			return
   142  		}
   143  	}
   144  }