From fa093d9b09e32a697783020e9d84c76932b1679d Mon Sep 17 00:00:00 2001
From: gingerBill <bill@gingerbill.org>
Date: Tue, 6 Feb 2024 23:58:22 +0000
Subject: [PATCH] Add `runtime.map_insert_and_check_for_previous`

---
 base/runtime/core_builtin.odin         | 14 +++++++
 base/runtime/dynamic_map_internal.odin | 51 ++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/base/runtime/core_builtin.odin b/base/runtime/core_builtin.odin
index 5c408e57a..82b3eb0de 100644
--- a/base/runtime/core_builtin.odin
+++ b/base/runtime/core_builtin.odin
@@ -824,6 +824,20 @@ map_insert :: proc(m: ^$T/map[$K]$V, key: K, value: V, loc := #caller_location)
 	return (^V)(__dynamic_map_set_without_hash((^Raw_Map)(m), map_info(T), rawptr(&key), rawptr(&value), loc))
 }
 
+// Explicitly inserts a key and value into a map `m`, the same as `map_insert`, but the return values differ.
+// - `prev_key_ptr` will return the previous pointer of a key if it exists, and `nil` otherwise.
+// - `value_ptr` will return the pointer of the memory where the insertion happens, and `nil` if the map failed to resize
+// - `found_previous` will be true if `prev_key_ptr != nil`
+@(require_results)
+map_insert_and_check_for_previous :: proc(m: ^$T/map[$K]$V, key: K, value: V, loc := #caller_location) -> (prev_key_ptr: ^K, value_ptr: ^V, found_previous: bool) {
+	key, value := key, value
+	kp, vp := __dynamic_map_set_extra_without_hash((^Raw_Map)(m), map_info(T), rawptr(&key), rawptr(&value), loc)
+	prev_key_ptr   = (^K)(kp)
+	value_ptr      = (^V)(vp)
+	found_previous = kp != nil
+	return
+}
+
 
 @builtin
 card :: proc "contextless" (s: $S/bit_set[$E; $U]) -> int {
diff --git a/base/runtime/dynamic_map_internal.odin b/base/runtime/dynamic_map_internal.odin
index 64cb02586..6955f4a1e 100644
--- a/base/runtime/dynamic_map_internal.odin
+++ b/base/runtime/dynamic_map_internal.odin
@@ -841,6 +841,33 @@ __dynamic_map_get :: proc "contextless" (#no_alias m: ^Raw_Map, #no_alias info:
 	}
 }
 
+__dynamic_map_get_key_and_value :: proc "contextless" (#no_alias m: ^Raw_Map, #no_alias info: ^Map_Info, h: Map_Hash, key: rawptr) -> (key_ptr, value_ptr: rawptr) {
+	if m.len == 0 {
+		return nil, nil
+	}
+	pos := map_desired_position(m^, h)
+	distance := uintptr(0)
+	mask := (uintptr(1) << map_log2_cap(m^)) - 1
+	ks, vs, hs, _, _ := map_kvh_data_dynamic(m^, info)
+	for {
+		element_hash := hs[pos]
+		if map_hash_is_empty(element_hash) {
+			return nil, nil
+		} else if distance > map_probe_distance(m^, element_hash, pos) {
+			return nil, nil
+		} else if element_hash == h {
+			other_key := rawptr(map_cell_index_dynamic(ks, info.ks, pos))
+			if info.key_equal(key, other_key) {
+				key_ptr   = other_key
+				value_ptr = rawptr(map_cell_index_dynamic(vs, info.vs, pos))
+				return
+			}
+		}
+		pos = (pos + 1) & mask
+		distance += 1
+	}
+}
+
 // IMPORTANT: USED WITHIN THE COMPILER
 __dynamic_map_check_grow :: proc "odin" (#no_alias m: ^Raw_Map, #no_alias info: ^Map_Info, loc := #caller_location) -> (err: Allocator_Error, has_grown: bool) {
 	if m.len >= map_resize_threshold(m^) {
@@ -874,6 +901,30 @@ __dynamic_map_set :: proc "odin" (#no_alias m: ^Raw_Map, #no_alias info: ^Map_In
 	m.len += 1
 	return rawptr(result)
 }
+__dynamic_map_set_extra_without_hash :: proc "odin" (#no_alias m: ^Raw_Map, #no_alias info: ^Map_Info, key, value: rawptr, loc := #caller_location) -> (prev_key_ptr, value_ptr: rawptr) {
+	return __dynamic_map_set_extra(m, info, info.key_hasher(key, map_seed(m^)), key, value, loc)
+}
+
+__dynamic_map_set_extra :: proc "odin" (#no_alias m: ^Raw_Map, #no_alias info: ^Map_Info, hash: Map_Hash, key, value: rawptr, loc := #caller_location) -> (prev_key_ptr, value_ptr: rawptr) {
+	if prev_key_ptr, value_ptr = __dynamic_map_get_key_and_value(m, info, hash, key); value_ptr != nil {
+		intrinsics.mem_copy_non_overlapping(value_ptr, value, info.vs.size_of_type)
+		return
+	}
+
+	hash := hash
+	err, has_grown := __dynamic_map_check_grow(m, info, loc)
+	if err != nil {
+		return nil, nil
+	}
+	if has_grown {
+		hash = info.key_hasher(key, map_seed(m^))
+	}
+
+	result := map_insert_hash_dynamic(m, info, hash, uintptr(key), uintptr(value))
+	m.len += 1
+	return nil, rawptr(result)
+}
+
 
 // IMPORTANT: USED WITHIN THE COMPILER
 @(private)