🦊

【Go】Genericsを使ってCollection操作をする

2022/08/21に公開

tech

概要

Go1.18からGenericsが導入されたので、よくあるCollection操作をGenericsにして共通化してみた。
プロジェクトでutil関数として置いてみたら、

・コードの書き方に統一性が生まれた
・makeでcap確保し忘れ等も防げた

のでオススメ。

Filter

スライスの要素から条件に引っかかるもののみ抜き出してくる。

「SQLで取ってきたデータを、アプリケーション側でさらにフィルタしてクライアントに返したい」と言った場合に便利。
一番よく使う関数かも。

func Filter[S ~[]T, T any](s S, f func(T) bool) S {
	ret := make(S, 0, len(s))
	for _, e := range s {
		if f(e) {
			ret = append(ret, e)
		}
	}
	return ret
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:  20,
		},
		{
			userID: "B",
			age:  30,
		},
	}

        // ageが20以下に絞る
	newUsers := Filter(users, func(u *user) bool {
		return u.age <= 20
	})

	fmt.Printf("len=%d\n", len(newUsers))
	for _, u := range newUsers {
		fmt.Printf("user=%v\n", u)
	}
	/*
	 len=1
	 user=&{A 20}
	*/
}

MapBy

要素に含まれるフィールドをkey、要素自体をvalueとしてmapを生成する。

userのIDをkeyとしてmapを作りたい、といった場合に便利。

func MapBy[T any, V comparable](s []T, f func(T) V) map[V]T {
	ret := make(map[V]T, len(s))
	for _, e := range s {
		ret[f(e)] = e
	}
	return ret
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:  20,
		},
		{
			userID: "B",
			age:  30,
		},
	}

	mapByUserID := ToMap(users, func(u *user) string {
		return u.userID
	})

	fmt.Printf("len=%d\n", len(mapByUserID))
	for key, value := range mapByUserID {
		fmt.Printf("key=%s\nvalue=%v\n", key, value)
	}
	
	/*
	 len=2

	 key=A
	 value=&{A 20}

	 key=B
	 value=&{B 30}
	*/
}

Select

要素に含まれるフィールドの値を新たなスライスとして取得する。

複数のuserIDを取得したい、といった場合に便利。

func Select[T, V any](s []T, f func(T) V) []V {
	ret := make([]V, 0, len(s))
	for _, e := range s {
		ret = append(ret, f(e))
	}
	return ret
}

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
	}

	userIDs := Select(users, func(u *user) string {
		return u.userID
	})

	fmt.Printf("len=%d\n", len(userIDs))
	fmt.Printf("%s\n", userIDs)
	/*
		len=2
		[A B]
	*/
}

First

要素の最初を取得する。
ぶっちゃけ、indexを0でやるのと記述量変わらないが、コードが長くなってくると意外と可読性が高かったりするので重宝してる。
あとはlenのチェック入れればindex out of rangeにならないので、この時点でのpanicも防げたり。

func First[T any](s []T) T {
	if len(s) == 0 {
		var v T
		return v
	}
	return s[0]
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
	}

	user := First(users)

	fmt.Printf("%v\n", user)
	/*
		&{A  20}
	*/
}

Last

要素の最後を取得する。
(Firstの逆)

func Last[T any](s []T) T {
	if len(s) == 0 {
		var v T
		return v
	}
	return s[len(s)-1]
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
	}

	user := Last(users)

	fmt.Printf("%v\n", user)
	/*
		&{B  30}
	*/
}

Sort

文字通りソートをしたい時に使う。

下記のようにすれば「元のスライスの並びは変更したくなくて、ソートされた状態で新しいスライスを生成したい」など小回り効くようになる。

func Sort[S ~[]T, T any](s S, f func(e1, e2 T) bool) S {
	ret := make(S, 0, len(s))
	ret = append(ret, s...)
	sort.Slice(ret, func(i, j int) bool {
		return f(ret[i], ret[j])
	})
	return ret
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
	}

        // 年取ってる順に並べる
	newUsers := Sort(users, func(u1, u2 *user) bool {
		return u1.age > u2.age
	})

	fmt.Printf("len=%d\n", len(newUsers))
	for _, u := range newUsers {
		fmt.Printf("user=%v\n", u)
	}
	/*
		len=2
		user=&{B  30}
		user=&{A  20}
	*/
}

Shuffle

要素の並びをランダムにする。
ゲームみたいなサービスだと結構使うかも。

func init() {
        // init時のSeedの設定を忘れずに
	seed, _ := crand.Int(crand.Reader, big.NewInt(math.MaxInt64))
	rand.Seed(seed.Int64())
}

func Shuffle[S ~[]T, T any](s S) S {
	ret := make(S, 0, len(s))
	ret = append(ret, s...)
	for i := len(s); i > 1; i-- {
		j := rand.Intn(i)
		ret[i-1], ret[j] = ret[j], ret[i-1]
	}
	return ret
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
		{
			userID: "C",
			age:    40,
		},
	}

	newUsers := Shuffle(users)

	fmt.Printf("len=%d\n", len(newUsers))
	for _, u := range newUsers {
		fmt.Printf("user=%v\n", u)
	}
	/*
		len=3
		user=&{C  40}
		user=&{A  20}
		user=&{B  30}
	*/
}

Split

指定した要素数のスライスに分割する。

func Split[S ~[]T, T any](s S, size int) []S {
	length := len(s)
	splits := make([]S, 0, length/size+1)
	for i := 0; i < length; i += size {
		end := i + size
		if length < end {
			end = length
		}
		splits = append(splits, s[i:end])
	}
	return splits
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
		{
			userID: "C",
			age:    40,
		},
	}

	newUserSplits := Split(users, 2)

	fmt.Printf("len=%d\n", len(newUserSplits))
	for _, newUsers := range newUserSplits {
		fmt.Printf("newUsers=%v\n", newUsers)
		for _, u := range newUsers {
			fmt.Printf("user=%v\n", u)
		}
	}
	/*
		len=2

		newUsers=[0xc00010d470 0xc00010d4a0]
		user=&{A  20}
		user=&{B  30}

		newUsers=[0xc00010d4d0]
		user=&{C  40}
	*/
}

Includes

引数で渡した関数に引っかかる要素があるかどうかを返す。

func Includes[S ~[]T, T any](s S, f func(T) bool) bool {
	for _, e := range s {
		if f(e) {
			return true
		}
	}
	return false
}

使用例

type user struct {
   userID string
   age  int
}

type users []*user

func main() {
	users := users{
		{
			userID: "A",
			age:    20,
		},
		{
			userID: "B",
			age:    30,
		},
		{
			userID: "C",
			age:    40,
		},
	}

	includesA := Includes(users, func(u *user) bool {
		return u.userID == "A"
	})

	fmt.Printf("%v\n", includesA)
	/*
		true
	*/
}

まとめ

・genericsを使うことによって、煩雑だったcollection操作が楽になる

・コードが統一されるので、プロジェクト初期の段階で導入しておくと良さそう