From 5a333b0fdc9864fd27b9cfaa29e68eae8209e839 Mon Sep 17 00:00:00 2001
From: Delyan Angelov <delian66@gmail.com>
Date: Sat, 20 Feb 2021 19:39:25 +0200
Subject: [PATCH] gen,parser: allow enums as map keys

---
 vlib/v/gen/c/cgen.v               |  2 +-
 vlib/v/parser/parse_type.v        |  8 +++++---
 vlib/v/tests/map_enum_keys_test.v | 17 +++++++++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)
 create mode 100644 vlib/v/tests/map_enum_keys_test.v

diff --git a/vlib/v/gen/c/cgen.v b/vlib/v/gen/c/cgen.v
index ca10468a1f..e52312e326 100644
--- a/vlib/v/gen/c/cgen.v
+++ b/vlib/v/gen/c/cgen.v
@@ -2447,7 +2447,7 @@ fn (mut g Gen) map_fn_ptrs(key_typ table.TypeSymbol) (string, string, string, st
 			key_eq_fn = '&map_eq_int_2'
 			clone_fn = '&map_clone_int_2'
 		}
-		.int, .u32, .rune, .f32 {
+		.int, .u32, .rune, .f32, .enum_ {
 			hash_fn = '&map_hash_int_4'
 			key_eq_fn = '&map_eq_int_4'
 			clone_fn = '&map_clone_int_4'
diff --git a/vlib/v/parser/parse_type.v b/vlib/v/parser/parse_type.v
index e35117f583..61dba2b1e4 100644
--- a/vlib/v/parser/parse_type.v
+++ b/vlib/v/parser/parse_type.v
@@ -72,7 +72,8 @@ pub fn (mut p Parser) parse_map_type() table.Type {
 	}
 	p.check(.lsbr)
 	key_type := p.parse_type()
-	is_alias := p.table.get_type_symbol(key_type).kind == .alias
+	key_sym := p.table.get_type_symbol(key_type)
+	is_alias := key_sym.kind == .alias
 	if key_type.idx() == 0 {
 		// error is reported in parse_type
 		return 0
@@ -83,9 +84,10 @@ pub fn (mut p Parser) parse_map_type() table.Type {
 		return 0
 	}
 	if !(key_type in [table.string_type_idx, table.voidptr_type_idx]
-		|| ((key_type.is_int() || key_type.is_float() || is_alias) && !key_type.is_ptr())) {
+		|| key_sym.kind == .enum_ || ((key_type.is_int() || key_type.is_float()
+		|| is_alias) && !key_type.is_ptr())) {
 		s := p.table.type_to_str(key_type)
-		p.error_with_pos('maps only support string, integer, float, rune or voidptr keys for now (not `$s`)',
+		p.error_with_pos('maps only support string, integer, float, rune, enum or voidptr keys for now (not `$s`)',
 			p.tok.position())
 		return 0
 	}
diff --git a/vlib/v/tests/map_enum_keys_test.v b/vlib/v/tests/map_enum_keys_test.v
new file mode 100644
index 0000000000..8016a57255
--- /dev/null
+++ b/vlib/v/tests/map_enum_keys_test.v
@@ -0,0 +1,17 @@
+enum Token {
+	aa = 2
+	bb
+	cc
+}
+
+fn test_map_with_enum_keys() {
+	mut m := map[Token]string{}
+	m[Token.aa] = 'abc'
+	m[Token.bb] = 'def'
+	assert m[Token.aa] == 'abc'
+	assert m[Token.bb] == 'def'
+	//
+	s := '$m'
+	assert s == "{aa: 'abc', bb: 'def'}"
+	println(m)
+}