go.mondoo.com/cnquery@v0.0.0-20231005093811-59568235f6ea/types/maps.go (about)

     1  // Copyright (c) Mondoo, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  package types
     5  
     6  import "sync"
     7  
     8  // StringSet is a map that contains unique strings
     9  type StringSet struct{ sync.Map }
    10  
    11  // Store a string to the set
    12  func (c *StringSet) Store(k string) { c.Map.Store(k, struct{}{}) }
    13  
    14  // Delete a key from the set
    15  func (c *StringSet) Delete(k string) { c.Map.Delete(k) }
    16  
    17  // Exist returns true if the string exists in the set
    18  func (c *StringSet) Exist(k string) bool {
    19  	_, ok := c.Map.Load(k)
    20  	return ok
    21  }
    22  
    23  // Range walks the list of keys in the set and executes the function
    24  // as long as it returns true
    25  func (c *StringSet) Range(f func(string) bool) {
    26  	c.Map.Range(func(key, value interface{}) bool {
    27  		return f(key.(string))
    28  	})
    29  }
    30  
    31  // List all keys
    32  func (c *StringSet) List() []string {
    33  	res := []string{}
    34  	c.Map.Range(func(key, value interface{}) bool {
    35  		res = append(res, key.(string))
    36  		return true
    37  	})
    38  	return res
    39  }
    40  
    41  // StringToStrings is a map that contains a list of strings for every string stored
    42  type StringToStrings struct{ sync.Map }
    43  
    44  // Store a string association to the set
    45  func (s *StringToStrings) Store(key string, value string) {
    46  	v, ok := s.Map.Load(key)
    47  	var list *StringSet
    48  	if !ok {
    49  		list = &StringSet{}
    50  		s.Map.Store(key, list)
    51  	} else {
    52  		list = v.(*StringSet)
    53  	}
    54  	list.Store(value)
    55  }
    56  
    57  // Exist a key-value connection
    58  func (s *StringToStrings) Exist(key string, value string) bool {
    59  	v, ok := s.Map.Load(key)
    60  	if !ok {
    61  		return false
    62  	}
    63  	return v.(*StringSet).Exist(value)
    64  }
    65  
    66  // List all keys and their associations
    67  func (s *StringToStrings) List() map[string][]string {
    68  	res := make(map[string][]string)
    69  	s.Map.Range(func(key, value interface{}) bool {
    70  		res[key.(string)] = value.(*StringSet).List()
    71  		return true
    72  	})
    73  	return res
    74  }
    75  
    76  // ListKey in the set
    77  func (s *StringToStrings) ListKey(key string) []string {
    78  	v, ok := s.Load(key)
    79  	if !ok {
    80  		return nil
    81  	}
    82  	return v.List()
    83  }
    84  
    85  // Load a key
    86  func (s *StringToStrings) Load(key string) (*StringSet, bool) {
    87  	v, ok := s.Map.Load(key)
    88  	if !ok {
    89  		return nil, false
    90  	}
    91  	return v.(*StringSet), true
    92  }
    93  
    94  // Delete a key-value connection
    95  func (s *StringToStrings) Delete(key string, value string) {
    96  	set, ok := s.Load(key)
    97  	if !ok {
    98  		return
    99  	}
   100  
   101  	set.Delete(value)
   102  
   103  	empty := true
   104  	set.Range(func(_ string) bool {
   105  		empty = false
   106  		return false
   107  	})
   108  
   109  	if empty {
   110  		s.Map.Delete(key)
   111  	}
   112  }
   113  
   114  // mapValuesToSlice returns a slice with the values of the map
   115  func MapValuesToSlice[K comparable, V any](m map[K]V) []V {
   116  	var slice []V
   117  	for _, v := range m {
   118  		slice = append(slice, v)
   119  	}
   120  	return slice
   121  }
   122  
   123  // mergeMaps merges 2 maps. If there are duplicate keys the values from m2 will override
   124  // the values from m1.
   125  func MergeMaps[K comparable, V any](m1 map[K]V, m2 map[K]V) map[K]V {
   126  	for k, v := range m2 {
   127  		m1[k] = v
   128  	}
   129  	return m1
   130  }