From f675f6433c698777c203eb8cdea180251cec71fc Mon Sep 17 00:00:00 2001
From: rw <me@rwinslow.com>
Date: Thu, 26 Jul 2018 16:40:56 -0700
Subject: [PATCH] go: give enums their own scalar types

---
 src/idl_gen_go.cpp            | 17 +++++++++++++++-
 tests/MyGame/Example/Any.go   | 11 ++++++-----
 tests/MyGame/Example/Color.go |  9 +++++----
 tests/go_test.go              | 45 ++++++++++++++++++++-----------------------
 4 files changed, 48 insertions(+), 34 deletions(-)

diff --git a/src/idl_gen_go.cpp b/src/idl_gen_go.cpp
index 8c95e71..c5767b4 100644
--- a/src/idl_gen_go.cpp
+++ b/src/idl_gen_go.cpp
@@ -86,6 +86,18 @@ static void BeginClass(const StructDef &struct_def, std::string *code_ptr) {
   code += "\n}\n\n";
 }
 
+// Construct the name of the type alias for this enum.
+std::string GetEnumTypeName(const EnumDef &enum_def) {
+  return GoIdentity(enum_def.name);
+}
+
+// Create a type for the enum values.
+static void GenEnumType(const EnumDef &enum_def, std::string *code_ptr) {
+  std::string &code = *code_ptr;
+  code += "type " + GetEnumTypeName(enum_def) + " = ";
+  code += GenTypeBasic(enum_def.underlying_type) + "\n";
+}
+
 // Begin enum code with a class declaration.
 static void BeginEnum(std::string *code_ptr) {
   std::string &code = *code_ptr;
@@ -99,6 +111,8 @@ static void EnumMember(const EnumDef &enum_def, const EnumVal ev,
   code += "\t";
   code += enum_def.name;
   code += ev.name;
+  code += " ";
+  code += GetEnumTypeName(enum_def);
   code += " = ";
   code += NumToString(ev.value) + "\n";
 }
@@ -114,7 +128,7 @@ static void BeginEnumNames(const EnumDef &enum_def, std::string *code_ptr) {
   std::string &code = *code_ptr;
   code += "var EnumNames";
   code += enum_def.name;
-  code += " = map[int]string{\n";
+  code += " = map[" + GetEnumTypeName(enum_def) + "]string{\n";
 }
 
 // A single enum name member.
@@ -636,6 +650,7 @@ static void GenEnum(const EnumDef &enum_def, std::string *code_ptr) {
   if (enum_def.generated) return;
 
   GenComment(enum_def.doc_comment, code_ptr, nullptr);
+  GenEnumType(enum_def, code_ptr);
   BeginEnum(code_ptr);
   for (auto it = enum_def.vals.vec.begin(); it != enum_def.vals.vec.end();
        ++it) {
diff --git a/tests/MyGame/Example/Any.go b/tests/MyGame/Example/Any.go
index d69b141..06c238e 100644
--- a/tests/MyGame/Example/Any.go
+++ b/tests/MyGame/Example/Any.go
@@ -2,14 +2,15 @@
 
 package Example
 
+type Any = byte
 const (
-	AnyNONE = 0
-	AnyMonster = 1
-	AnyTestSimpleTableWithEnum = 2
-	AnyMyGame_Example2_Monster = 3
+	AnyNONE Any = 0
+	AnyMonster Any = 1
+	AnyTestSimpleTableWithEnum Any = 2
+	AnyMyGame_Example2_Monster Any = 3
 )
 
-var EnumNamesAny = map[int]string{
+var EnumNamesAny = map[Any]string{
 	AnyNONE:"NONE",
 	AnyMonster:"Monster",
 	AnyTestSimpleTableWithEnum:"TestSimpleTableWithEnum",
diff --git a/tests/MyGame/Example/Color.go b/tests/MyGame/Example/Color.go
index d1eac54..7318cbc 100644
--- a/tests/MyGame/Example/Color.go
+++ b/tests/MyGame/Example/Color.go
@@ -2,13 +2,14 @@
 
 package Example
 
+type Color = int8
 const (
-	ColorRed = 1
-	ColorGreen = 2
-	ColorBlue = 8
+	ColorRed Color = 1
+	ColorGreen Color = 2
+	ColorBlue Color = 8
 )
 
-var EnumNamesColor = map[int]string{
+var EnumNamesColor = map[Color]string{
 	ColorRed:"Red",
 	ColorGreen:"Green",
 	ColorBlue:"Blue",
diff --git a/tests/go_test.go b/tests/go_test.go
index 1eb92a0..946ba04 100644
--- a/tests/go_test.go
+++ b/tests/go_test.go
@@ -1356,31 +1356,28 @@ func CheckFinishedBytesError(fail func(string, ...interface{})) {
 
 // CheckEnumNames checks that the generated enum names are correct.
 func CheckEnumNames(fail func(string, ...interface{})) {
-	type testEnumNames struct {
-		EnumNames map[int]string
-		Expected  map[int]string
-	}
-	data := [...]testEnumNames{
-		{example.EnumNamesAny,
-			map[int]string{
-				example.AnyNONE:                    "NONE",
-				example.AnyMonster:                 "Monster",
-				example.AnyTestSimpleTableWithEnum: "TestSimpleTableWithEnum",
-			},
-		},
-		{example.EnumNamesColor,
-			map[int]string{
-				example.ColorRed:   "Red",
-				example.ColorGreen: "Green",
-				example.ColorBlue:  "Blue",
-			},
-		},
+	{
+
+		want := map[example.Any]string{
+			example.AnyNONE:                    "NONE",
+			example.AnyMonster:                 "Monster",
+			example.AnyTestSimpleTableWithEnum: "TestSimpleTableWithEnum",
+			example.AnyMyGame_Example2_Monster: "MyGame_Example2_Monster",
+		}
+		got := example.EnumNamesAny
+		if !reflect.DeepEqual(got, want) {
+			fail("enum name is not equal")
+		}
 	}
-	for _, t := range data {
-		for val, name := range t.Expected {
-			if name != t.EnumNames[val] {
-				fail("enum name is not equal")
-			}
+	{
+		want := map[example.Color]string{
+			example.ColorRed:   "Red",
+			example.ColorGreen: "Green",
+			example.ColorBlue:  "Blue",
+		}
+		got := example.EnumNamesColor
+		if !reflect.DeepEqual(got, want) {
+			fail("enum name is not equal")
 		}
 	}
 }
-- 
2.7.4