insert_ordered_map.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package sgutil
     6  
     7  import "iter"
     8  
     9  type InsertMap[K comparable, V any] struct {
    10  	m map[K]uint32 // Maps keys to their insertion index in the values slice.
    11  	v []V          // Stores values in their insertion order.
    12  	k []K          // Lazy slice of keys in insertion order, for iterators.
    13  }
    14  
    15  // Put inserts or updates a key-value pair in the map.
    16  // If the key already exists, its value is updated while its insertion order
    17  // remains unchanged. Returns the old value if there is one,
    18  // and a boolean indicating whether an update occurred.
    19  func (im *InsertMap[K, V]) Put(key K, val V) (old V, updated bool) {
    20  	if im.m == nil {
    21  		im.m = make(map[K]uint32)
    22  	}
    23  
    24  	var index uint32
    25  	index, updated = im.m[key]
    26  	if updated {
    27  		// Key already exists; update its value in-place in the values slice.
    28  		old = im.v[index]
    29  		im.v[index] = val
    30  		return
    31  	}
    32  
    33  	// Keep the keys slice synchronized if it has already been initialized.
    34  	if len(im.v) > 0 && len(im.k) == len(im.v) {
    35  		im.k = append(im.k, key)
    36  	}
    37  
    38  	// Record the new key's insertion index and append the value.
    39  	im.m[key] = uint32(len(im.v))
    40  	im.v = append(im.v, val)
    41  	return
    42  }
    43  
    44  // updateKeys synchronizes the lazy keys slice (im.k) with the values slice (im.v).
    45  // It populates im.k with keys at their respective insertion indices using
    46  // the mapping stored in im.m.
    47  func (im *InsertMap[K, V]) updateKeys() {
    48  	if len(im.k) == len(im.v) {
    49  		return
    50  	}
    51  	im.k = make([]K, len(im.v))
    52  	for k, i := range im.m {
    53  		im.k[i] = k
    54  	}
    55  }
    56  
    57  // Contains reports whether the map contains the specified key.
    58  func (im *InsertMap[K, V]) Contains(key K) bool {
    59  	if im.m == nil {
    60  		return false
    61  	}
    62  	_, ok := im.m[key]
    63  	return ok
    64  }
    65  
    66  // Get returns the value associated with the specified key.
    67  // If the key does not exist, the zero value of type V is returned.
    68  func (im *InsertMap[K, V]) Get(key K) V {
    69  	if im.m != nil {
    70  		index, ok := im.m[key]
    71  		if ok {
    72  			return im.v[index]
    73  		}
    74  	}
    75  	var v V
    76  	return v
    77  }
    78  
    79  // GetOk returns the value associated with the specified key and
    80  // a boolean indicating whether the key exists in the map.
    81  // If the key is not in the map, the zero value is returned along with false.
    82  func (im *InsertMap[K, V]) GetOk(key K) (V, bool) {
    83  	if im.m != nil {
    84  		index, ok := im.m[key]
    85  		if ok {
    86  			return im.v[index], true
    87  		}
    88  	}
    89  	var v V
    90  	return v, false
    91  }
    92  
    93  // Compare compares two keys based on their insertion order.
    94  // It returns:
    95  //   - -1 if 'a' was inserted before 'b'
    96  //   - 1 if 'a' was inserted after 'b'
    97  //   - 0 if 'a' and 'b' are equal, or if neither key exists in the map.
    98  //
    99  // If only one of the keys exists in the map, the existing key is considered
   100  // to be "before" (less than) the non-existing key, returning -1 if 'a' exists,
   101  // and 1 if 'b' exists.
   102  func (im *InsertMap[K, V]) Compare(a, b K) int {
   103  	if im.m == nil {
   104  		return 0
   105  	}
   106  	ai, aok := im.m[a]
   107  	bi, bok := im.m[b]
   108  	if aok != bok {
   109  		if aok {
   110  			return -1
   111  		}
   112  		return 1
   113  	}
   114  	if !aok {
   115  		return 0
   116  	}
   117  	if ai < bi {
   118  		return -1
   119  	}
   120  	if ai > bi {
   121  		return 1
   122  	}
   123  	return 0
   124  }
   125  
   126  // All returns an iterator over the key-value pairs of the map in insertion order.
   127  func (im *InsertMap[K, V]) All() iter.Seq2[K, V] {
   128  	return func(yield func(K, V) bool) {
   129  		im.updateKeys()
   130  		for i, k := range im.k {
   131  			if !yield(k, im.v[i]) {
   132  				return
   133  			}
   134  		}
   135  	}
   136  }
   137  
   138  // Keys returns an iterator over the keys of the map in insertion order.
   139  func (im *InsertMap[K, V]) Keys() iter.Seq[K] {
   140  	return func(yield func(K) bool) {
   141  		im.updateKeys()
   142  		for _, k := range im.k {
   143  			if !yield(k) {
   144  				return
   145  			}
   146  		}
   147  	}
   148  }
   149  
   150  // Values returns an iterator over the values of the map in insertion order.
   151  func (im *InsertMap[K, V]) Values() iter.Seq[V] {
   152  	return func(yield func(V) bool) {
   153  		for _, v := range im.v {
   154  			if !yield(v) {
   155  				return
   156  			}
   157  		}
   158  	}
   159  }
   160
View as plain text