github.com/alimy/mir/v4@v4.1.0/internal/utils/set.go (about)

     1  // Copyright 2022 Michael Li <alimy@gility.net>. All rights reserved.
     2  // Use of this source code is governed by Apache License 2.0 that
     3  // can be found in the LICENSE file.
     4  
     5  package utils
     6  
     7  import (
     8  	"errors"
     9  	"sync"
    10  )
    11  
    12  var (
    13  	errExist = errors.New("element had exist")
    14  )
    15  
    16  // OnceFunc used after OnceSet add it
    17  type OnceFunc func(string) error
    18  
    19  // Set once set
    20  type Set interface {
    21  	Add(string) error
    22  	Exist(string) bool
    23  	List() []string
    24  }
    25  
    26  type onceSet struct {
    27  	inSet    map[string]struct{}
    28  	onceFunc OnceFunc
    29  }
    30  
    31  type muxSet struct {
    32  	mu    *sync.RWMutex
    33  	inSet map[string]struct{}
    34  }
    35  
    36  type strSet map[string]struct{}
    37  
    38  // Add add a item to set
    39  func (s *onceSet) Add(it string) error {
    40  	if _, exist := s.inSet[it]; !exist {
    41  		err := s.onceFunc(it)
    42  		if err != nil {
    43  			return err
    44  		}
    45  		s.inSet[it] = struct{}{}
    46  	}
    47  	return nil
    48  }
    49  
    50  // Exist whether it exist
    51  func (s *onceSet) Exist(it string) bool {
    52  	_, exist := s.inSet[it]
    53  	return exist
    54  }
    55  
    56  // List return all items in sets
    57  func (s *onceSet) List() []string {
    58  	res := make([]string, 0, len(s.inSet))
    59  	for item := range s.inSet {
    60  		res = append(res, item)
    61  	}
    62  	return res
    63  }
    64  
    65  // Add add a item to set
    66  func (s *muxSet) Add(it string) error {
    67  	s.mu.Lock()
    68  	defer s.mu.Unlock()
    69  
    70  	if _, exist := s.inSet[it]; exist {
    71  		return errExist
    72  	}
    73  	s.inSet[it] = struct{}{}
    74  	return nil
    75  }
    76  
    77  // Exist whether it exist
    78  func (s *muxSet) Exist(it string) bool {
    79  	s.mu.RLock()
    80  	defer s.mu.RUnlock()
    81  
    82  	_, exist := s.inSet[it]
    83  	return exist
    84  }
    85  
    86  // List return all items in sets
    87  func (s *muxSet) List() []string {
    88  	s.mu.RLock()
    89  	defer s.mu.RUnlock()
    90  
    91  	res := make([]string, 0, len(s.inSet))
    92  	for item := range s.inSet {
    93  		res = append(res, item)
    94  	}
    95  	return res
    96  }
    97  
    98  // Add add a item to set
    99  func (s strSet) Add(item string) error {
   100  	if _, exist := s[item]; exist {
   101  		return errExist
   102  	}
   103  	s[item] = struct{}{}
   104  	return nil
   105  }
   106  
   107  // Exist whether it exist
   108  func (s strSet) Exist(item string) bool {
   109  	_, exist := s[item]
   110  	return exist
   111  }
   112  
   113  // List return all items in sets
   114  func (s strSet) List() []string {
   115  	methods := make([]string, 0, len(s))
   116  	for m := range s {
   117  		methods = append(methods, m)
   118  	}
   119  	return methods
   120  }
   121  
   122  // NewOnceSet return an OnceSet instance
   123  // Warning: OnceSet is not goroutine safe
   124  func NewOnceSet(onceFunc OnceFunc) Set {
   125  	return &onceSet{
   126  		inSet:    make(map[string]struct{}),
   127  		onceFunc: onceFunc,
   128  	}
   129  }
   130  
   131  // NewMuxSet return a goroutine safe set
   132  func NewMuxSet() Set {
   133  	return &muxSet{
   134  		mu:    &sync.RWMutex{},
   135  		inSet: make(map[string]struct{}),
   136  	}
   137  }
   138  
   139  // NewStrSet return a str sets
   140  func NewStrSet() Set {
   141  	return make(strSet)
   142  }