Skip to content

Commit

Permalink
feat: preserve type alias of slices and maps (see #365)
Browse files Browse the repository at this point in the history
  • Loading branch information
samber committed Jun 29, 2024
1 parent 3ba93a1 commit 93686db
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 18 deletions.
8 changes: 4 additions & 4 deletions parallel/slice.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
42 changes: 28 additions & 14 deletions parallel/slice_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"})
Expand All @@ -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"
Expand All @@ -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]
Expand All @@ -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")
}

0 comments on commit 93686db

Please sign in to comment.