fix N=1 and cleanup tests

2 files changed, 67 insertions, 26 deletions

diff --git a/base/runtime/dynamic_map_internal.odin b/base/runtime/dynamic_map_internal.odin
index 281d4b88e..4e22aa25c 100644
--- a/base/runtime/dynamic_map_internal.odin
+++ b/base/runtime/dynamic_map_internal.odin
@@ -158,6 +158,11 @@ map_cell_index_static :: #force_inline proc "contextless" (cells: [^]Map_Cell($T
 	} else when (N & (N - 1)) == 0 && N <= 8*size_of(uintptr) {
 		// Likely case, N is a power of two because T is a power of two.
 
+		// Unique case, no need to index data here since only one element.
+		when N == 1 {
+			return &cells[index].data[0]
+		}
+
 		// Compute the integer log 2 of N, this is the shift amount to index the
 		// correct cell. Odin's intrinsics.count_leading_zeros does not produce a
 		// constant, hence this approach. We only need to check up to N = 64.
@@ -167,12 +172,7 @@ map_cell_index_static :: #force_inline proc "contextless" (cells: [^]Map_Cell($T
 		         4 when N == 16 else
 		         5 when N == 32 else 6
 		#assert(SHIFT <= MAP_CACHE_LINE_LOG2)
-		// Unique case, no need to index data here since only one element.
-		when N == 1 {
-			return &cells[index >> SHIFT].data[0]
-		} else {
-			return &cells[index >> SHIFT].data[index & (N - 1)]
-		}
+		return &cells[index >> SHIFT].data[index & (N - 1)]
 	} else {
 		// Least likely (and worst case), we pay for a division operation but we
 		// assume the compiler does not actually generate a division. N will be in the
diff --git a/tests/core/runtime/test_core_runtime.odin b/tests/core/runtime/test_core_runtime.odin
index ccadcec27..d2ced2504 100644
--- a/tests/core/runtime/test_core_runtime.odin
+++ b/tests/core/runtime/test_core_runtime.odin
@@ -67,25 +67,66 @@ test_init_cap_map_dynarray :: proc(t: ^testing.T) {
 
 @(test)
 test_map_get :: proc(t: ^testing.T) {
-	m := map[int][3]int{
-		1 = {10, 100, 1000},
-		2 = {20, 200, 2000},
-		3 = {30, 300, 3000},
+	check :: proc(t: ^testing.T, m: map[$K]$V, loc := #caller_location) {
+		for k, v in m {
+			got_key, got_val, ok := runtime.map_get(m, k)
+			testing.expect_value(t, got_key, k, loc = loc)
+			testing.expect_value(t, got_val, v, loc = loc)
+			testing.expect(t, ok, loc = loc)
+		}
+	}
+
+	// small keys & values
+	{
+		m := map[int]int{
+			1 = 10,
+			2 = 20,
+			3 = 30,
+		}
+		defer delete(m)
+		check(t, m)
+	}
+
+	// small keys; 2 values per cell
+	{
+		m := map[int][3]int{
+			1 = [3]int{10, 100, 1000},
+			2 = [3]int{20, 200, 2000},
+			3 = [3]int{30, 300, 3000},
+		}
+		defer delete(m)
+		check(t, m)
+	}
+
+	// 2 keys per cell; small values
+	{
+		m := map[[3]int]int{
+			[3]int{10, 100, 1000} = 1,
+			[3]int{20, 200, 2000} = 2,
+			[3]int{30, 300, 3000} = 3,
+		}
+		defer delete(m)
+		check(t, m)
+	}
+
+	// small keys; value bigger than a chacheline
+	{
+		m := map[int][9]int{
+			1 = [9]int{10, 100, 1000, 10000, 100000, 1000000, 10000000, 100000000, 1000000000},
+			2 = [9]int{20, 200, 2000, 20000, 200000, 2000000, 20000000, 200000000, 2000000000},
+			3 = [9]int{30, 300, 3000, 30000, 300000, 3000000, 30000000, 300000000, 3000000000},
+		}
+		defer delete(m)
+		check(t, m)
+	}
+	// keys bigger than a chacheline; small values
+	{
+		m := map[[9]int]int{
+			[9]int{10, 100, 1000, 10000, 100000, 1000000, 10000000, 100000000, 1000000000} = 1,
+			[9]int{20, 200, 2000, 20000, 200000, 2000000, 20000000, 200000000, 2000000000} = 2,
+			[9]int{30, 300, 3000, 30000, 300000, 3000000, 30000000, 300000000, 3000000000} = 3,
+		}
+		defer delete(m)
+		check(t, m)
 	}
-	defer delete(m)
-
-	k1, v1, ok1 := runtime.map_get(m, 1)
-	testing.expect_value(t, k1, 1)
-	testing.expect_value(t, v1, [3]int{10, 100, 1000})
-	testing.expect_value(t, ok1, true)
-
-	k2, v2, ok2 := runtime.map_get(m, 2)
-	testing.expect_value(t, k2, 2)
-	testing.expect_value(t, v2, [3]int{20, 200, 2000})
-	testing.expect_value(t, ok2, true)
-
-	k3, v3, ok3 := runtime.map_get(m, 3)
-	testing.expect_value(t, k3, 3)
-	testing.expect_value(t, v3, [3]int{30, 300, 3000})
-	testing.expect_value(t, ok3, true)
 }