CopyMapSubset

Author: muir

map.go

@@ -26,6 +26,7 @@ func CompareKeys[K comparable, V any](a, b map[K]V) ([]K, []K) {
 
 // MissingKeys returns the keys that are in a but not b
 func MissingKeys[K comparable, V any](a, b map[K]V) []K {
+	// Pre-allocate with capacity of a since that's the maximum possible size
 	onlyA := make([]K, 0, len(a))
 	for k := range a {
 		if _, ok := b[k]; !ok {
@@ -49,19 +50,47 @@ func CopyMap[K comparable, V any](m map[K]V) map[K]V {
 	if m == nil {
 		return nil
 	}
-	newM := make(map[K]V)
+	// Pre-allocate with capacity of m
+	newM := make(map[K]V, len(m))
 	for k, v := range m {
 		newM[k] = v
 	}
 	return newM
 }
 
+// CopyMapSubset creates a new map containing only the specified keys from the original map
+func CopyMapSubset[K comparable, V any](m map[K]V, keys []K) map[K]V {
+	if m == nil {
+		return nil
+	}
+
+	if len(m) == 0 {
+		return make(map[K]V)
+	}
+
+	// Pre-allocate with capacity of keys since that's the maximum possible size
+	result := make(map[K]V, len(keys))
+	for _, k := range keys {
+		if v, ok := m[k]; ok {
+			result[k] = v
+		}
+	}
+	return result
+}
+
 // Merge copies b onto a, overriding any common keys:
 // a is modified and returned.
 func Merge[K comparable, V any](a, b map[K]V) map[K]V {
 	if a == nil {
 		return b
 	}
+	// If a is nil, create a new map with capacity of b
+	if a == nil {
+		if b == nil {
+			return nil
+		}
+		a = make(map[K]V, len(b))
+	}
 	for k, v := range b {
 		a[k] = v
 	}

map_test.go

@@ -344,6 +344,119 @@ func TestCopyMap(t *testing.T) {
 	})
 }
 
+func TestCopyMapSubset(t *testing.T) {
+	t.Parallel()
+
+	t.Run("copies specified keys", func(t *testing.T) {
+		t.Parallel()
+
+		original := map[string]int{
+			"a": 1,
+			"b": 2,
+			"c": 3,
+			"d": 4,
+		}
+
+		keys := []string{"a", "c"}
+		subset := generic.CopyMapSubset(original, keys)
+
+		t.Log("Should copy only specified keys from the original map")
+		assert.Len(t, subset, 2)
+		assert.Equal(t, 1, subset["a"])
+		assert.Equal(t, 3, subset["c"])
+		assert.NotContains(t, subset, "b")
+		assert.NotContains(t, subset, "d")
+	})
+
+	t.Run("ignores keys not in original map", func(t *testing.T) {
+		t.Parallel()
+
+		original := map[string]int{
+			"a": 1,
+			"b": 2,
+		}
+
+		keys := []string{"a", "c", "d"}
+		subset := generic.CopyMapSubset(original, keys)
+
+		t.Log("Should only copy keys that exist in the original map")
+		assert.Len(t, subset, 1)
+		assert.Equal(t, 1, subset["a"])
+		assert.NotContains(t, subset, "c")
+		assert.NotContains(t, subset, "d")
+	})
+
+	t.Run("returns empty map for nil input", func(t *testing.T) {
+		t.Parallel()
+
+		var original map[string]int = nil
+		keys := []string{"a", "b"}
+		subset := generic.CopyMapSubset(original, keys)
+
+		t.Log("Should return empty map when input map is nil")
+		assert.Nil(t, subset)
+	})
+
+	t.Run("handles empty keys", func(t *testing.T) {
+		t.Parallel()
+
+		original := map[string]int{
+			"a": 1,
+			"b": 2,
+		}
+
+		var keys []string = nil
+		subset := generic.CopyMapSubset(original, keys)
+
+		t.Log("Should return empty map when keys slice is nil")
+		assert.Empty(t, subset)
+		assert.NotNil(t, subset)
+
+		keys = []string{}
+		subset = generic.CopyMapSubset(original, keys)
+
+		t.Log("Should return empty map when keys slice is empty")
+		assert.Empty(t, subset)
+		assert.NotNil(t, subset)
+	})
+
+	t.Run("handles empty map", func(t *testing.T) {
+		t.Parallel()
+
+		original := make(map[string]int)
+		keys := []string{"a", "b"}
+		subset := generic.CopyMapSubset(original, keys)
+
+		t.Log("Should return empty map when input map is empty")
+		assert.Empty(t, subset)
+		assert.NotNil(t, subset)
+	})
+
+	t.Run("handles complex values", func(t *testing.T) {
+		t.Parallel()
+
+		type User struct {
+			Name  string
+			Admin bool
+		}
+
+		original := map[string]User{
+			"a": {Name: "Alice", Admin: false},
+			"b": {Name: "Bob", Admin: true},
+			"c": {Name: "Charlie", Admin: false},
+		}
+
+		keys := []string{"a", "c"}
+		subset := generic.CopyMapSubset(original, keys)
+
+		t.Log("Should work with complex value types")
+		assert.Len(t, subset, 2)
+		assert.Equal(t, User{Name: "Alice", Admin: false}, subset["a"])
+		assert.Equal(t, User{Name: "Charlie", Admin: false}, subset["c"])
+		assert.NotContains(t, subset, "b")
+	})
+}
+
 func TestMerge(t *testing.T) {
 	t.Parallel()
 

set.go

@@ -1,7 +1,7 @@
 package generic
 
 func ToSet[T comparable](slice []T) map[T]struct{} {
-	m := make(map[T]struct{})
+	m := make(map[T]struct{}, len(slice))
 	for _, item := range slice {
 		m[item] = struct{}{}
 	}