add GetOrSet() functionality

e2e_test.go

@@ -137,6 +137,24 @@ func TestDelete2(t *testing.T) {
 	})
 }
 
+// from https://pkg.go.dev/sync#Map.LoadOrStore
+func TestGetOrSet(t *testing.T) {
+	var (
+		m    = New[int, string]()
+		data = "one"
+	)
+	if val, loaded := m.GetOrSet(1, data); loaded {
+		t.Error("Value should not have been present")
+	} else if val != data {
+		t.Error("Returned value should be the same as given value if absent")
+	}
+	if val, loaded := m.GetOrSet(1, data); !loaded {
+		t.Error("Value should have been present")
+	} else if val != data {
+		t.Error("Returned value should be the same as given value")
+	}
+}
+
 func TestIterator(t *testing.T) {
 	m := New[int, *Animal]()
 

map.go

@@ -164,6 +164,58 @@ func (m *Map[K, V]) Set(key K, value V) {
 	}
 }
 
+// GetOrSet returns the existing value for the key if present
+// Otherwise, it stores and returns the given value
+// The loaded result is true if the value was loaded, false if stored
+func (m *Map[K, V]) GetOrSet(key K, value V) (actual V, loaded bool) {
+	var (
+		h        = m.hasher(key)
+		data     = m.metadata.Load()
+		existing = data.indexElement(h)
+	)
+	// try to get the element if present
+	for elem := existing; elem != nil; elem = elem.nextPtr.Load() {
+		if elem.keyHash == h && elem.key == key {
+			actual, loaded = *elem.value.Load(), true
+			return
+		}
+		if elem.keyHash <= h || elem.keyHash == marked {
+			continue
+		} else {
+			break
+		}
+	}
+	// Get() failed because element is absent
+	// store the value given by user
+	actual, loaded = value, false
+
+	var (
+		alloc   *element[K, V]
+		created = false
+		valPtr  = &value
+	)
+	if existing == nil || existing.keyHash > h {
+		existing = m.listHead
+	}
+	if alloc, created = existing.inject(h, key, valPtr); alloc != nil {
+		if created {
+			m.numItems.Add(1)
+		}
+	} else {
+		for existing = m.listHead; alloc == nil; alloc, created = existing.inject(h, key, valPtr) {
+		}
+		if created {
+			m.numItems.Add(1)
+		}
+	}
+
+	count := data.addItemToIndex(alloc)
+	if resizeNeeded(uintptr(len(data.index)), count) && m.resizing.CompareAndSwap(notResizing, resizingInProgress) {
+		m.grow(0) // double in size
+	}
+	return
+}
+
 // ForEach iterates over key-value pairs and executes the lambda provided for each such pair
 // lambda must return `true` to continue iteration and `false` to break iteration
 func (m *Map[K, V]) ForEach(lambda func(K, V) bool) {