From 280ecb5cef295ed877e1791f921392f7e619abe7 Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Fri, 2 Dec 2022 14:53:57 +0800
Subject: [PATCH] feat: add SortBy function for slice

---
 slice/slice.go          |  6 ++++++
 slice/slice_internal.go | 32 +++++++++++++++++++++++++++++---
 slice/slice_test.go     | 30 ++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/slice/slice.go b/slice/slice.go
index db50fc2..307ef34 100644
--- a/slice/slice.go
+++ b/slice/slice.go
@@ -725,6 +725,12 @@ func Sort[T lancetconstraints.Ordered](slice []T, sortOrder ...string) {
 	}
 }
 
+// SortBy sorts the slice in ascending order as determined by the less function.
+// This sort is not guaranteed to be stable
+func SortBy[T any](slice []T, less func(a, b T) bool) {
+	quickSortBy(slice, 0, len(slice)-1, less)
+}
+
 // SortByField return sorted slice by field
 // slice element should be struct, field type should be int, uint, string, or bool
 // default sortType is ascending (asc), if descending order, set sortType to desc
diff --git a/slice/slice_internal.go b/slice/slice_internal.go
index 9ff8996..9954325 100644
--- a/slice/slice_internal.go
+++ b/slice/slice_internal.go
@@ -29,14 +29,14 @@ func sliceElemType(reflectType reflect.Type) reflect.Type {
 
 func quickSort[T lancetconstraints.Ordered](slice []T, lowIndex, highIndex int, order string) {
 	if lowIndex < highIndex {
-		p := partition(slice, lowIndex, highIndex, order)
+		p := partitionOrderedSlice(slice, lowIndex, highIndex, order)
 		quickSort(slice, lowIndex, p-1, order)
 		quickSort(slice, p+1, highIndex, order)
 	}
 }
 
-// partition split slice into two parts for quick sort
-func partition[T lancetconstraints.Ordered](slice []T, lowIndex, highIndex int, order string) int {
+// partitionOrderedSlice split ordered slice into two parts for quick sort
+func partitionOrderedSlice[T lancetconstraints.Ordered](slice []T, lowIndex, highIndex int, order string) int {
 	p := slice[highIndex]
 	i := lowIndex
 
@@ -59,6 +59,32 @@ func partition[T lancetconstraints.Ordered](slice []T, lowIndex, highIndex int,
 	return i
 }
 
+func quickSortBy[T any](slice []T, lowIndex, highIndex int, less func(a, b T) bool) {
+	if lowIndex < highIndex {
+		p := partitionAnySlice(slice, lowIndex, highIndex, less)
+		quickSortBy(slice, lowIndex, p-1, less)
+		quickSortBy(slice, p+1, highIndex, less)
+	}
+}
+
+// partitionAnySlice split any slice into two parts for quick sort
+func partitionAnySlice[T any](slice []T, lowIndex, highIndex int, less func(a, b T) bool) int {
+	p := slice[highIndex]
+	i := lowIndex
+
+	for j := lowIndex; j < highIndex; j++ {
+
+		if less(slice[j], p) {
+			swap(slice, i, j)
+			i++
+		}
+	}
+
+	swap(slice, i, highIndex)
+
+	return i
+}
+
 // swap two slice value at index i and j
 func swap[T any](slice []T, i, j int) {
 	slice[i], slice[j] = slice[j], slice[i]
diff --git a/slice/slice_test.go b/slice/slice_test.go
index e0fcb03..36c7459 100644
--- a/slice/slice_test.go
+++ b/slice/slice_test.go
@@ -558,6 +558,36 @@ func TestSort(t *testing.T) {
 	assert.Equal([]string{"e", "d", "c", "b", "a"}, strings)
 }
 
+func TestSortBy(t *testing.T) {
+	assert := internal.NewAssert(t, "TestSortBy")
+
+	numbers := []int{1, 4, 3, 2, 5}
+
+	SortBy(numbers, func(a, b int) bool {
+		return a < b
+	})
+	assert.Equal([]int{1, 2, 3, 4, 5}, numbers)
+
+	type User struct {
+		Name string
+		Age  uint
+	}
+
+	users := []User{
+		{Name: "a", Age: 21},
+		{Name: "b", Age: 15},
+		{Name: "c", Age: 100}}
+
+	SortBy(users, func(a, b User) bool {
+		return a.Age < b.Age
+	})
+
+	t.Logf("sort users by age: %v", users)
+
+	// output
+	// [{b 15} {a 21} {c 100}]
+}
+
 func TestSortByFielDesc(t *testing.T) {
 	assert := internal.NewAssert(t, "TestSortByFielDesc")