From 191a167f42b2e9a84f1c0dbf8ff1a79e01cecec3 Mon Sep 17 00:00:00 2001
From: zakuro <z@kuro.red>
Date: Thu, 29 Apr 2021 14:33:41 +0900
Subject: [PATCH] checker: allow omitting the enum name of map key (#9919)

---
 vlib/v/checker/checker.v          | 4 +++-
 vlib/v/tests/map_enum_keys_test.v | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 5d23ec254d..8b7256f308 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -6109,14 +6109,16 @@ pub fn (mut c Checker) index_expr(mut node ast.IndexExpr) ast.Type {
 			typ = typ.set_nr_muls(0)
 		}
 	} else { // [1]
-		index_type := c.expr(node.index)
 		if typ_sym.kind == .map {
 			info := typ_sym.info as ast.Map
+			c.expected_type = info.key_type
+			index_type := c.expr(node.index)
 			if !c.check_types(index_type, info.key_type) {
 				err := c.expected_msg(index_type, info.key_type)
 				c.error('invalid key: $err', node.pos)
 			}
 		} else {
+			index_type := c.expr(node.index)
 			c.check_index(typ_sym, node.index, index_type, node.pos, false)
 		}
 		value_type := c.table.value_type(typ)
diff --git a/vlib/v/tests/map_enum_keys_test.v b/vlib/v/tests/map_enum_keys_test.v
index 8016a57255..a83874a60a 100644
--- a/vlib/v/tests/map_enum_keys_test.v
+++ b/vlib/v/tests/map_enum_keys_test.v
@@ -6,10 +6,10 @@ enum Token {
 
 fn test_map_with_enum_keys() {
 	mut m := map[Token]string{}
-	m[Token.aa] = 'abc'
+	m[.aa] = 'abc'
 	m[Token.bb] = 'def'
 	assert m[Token.aa] == 'abc'
-	assert m[Token.bb] == 'def'
+	assert m[.bb] == 'def'
 	//
 	s := '$m'
 	assert s == "{aa: 'abc', bb: 'def'}"