From a24352d043b55f742c0ae55543610e3c8b06a8c0 Mon Sep 17 00:00:00 2001
From: xing <fthvgb1>
Date: Mon, 10 Oct 2022 21:33:41 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 actions/common/comments.go       |  4 ++--
 helper/func.go                   | 14 ++++++++----
 helper/func_test.go              | 37 +++++++++++++++++++++++++++-----
 middleware/validateservername.go |  4 ++--
 4 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/actions/common/comments.go b/actions/common/comments.go
index 9c43df4..69b772e 100644
--- a/actions/common/comments.go
+++ b/actions/common/comments.go
@@ -58,7 +58,7 @@ func getCommentByIds(args ...any) (map[uint64]models.WpComments, error) {
 	if err != nil {
 		return m, err
 	}
-	return helper.SliceToMap(r, func(t models.WpComments) uint64 {
+	return helper.SimpleSliceToMap(r, func(t models.WpComments) uint64 {
 		return t.CommentId
-	}, true), err
+	}), err
 }
diff --git a/helper/func.go b/helper/func.go
index 99607e1..0cd5d28 100644
--- a/helper/func.go
+++ b/helper/func.go
@@ -231,16 +231,22 @@ func SliceSelfReverse[T any](arr []T) []T {
 	return arr
 }
 
-func SliceToMap[K comparable, V any](arr []V, fn func(V) K, isCoverPrev bool) map[K]V {
-	m := make(map[K]V)
+func SimpleSliceToMap[K comparable, V any](arr []V, fn func(V) K) map[K]V {
+	return SliceToMap(arr, func(v V) (K, V) {
+		return fn(v), v
+	}, true)
+}
+
+func SliceToMap[K comparable, V, T any](arr []V, fn func(V) (K, T), isCoverPrev bool) map[K]T {
+	m := make(map[K]T)
 	for _, v := range arr {
-		k := fn(v)
+		k, r := fn(v)
 		if !isCoverPrev {
 			if _, ok := m[k]; ok {
 				continue
 			}
 		}
-		m[k] = v
+		m[k] = r
 	}
 	return m
 }
diff --git a/helper/func_test.go b/helper/func_test.go
index 539fe8a..11eba36 100644
--- a/helper/func_test.go
+++ b/helper/func_test.go
@@ -513,7 +513,7 @@ func TestSliceToMap(t *testing.T) {
 	}
 	type args struct {
 		arr         []ss
-		fn          func(ss) int
+		fn          func(ss) (int, ss)
 		isCoverPrev bool
 	}
 	tests := []struct {
@@ -525,8 +525,8 @@ func TestSliceToMap(t *testing.T) {
 			name: "t1",
 			args: args{
 				arr: []ss{{1, "k1"}, {2, "v2"}, {2, "v3"}},
-				fn: func(s ss) int {
-					return s.id
+				fn: func(s ss) (int, ss) {
+					return s.id, s
 				},
 				isCoverPrev: true,
 			},
@@ -535,8 +535,8 @@ func TestSliceToMap(t *testing.T) {
 			name: "t2",
 			args: args{
 				arr: []ss{{1, "k1"}, {2, "v2"}, {2, "v3"}},
-				fn: func(s ss) int {
-					return s.id
+				fn: func(s ss) (int, ss) {
+					return s.id, s
 				},
 				isCoverPrev: false,
 			},
@@ -551,3 +551,30 @@ func TestSliceToMap(t *testing.T) {
 		})
 	}
 }
+
+func TestSimpleSliceToMap(t *testing.T) {
+	type args struct {
+		arr []int
+		fn  func(int) int
+	}
+	tests := []struct {
+		name string
+		args args
+		want map[int]int
+	}{
+		{
+			name: "t1",
+			args: args{arr: []int{1, 2, 3}, fn: func(i int) int {
+				return i
+			}},
+			want: map[int]int{1: 1, 2: 2, 3: 3},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := SimpleSliceToMap(tt.args.arr, tt.args.fn); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("SimpleSliceToMap() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/middleware/validateservername.go b/middleware/validateservername.go
index 57febb2..b21401d 100644
--- a/middleware/validateservername.go
+++ b/middleware/validateservername.go
@@ -9,9 +9,9 @@ import (
 )
 
 func ValidateServerNames() func(ctx *gin.Context) {
-	serverName := helper.SliceToMap(vars.Conf.TrustServerNames, func(v string) string {
+	serverName := helper.SimpleSliceToMap(vars.Conf.TrustServerNames, func(v string) string {
 		return v
-	}, true)
+	})
 	return func(c *gin.Context) {
 		if len(serverName) > 0 {
 			if _, ok := serverName[strings.Split(c.Request.Host, ":")[0]]; !ok {