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

     1  // Copyright (c) Mondoo, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  package llx
     5  
     6  import "sync"
     7  
     8  // DictTracker for tracking simple keys in a map
     9  type DictTracker struct {
    10  	mu   sync.Mutex
    11  	data map[string]struct{}
    12  }
    13  
    14  // NewDictTracker initializes a new dict tracker
    15  func NewDictTracker() *DictTracker {
    16  	return &DictTracker{
    17  		data: map[string]struct{}{},
    18  	}
    19  }
    20  
    21  // Set the key in the tracker
    22  func (c *DictTracker) Set(k string) {
    23  	c.mu.Lock()
    24  	c.data[k] = struct{}{}
    25  	c.mu.Unlock()
    26  }
    27  
    28  // Exists checks if key exists in the map
    29  func (c *DictTracker) Exists(k string) bool {
    30  	c.mu.Lock()
    31  	_, ok := c.data[k]
    32  	c.mu.Unlock()
    33  	return ok
    34  }
    35  
    36  // CheckOrSet a key, returns true if the key already exists
    37  func (c *DictTracker) CheckOrSet(k string) bool {
    38  	c.mu.Lock()
    39  	_, ok := c.data[k]
    40  	if !ok {
    41  		c.data[k] = struct{}{}
    42  	}
    43  	c.mu.Unlock()
    44  	return ok
    45  }
    46  
    47  // Unset a key in the tracker
    48  func (c *DictTracker) Unset(k string) {
    49  	c.mu.Lock()
    50  	delete(c.data, k)
    51  	c.mu.Unlock()
    52  }
    53  
    54  // GetAndClear retrieves all keys and clears out the dict
    55  func (c *DictTracker) GetAndClear() []string {
    56  	c.mu.Lock()
    57  	res := make([]string, len(c.data))
    58  	var i int
    59  	for k := range c.data {
    60  		res[i] = k
    61  		i++
    62  	}
    63  	c.data = map[string]struct{}{}
    64  	c.mu.Unlock()
    65  
    66  	return res
    67  }
    68  
    69  type DictGroupTracker struct {
    70  	mutex sync.Mutex
    71  	data  map[string]map[string]struct{}
    72  }
    73  
    74  func (g *DictGroupTracker) Clear() {
    75  	g.mutex.Lock()
    76  	g.data = map[string]map[string]struct{}{}
    77  	g.mutex.Unlock()
    78  }
    79  
    80  func (g *DictGroupTracker) ClearGroup(name string) {
    81  	g.mutex.Lock()
    82  	g.data[name] = map[string]struct{}{}
    83  	g.mutex.Unlock()
    84  }
    85  
    86  // CheckOrSet a key, returns true if the key already exists
    87  func (g *DictGroupTracker) CheckOrSet(group string, key string) bool {
    88  	g.mutex.Lock()
    89  	grp, ok := g.data[group]
    90  	if !ok {
    91  		grp = map[string]struct{}{}
    92  		g.data[group] = grp
    93  	}
    94  
    95  	_, ok = grp[key]
    96  	if !ok {
    97  		grp[key] = struct{}{}
    98  	}
    99  	g.mutex.Unlock()
   100  
   101  	return ok
   102  }