perf: preallocate in Assign (#484)

samber · Jun 30, 2024 · 57c2d6d · 57c2d6d
1 parent 3b29b7d
commit 57c2d6d
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 2 deletions.
diff --git a/map.go b/map.go
@@ -171,8 +171,12 @@ func Invert[K comparable, V comparable](in map[K]V) map[V]K {
 // Assign merges multiple maps from left to right.
 // Play: https://go.dev/play/p/VhwfJOyxf5o
 func Assign[K comparable, V any, Map ~map[K]V](maps ...Map) Map {
-	out := Map{}
+	count := 0
+	for i := range maps {
+		count += len(maps[i])
+	}
 
+	out := make(Map, count)
 	for i := range maps {
 		for k := range maps[i] {
 			out[k] = maps[i][k]

diff --git a/map_test.go b/map_test.go
@@ -359,7 +359,8 @@ func TestMapEntries(t *testing.T) {
 		}{"1-11-1": {name: "foo", age: 1}, "2-22-2": {name: "bar", age: 2}}, func(k string, v struct {
 			name string
 			age  int
-		}) (string, string) {
+		},
+		) (string, string) {
 			return v.name, k
 		}, map[string]string{"bar": "2-22-2", "foo": "1-11-1"})
 	}
@@ -381,3 +382,67 @@ func TestMapToSlice(t *testing.T) {
 	is.ElementsMatch(result1, []string{"1_5", "2_6", "3_7", "4_8"})
 	is.ElementsMatch(result2, []string{"1", "2", "3", "4"})
 }
+
+func BenchmarkAssign(b *testing.B) {
+	counts := []int{32768, 1024, 128, 32, 2}
+
+	allDifferentMap := func(b *testing.B, n int) []map[string]int {
+		defer b.ResetTimer()
+		m := make([]map[string]int, 0)
+		for i := 0; i < n; i++ {
+			m = append(m, map[string]int{
+				strconv.Itoa(i): i,
+				strconv.Itoa(i): i,
+				strconv.Itoa(i): i,
+				strconv.Itoa(i): i,
+				strconv.Itoa(i): i,
+				strconv.Itoa(i): i,
+			},
+			)
+		}
+		return m
+	}
+
+	allTheSameMap := func(b *testing.B, n int) []map[string]int {
+		defer b.ResetTimer()
+		m := make([]map[string]int, 0)
+		for i := 0; i < n; i++ {
+			m = append(m, map[string]int{
+				"a": 1,
+				"b": 2,
+				"c": 3,
+				"d": 4,
+				"e": 5,
+				"f": 6,
+			},
+			)
+		}
+		return m
+	}
+
+	for _, count := range counts {
+		differentMap := allDifferentMap(b, count)
+		sameMap := allTheSameMap(b, count)
+
+		b.Run(fmt.Sprintf("%d", count), func(b *testing.B) {
+			testcase := []struct {
+				name string
+				maps []map[string]int
+			}{
+				{"different", differentMap},
+				{"same", sameMap},
+			}
+
+			for _, tc := range testcase {
+				b.Run(tc.name, func(b *testing.B) {
+					b.ResetTimer()
+					for n := 0; n < b.N; n++ {
+						result := Assign(tc.maps...)
+						_ = result
+					}
+				})
+			}
+		})
+
+	}
+}