Zenn
🗺️

【Go】ジェネリクスと双方向マップを使った enum の実装

2022/12/09に公開
Go
Go
tech
gRPC
tech

この記事は CyberAgent AI tech studio | Go Advent Calendar 2022 8日目 の記事です。

Goのenum

Goには標準でenum（列挙型）のような機能がありません。多言語から入ったので最初戸惑いましたが、とはいえ使いたい場面が出てくるので実装すると大体は以下の2パターンが多いようでした。

実装パターン1：独自型のintを定義して使う

Go言語 - enumを定義する - 覚えたら書く

type Animal int

const (
	Dog Animal = iota + 1
	Cat
	Bear
	Monkey
)

func (a Animal) String() string {
	switch a {
	case Dog:
		return "Dog"
	case Cat:
		return "Cat"
	case Bear:
		return "Bear"
	case Monkey:
		return "Monkey"
	default:
		return "Unknown"
	}
}

int独自型のiota を使用して連続する整数定数を当てはめていき、実質的に番号でアクセスできるようにしています。

また、enumで列挙を定義しただけではなく、どうしてもそれを数字ではなく文字列化したものが欲しくなります。その場合に書かされるのがこの switch 文です。

最初に見たときに、

  • switch文に愚直にマッピングを書かないといけないので記述量が多い
  • （個人的に）可読性が悪くて見にくい

と感じてなんとかならないかなと思いました。

ちなみに iota + 1 にしているのは、以下のように初期化した際に意図せず一つ目に項目になってしまうのを防ぐためです。あえて 0 に UnknownAnimalなど入れておくパターンもあります。

package main

import (
	"fmt"
)

func main() {
	var animal1 Animal
	// 初期値なし変数なので特に何かを入れている意図はない
	fmt.Println(animal1.String())

	// 意図的にDogを入れている
	animal2 := Dog
	fmt.Println(animal2.String())
}

type Animal int

const (
	Dog Animal = iota
	Cat
	Bear
	Monkey
)

func (a Animal) String() string {
	switch a {
	case Dog:
		return "Dog"
	case Cat:
		return "Cat"
	case Bear:
		return "Bear"
	case Monkey:
		return "Monkey"
	default:
		return "Unknown"
	}
}

> go run main.go
Dog
Dog

実装パターン2：intではなく string の独自型をマッピングするパターン

goのenumは文字列でいいのではないかという話

enumとして列挙型は作成したいものの、アクセスする際及び取得して使う際は基本的に string であることが多いため、そもそもマッピングする独自型をstringにすればいいのではという発想です。

type Animal string

const (
	Dog    = Animal("Dog")
	Cat    = Animal("Cat")
	Bear   = Animal("Bear")
	Monkey = Animal("Monkey")
)

func main() {
	animal1 := Dog
	fmt.Println(animal1)
}

確かにこれは圧倒的に記述量が減りました。constに対応関係を書いていくだけなのでシンプルで一見良さそうです。ですが少々欠点があります。それは、

  1. enum用の独自型をpublicにしてしまうと、項目にない要素を自由に作成できてしまう
  2. enum用の独自型をprivateにしてしまうと、変数の引数としてenumの型を指定することがができない

というものです。

enum用の独自型をpublicにしてしまうと、項目にない要素を自由に作成できてしまう

この Animal は string なので、以下のようなことができてしまいます。

func main() {
	animal1 := animal.Dog
	chimera := animal.Animal("Chimera")
	fmt.Println(animal1)
	fmt.Println(chimera)
}

型がstringゆえに、このように未知の動物を作り出すことができてしまいます。しかも丁寧に string も取得できます。

これでは折角 enum で登録する項目を絞ったにも関わらず、新たな項目が外部で作成される可能性があります。（ちなみに int の場合でも外部で作成することはできますが、その int の数字がどの string に対応しているかは switch 文を変更しないとできないので string よりは安全です）

カジュアルに enum を使いたいぐらいの時はいいかもしれませんが、「この enum に入っている項目以外を絶対Animal型にさせたくない！（Animal型から意図せぬStringを取れるようにしたくない）」という場合はちょっと微妙です。

enum用の独自型をprivateにしてしまうと、変数の引数としてenumの型を指定することがができない

「じゃあ外部から新しく作成されないように、animal をプライベートにしよう！」となるかと思いますが、これだとちょっと困ることがあります。それはこの enum 型を引数にとる関数が作成できないことです。

例えば同様に animal.go のパッケージを切って呼び出す場合、publicにしているとこんな感じのことができます。

animal/animal.go

package animal

type Animal string

const (
	Dog    = Animal("Dog")
	Cat    = Animal("Cat")
	Bear   = Animal("Bear")
	Monkey = Animal("Monkey")
)

main.go

package main

import (
	"fmt"

	"github.com/BIwashi/go-play-ground/animal"
)

func main() {
	zoo(animal.Dog, animal.Cat)
}

func zoo(an1, an2 animal.Animal) {
	fmt.Println(an1, an2)
}

関数 zoo の引数に animal.Animal の型を指定することができます。しかし、この型をprivateにしてしまうと引数として型を指定することができなくなってしまいます。

今回の zoo くらいであればanimal/animal.goに書けばいいですが、たくさん引数があるうちの一つにどうしても enum型の値を入れたとなると厳しいです。

ここまで

という感じで、現状色々なところで使用されているenumの実装において個人的に欲しいのは以下の項目でした。

  • int の連番でenumを定義するのはいいが、もう少しスマートにstringを取得できるようマッピングしたい
  • 新しいeunmの項目を外部から勝手に作成されたくはないが、関数等の引数に型指定としては使いたい

これらを満たしたくて考えたのは、int ⇄ string の変換をシームレスに行える用の「双方向マップ」を作成するというアプローチです。

双方向マップ（Bidirectional map）

双方向マップについてです。

双方向マップは、1対1の双方向に変換可能なmapのことを指します。

Goの map の場合は、基本的に key を指定して value を取得する一方通行になります。value を指定して key をとってくることはできません。また、key が違えば value の値が重複していても問題はありません。双方向mapの場合は、key, value がそれぞれ重複のない形で存在し、どちらかを指定すればどちらかの値が確定するというものです。

Bidirectional map - Wikipedia

双方向マップを使用して、enumのintをstringを紐付ける

というわけで要するに、enumで指定したint型（例でいう type Animal int ）とその数字に対応する string を双方向マップで関連付けてしまえばよいのでは？という発想です。というわけでこんな感じです。

Go Playground - The Go Programming Language

animal/animal.go

package animal

import "fmt"

type Animal int

const (
	Dog Animal = iota + 1
	Cat
	Bear
	Monkey
)

type (
	animalBimap struct {
		forwardMap map[string]Animal
		reverseMap map[Animal]string
	}
)

var animalb = NewMustAnimalBimap(
	map[string]Animal{
		"dog":    Dog,
		"cat":    Cat,
		"bear":   Bear,
		"monkey": Monkey,
	},
)

func reverseMap(forwardMap map[string]Animal) (map[Animal]string, error) {
	reverseMap := make(map[Animal]string)
	for k, v := range forwardMap {
		if existKey, ok := reverseMap[v]; ok {
			// valueの重複があった場合エラー
			return nil, fmt.Errorf("duplicate key: %v, %v", existKey, k)
		}
		reverseMap[v] = k
	}
	return reverseMap, nil
}

func NewAnimalBimap(fmap map[string]Animal) (animalBimap, error) {
	rmap, err := reverseMap(fmap)
	if err != nil {
		return animalBimap{}, err
	}
	return animalBimap{
		forwardMap: fmap,
		reverseMap: rmap,
	}, nil
}

func NewMustAnimalBimap(fmap map[string]Animal) animalBimap {
	ab, err := NewAnimalBimap(fmap)
	if err != nil {
		panic(err)
	}
	return ab
}

func (ab *animalBimap) GetForwardMap() map[string]Animal {
	return ab.forwardMap
}

func (ab *animalBimap) GetReverseMap() map[Animal]string {
	return ab.reverseMap
}

func (a Animal) String() string {
	return animalb.GetReverseMap()[a]
}

func NewAnimal(s string) (Animal, error) {
	if a, ok := animalb.GetForwardMap()[s]; ok {
		return a, nil
	}
	return Animal(0), fmt.Errorf("invalid animal: %v", s)
}

main.go

package main

import (
	"fmt"

	"github.com/BIwashi/go-play-ground/animal"
)

func main() {

	newAnimal, err := animal.NewAnimal("dog")
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(newAnimal.String())

	newAnimal, err = animal.NewAnimal("bird")
	if err != nil {
		// enumに存在しない動物を渡すとエラー
		fmt.Println(err)
		return
	}
	fmt.Println(newAnimal.String())

}

ちょっと長いので順番に説明していきます。

（記述量が増えた部分は、最終的にジェネリクスを使っていい感じにします。）

type Animal int

const (
	Dog Animal = iota + 1
	Cat
	Bear
	Monkey
)

ここまでは独自型の int を使う場合と一緒です。

func reverseMap(forwardMap map[string]Animal) (map[Animal]string, error) {
	reverseMap := make(map[Animal]string)
	for k, v := range forwardMap {
		if existKey, ok := reverseMap[v]; ok {
			// valueの重複があった場合エラー
			return nil, fmt.Errorf("duplicate key: %v, %v", existKey, k)
		}
		reverseMap[v] = k
	}
	return reverseMap, nil
}

func NewAnimalBimap(fmap map[string]Animal) (animalBimap, error) {
	rmap, err := reverseMap(fmap)
	if err != nil {
		return animalBimap{}, err
	}
	return animalBimap{
		forwardMap: fmap,
		reverseMap: rmap,
	}, nil
}

NewAnimalBimap で map を引数に取り、内部でその map を反転させた map を作成する reverseMap を呼んでいます。最初の双方向マップの説明で記載した通り、mapにおいて key の重複は許されませんが、value は許されます。例えば

seafoodMap := map[string]string{
		"salmon": "sushi",
		"maguro": "sashimi",
		"tuna":   "sashimi",
		"ikura":  "sushi",
	}

それぞれどの食べ方が好きかどうかは置いておいて、このように value の料理名が重複している map を作成することはできます。しかし、これを入れ替えると value が今度 key になるので重複してはいけません。

// DuplicateLitKeyのエラーが出て作成できない
reverseSeafoodMap := map[string]string{
		"sushi":   "salmon",
		"sashimi": "maguro",
		"sashimi": "tuna",
		"sushi":   "ikura",
	}

そのため、反転させた map を作成する reverseMap では value に重複を判定して、存在していた場合はエラーになるようにしています。

type (
	animalBimap struct {
		forwardMap map[string]Animal
		reverseMap map[Animal]string
	}
)

func (ab *animalBimap) GetForwardMap() map[string]Animal {
	return ab.forwardMap
}

func (ab *animalBimap) GetReverseMap() map[Animal]string {
	return ab.reverseMap
}

そして、animalBimap の構造体の中に順方向と逆方向の map を入れて作成します。

これらの map はそれぞれ getter を使用してアクセスできるようにしています。

func (a Animal) String() string {
	return animalb.GetReverseMap()[a]
}

func NewAnimal(s string) (Animal, error) {
	if a, ok := animalb.GetForwardMap()[s]; ok {
		return a, nil
	}
	return Animal(0), fmt.Errorf("invalid animal: %v", s)
}

そしてこの双方向 map を利用して以下を行なっています。

  • NewAnimal(s string)：primitive型の string を使用して、int独自型Animalを作成する（mapに存在しない項目の場合は作成できない）
  • String()：int独自型から対応する string を取得する

これによって enum で実現したかった要素の限定と int ⇄ string の相互変換を実装できました。

…記述量増えてない…？（ジェネリクスの活用）

そうです。こんなに書くくらいなら switch 文を書いた方がまだ少ないです。というわけでこの双方向マップの作成をジェネリクスを使って行なってみます。

Go Playground - The Go Programming Language

bimap/bimap.go

package bimap

import "fmt"

type (
	// Reverseする前のmapもセットで引き回す用の構造体を定義
	Bimap[K, V comparable] struct {
		forwardMap map[K]V
		reverseMap map[V]K
	}
)

func ReverseMap[K, V comparable](forwardMap map[K]V) (map[V]K, error) {
	m := make(map[V]K, len(forwardMap))
	for k, v := range forwardMap {
		if existKey, ok := m[v]; ok {
			// valueの重複があった場合はエラー
			return nil, fmt.Errorf("duplicate key: %v, %v", existKey, k)
		} else {
			m[v] = k
		}
	}
	return m, nil
}

func (b Bimap[K, V]) ReverseMap() (map[V]K, error) {
	return ReverseMap(b.forwardMap)
}

// errorでハンドリングしたい場合
func NewBimap[K, V comparable](fmap map[K]V) (Bimap[K, V], error) {
	rmap, err := ReverseMap(fmap)
	if err != nil {
		return Bimap[K, V]{}, fmt.Errorf("unable to create bimap: %w", err)
	}
	return Bimap[K, V]{
		forwardMap: fmap,
		reverseMap: rmap,
	}, nil
}

// panicで落としたい場合
func NewMustBimap[K, V comparable](forwardMap map[K]V) Bimap[K, V] {
	bimap, err := NewBimap(forwardMap)
	if err != nil {
		panic(err)
	}
	return bimap
}

func (m *Bimap[K, V]) GetForwardMap() map[K]V {
	return m.forwardMap
}

func (m *Bimap[K, V]) GetReverseMap() map[V]K {
	return m.reverseMap
}

enum/animal.go

package enum

import (
	"fmt"

	"github.com/BIwashi/go-play-ground/bimap"
)

type Animal int

const (
	Dog Animal = iota + 1
	Cat
	Bear
	Monkey
)

var enumAnimal = bimap.NewMustBimap(
	map[string]Animal{
		"dog":    Dog,
		"cat":    Cat,
		"bear":   Bear,
		"monkey": Monkey,
	},
)

func (a Animal) String() (string, error) {
	if s, ok := enumAnimal.GetReverseMap()[a]; ok {
		return s, nil
	}
	return "", fmt.Errorf("invalid animal: %d", a)
}

func NewAnimal(animal string) (Animal, error) {
	if a, ok := enumAnimal.GetForwardMap()[animal]; ok {
		return a, nil
	}
	return Animal(0), fmt.Errorf("invalid animal: %s", animal)
}

main.go

package main

import (
	"fmt"

	"github.com/BIwashi/go-play-ground/enum"
)

func main() {

	newAnimal, err := enum.NewAnimal("dog")
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(newAnimal.String())

	fmt.Println(enum.Cat.String())

	// enumの項目にないので失敗する
	if _, err := enum.NewAnimal("chimera"); err != nil {
		fmt.Println(err)
		return
	}

}

双方向 Map を作成する処理は共通ですが、その map における key は enum に使用する int 独自型であるため、map[int]string を引数に取り ReverseMap を作成するという処理を作ることができません。というわけでジェネリクスを使うことで独自型を柔軟に入れられる形で実装してみました。

bimap の pkg は共通で使用し、新しく enum を作成したい際はenum/animal.go のように双方向 map を使用して int 独自型とstring型の相互変換を行えるようにしています。ちなみに NewMustBimapというエラーではなく panic で落とす方の関数を使用しているのは、compileでのエラーチェックではなくそもそもmapの値が不正（例：value の重複）である場合は起動させないようにし、かつグローバルで定義するためです（Regexp をグローバルで使用する際に MustCompileを使用するのと同じような使い方をしています）。

というわけで、最初の switch 文での実装と比較するとこんな感じです。

// switch
func (a Animal) String() string {
	switch a {
	case Dog:
		return "Dog"
	case Cat:
		return "Cat"
	case Bear:
		return "Bear"
	case Monkey:
		return "Monkey"
	default:
		return "Unknown"
	}
}

// Generics
var enumAnimal = bimap.NewMustBimap(
	map[string]Animal{
		"dog":    Dog,
		"cat":    Cat,
		"bear":   Bear,
		"monkey": Monkey,
	},
)

func (a Animal) String() (string, error) {
	if s, ok := enumAnimal.GetReverseMap()[a]; ok {
		return s, nil
	}
	return "", fmt.Errorf("invalid animal: %d", a)
}

独自型intからstringを取得する処理は…そんなに変わってない…！？ 気もしますが、map にまとめることで少し可読性が増した気がします。一方「そのstringがenumに登録されていた場合、独自型を作成する（登録されていない場合はエラー）」という場合はどうなるでしょうか？

// // switch
func NewAnimal(animal string) (Animal, error) {
	switch animal {
	case "dog":
		return Dog, nil
	case "cat":
		return Cat, nil
	case "bear":
		return Bear, nil
	case "monkey":
		return Monkey, nil
	default:
		return Animal(0), fmt.Errorf("invalid animal: %s", animal)
	}
}

// Generics
func NewAnimal(animal string) (Animal, error) {
	if a, ok := enumAnimal.GetForwardMap()[animal]; ok {
		return a, nil
	}
	return Animal(0), fmt.Errorf("invalid animal: %s", animal)
}

そうです。swicth 文での実装はまた逆向きのswitch 文を書かないといけません。これが少量ならいいですが、段々と数が増えてくるとほぼ同じ相互変換用の switch 文を加筆していく必要がありちょっとめんどくさいです。しかも、2種類の switch 文を同時にちゃんと追加したことを確認しないと意図せぬ enum が作成されてしまう可能性があります。

例えば NewAnimal(animal string) の方の switch 文には追加したものの、 String() の方の switch 文に書き忘れていたため、Animal型を作成できたのに string が取得できないという状況が発生してしまいます。

そういった点で、mapという key と value （ int と string ）をセットで最初に登録しないといけない、かつそれを違反していると起動できないという制約があるので、こうした実装漏れはなくなります。これが一番の利点だと思います。

proto の enum から アプリケーション enum にマッピングする

proto には enum 型が存在しますが、そのままアプリケーションでその型を引き回すのではなくアプリケーション用のenumに変換したい場合があります。その際もこの双方向マップが役に立ちます。

例えば以下のような enum 型を proto で定義します。

enum Animal {
  ANIMAL_UNSPECIFIED = 0;
  ANIMAL_DOG = 1;
  ANIMAL_CAT = 2;
  ANIMAL_BEAR = 3;
  ANIMAL_MONKEY = 4;
}

これをもとに生成される Go のコードは以下のようなものです。

type Animal int32

const (
	Animal_ANIMAL_UNSPECIFIED Animal = 0
	Animal_ANIMAL_DOG         Animal = 1
	Animal_ANIMAL_CAT         Animal = 2
	Animal_ANIMAL_BEAR        Animal = 3
	Animal_ANIMAL_MONKEY      Animal = 4
)

// Enum value maps for Animal.
var (
	Animal_name = map[int32]string{
		0: "ANIMAL_UNSPECIFIED",
		1: "ANIMAL_DOG",
		2: "ANIMAL_CAT",
		3: "ANIMAL_BEAR",
		4: "ANIMAL_MONKEY",
	}
	Animal_value = map[string]int32{
		"ANIMAL_UNSPECIFIED": 0,
		"ANIMAL_DOG":         1,
		"ANIMAL_CAT":         2,
		"ANIMAL_BEAR":        3,
		"ANIMAL_MONKEY":      4,
	}
)

func (x Animal) Enum() *Animal {
	p := new(Animal)
	*p = x
	return p
}

func (x Animal) String() string {
	return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x))
}

func (Animal) Descriptor() protoreflect.EnumDescriptor {
	return file_customer_v1_enum_proto_enumTypes[1].Descriptor()
}

func (Animal) Type() protoreflect.EnumType {
	return &file_customer_v1_enum_proto_enumTypes[1]
}

func (x Animal) Number() protoreflect.EnumNumber {
	return protoreflect.EnumNumber(x)
}

// Deprecated: Use Animal.Descriptor instead.
func (Animal) EnumDescriptor() ([]byte, []int) {
	return file_customer_v1_enum_proto_rawDescGZIP(), []int{1}
}

（おや… int32 → string と string → int32 のmap作ってるのでアプローチが似てる…？）

この enum を proto での引数などに指定すると、得られるのは ANIMAL_DOG などの const で設定された値です。これを先ほどのアプリケーション側で実装した独自型とマッピングしたくなってきます。

というわけでこんな感じです。

var animalProtoBimap = bimap.NewMustBimap(
	map[apiv1.Animal]enum.Animal{
		apiv1.Animal_ANIMAL_CAT:    enum.CAT,
		apiv1.Animal_ANIMAL_DOG:    enum.DOG,
		apiv1.Animal_ANIMAL_BEAR:   enum.BEAR,
		apiv1.Animal_ANIMAL_MONKEY: enum.MONKEY,
	},
)

func animalProtoToEnum(animal apiv1.Animal) (enum.Animal, error) {
	a, ok := animalBimap.GetForwardMap()[animal]
	if !ok {
		return "", fmt.Errorf("invalid argument error")
	}
	return a, nil
}

proto の enum 型を先ほど作成した enumパッケージの Animal 型のマッピングすることができました。

すでに const での値（ Animal_ANIMAL_CAT など）が作成されているので、map で対応関係を明示的に示すだけです。これは結構使い勝手がいいです。

proto で外部から来る enum に関してはこのようにマッピングすると、漏れも少なくなおかつスマートに実装できた気がします。

まとめ

今回は Go の enum をジェネリクスを使用した双方向マップを使用することで新しく実装してみました。動機は enum を簡潔に書きたいというものでしたが、結果的に双方向マップの実装は色々なところで使い道がありそうだなと思いました。

今後も使えそうな場面では、でも過剰には使わないという程度でジェネリクスを使用していきたいと思います。

Discussion