feat: preserve type alias of slices and maps (see #365)

parallel/slice.go

@@ -67,8 +67,8 @@ func Times[T any](count int, iteratee func(index int) T) []T {
 
 // GroupBy returns an object composed of keys generated from the results of running each element of collection through iteratee.
 // `iteratee` is call in parallel.
-func GroupBy[T any, U comparable](collection []T, iteratee func(item T) U) map[U][]T {
-	result := map[U][]T{}
+func GroupBy[T any, U comparable, Slice ~[]T](collection Slice, iteratee func(item T) U) map[U]Slice {
+	result := map[U]Slice{}
 
 	var mu sync.Mutex
 	var wg sync.WaitGroup
@@ -96,8 +96,8 @@ func GroupBy[T any, U comparable](collection []T, iteratee func(item T) U) map[U
 // determined by the order they occur in collection. The grouping is generated from the results
 // of running each element of collection through iteratee.
 // `iteratee` is call in parallel.
-func PartitionBy[T any, K comparable](collection []T, iteratee func(item T) K) [][]T {
-	result := [][]T{}
+func PartitionBy[T any, K comparable, Slice ~[]T](collection Slice, iteratee func(item T) K) []Slice {
+	result := []Slice{}
 	seen := map[K]int{}
 
 	var mu sync.Mutex

parallel/slice_test.go

@@ -5,20 +5,20 @@ import (
 	"strconv"
 	"sync/atomic"
 	"testing"
-	
+
 	"github.com/stretchr/testify/assert"
 )
 
 func TestMap(t *testing.T) {
 	is := assert.New(t)
-	
+
 	result1 := Map([]int{1, 2, 3, 4}, func(x int, _ int) string {
 		return "Hello"
 	})
 	result2 := Map([]int64{1, 2, 3, 4}, func(x int64, _ int) string {
 		return strconv.FormatInt(x, 10)
 	})
-	
+
 	is.Equal(len(result1), 4)
 	is.Equal(len(result2), 4)
 	is.Equal(result1, []string{"Hello", "Hello", "Hello", "Hello"})
@@ -27,52 +27,59 @@ func TestMap(t *testing.T) {
 
 func TestForEach(t *testing.T) {
 	is := assert.New(t)
-	
+
 	var counter uint64
 	collection := []int{1, 2, 3, 4}
 	ForEach(collection, func(x int, i int) {
 		atomic.AddUint64(&counter, 1)
 	})
-	
+
 	is.Equal(uint64(4), atomic.LoadUint64(&counter))
 }
 
 func TestTimes(t *testing.T) {
 	is := assert.New(t)
-	
+
 	result1 := Times(3, func(i int) string {
 		return strconv.FormatInt(int64(i), 10)
 	})
-	
+
 	is.Equal(len(result1), 3)
 	is.Equal(result1, []string{"0", "1", "2"})
 }
 
 func TestGroupBy(t *testing.T) {
 	is := assert.New(t)
-	
+
 	result1 := GroupBy([]int{0, 1, 2, 3, 4, 5}, func(i int) int {
 		return i % 3
 	})
-	
+
 	// order
 	for x := range result1 {
 		sort.Slice(result1[x], func(i, j int) bool {
 			return result1[x][i] < result1[x][j]
 		})
 	}
-	
+
 	is.EqualValues(len(result1), 3)
 	is.EqualValues(result1, map[int][]int{
 		0: {0, 3},
 		1: {1, 4},
 		2: {2, 5},
 	})
+
+	type myStrings []string
+	allStrings := myStrings{"", "foo", "bar"}
+	nonempty := GroupBy(allStrings, func(i string) int {
+		return 42
+	})
+	is.IsType(nonempty[42], allStrings, "type preserved")
 }
 
 func TestPartitionBy(t *testing.T) {
 	is := assert.New(t)
-	
+
 	oddEven := func(x int) string {
 		if x < 0 {
 			return "negative"
@@ -81,10 +88,10 @@ func TestPartitionBy(t *testing.T) {
 		}
 		return "odd"
 	}
-	
+
 	result1 := PartitionBy([]int{-2, -1, 0, 1, 2, 3, 4, 5}, oddEven)
 	result2 := PartitionBy([]int{}, oddEven)
-	
+
 	// order
 	sort.Slice(result1, func(i, j int) bool {
 		return result1[i][0] < result1[j][0]
@@ -94,7 +101,14 @@ func TestPartitionBy(t *testing.T) {
 			return result1[x][i] < result1[x][j]
 		})
 	}
-	
+
 	is.ElementsMatch(result1, [][]int{{-2, -1}, {0, 2, 4}, {1, 3, 5}})
 	is.Equal(result2, [][]int{})
+
+	type myStrings []string
+	allStrings := myStrings{"", "foo", "bar"}
+	nonempty := PartitionBy(allStrings, func(item string) int {
+		return len(item)
+	})
+	is.IsType(nonempty[0], allStrings, "type preserved")
 }