gen,parser: allow enums as map keys

vlang · Feb 20, 2021 · 5a333b0 · 5a333b0
1 parent 1e71c0e
commit 5a333b0
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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)
+}