From da4b6b934d9461d10dbf7c3bf52f8d8e83144d68 Mon Sep 17 00:00:00 2001
From: yuyi <yuyi98@163.com>
Date: Tue, 22 Jun 2021 00:24:42 +0800
Subject: [PATCH] checker: fix map assign array of interface values (#10532)

---
 vlib/v/checker/checker.v                      | 35 +++++++++++-----
 vlib/v/checker/tests/map_init_wrong_type.out  |  4 +-
 .../map_assign_array_of_interface_test.v      | 40 +++++++++++++++++++
 3 files changed, 66 insertions(+), 13 deletions(-)
 create mode 100644 vlib/v/tests/map_assign_array_of_interface_test.v

diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 3200d97392..962d337b8e 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -6872,30 +6872,43 @@ pub fn (mut c Checker) map_init(mut node ast.MapInit) ast.Type {
 		return node.typ
 	}
 	if node.keys.len > 0 && node.vals.len > 0 {
-		// `{'age': 20}`
-		mut key0_type := c.table.mktyp(c.expr(node.keys[0]))
-		if node.keys[0].is_auto_deref_var() {
-			key0_type = key0_type.deref()
-		}
-		mut val0_type := c.table.mktyp(c.expr(node.vals[0]))
-		if node.vals[0].is_auto_deref_var() {
-			val0_type = val0_type.deref()
+		mut key0_type := ast.void_type
+		mut val0_type := ast.void_type
+		use_expected_type := c.expected_type != ast.void_type && !c.inside_const
+			&& c.table.get_type_symbol(c.expected_type).kind == .map
+		if use_expected_type {
+			sym := c.table.get_type_symbol(c.expected_type)
+			info := sym.map_info()
+			key0_type = c.unwrap_generic(info.key_type)
+			val0_type = c.unwrap_generic(info.value_type)
+		} else {
+			// `{'age': 20}`
+			key0_type = c.table.mktyp(c.expr(node.keys[0]))
+			if node.keys[0].is_auto_deref_var() {
+				key0_type = key0_type.deref()
+			}
+			val0_type = c.table.mktyp(c.expr(node.vals[0]))
+			if node.vals[0].is_auto_deref_var() {
+				val0_type = val0_type.deref()
+			}
 		}
 		mut same_key_type := true
 		for i, key in node.keys {
-			if i == 0 {
+			if i == 0 && !use_expected_type {
 				continue
 			}
 			val := node.vals[i]
 			key_type := c.expr(key)
 			c.expected_type = val0_type
 			val_type := c.expr(val)
-			if !c.check_types(key_type, key0_type) {
+			if !c.check_types(key_type, key0_type) || (i == 0 && key_type.is_number()
+				&& key0_type.is_number() && key0_type != c.table.mktyp(key_type)) {
 				msg := c.expected_msg(key_type, key0_type)
 				c.error('invalid map key: $msg', key.position())
 				same_key_type = false
 			}
-			if !c.check_types(val_type, val0_type) {
+			if !c.check_types(val_type, val0_type) || (i == 0 && val_type.is_number()
+				&& val0_type.is_number() && val0_type != c.table.mktyp(val_type)) {
 				msg := c.expected_msg(val_type, val0_type)
 				c.error('invalid map value: $msg', val.position())
 			}
diff --git a/vlib/v/checker/tests/map_init_wrong_type.out b/vlib/v/checker/tests/map_init_wrong_type.out
index 283edf7bbd..579f53ef28 100644
--- a/vlib/v/checker/tests/map_init_wrong_type.out
+++ b/vlib/v/checker/tests/map_init_wrong_type.out
@@ -1,8 +1,8 @@
-vlib/v/checker/tests/map_init_wrong_type.vv:3:8: error: cannot assign to `a`: expected `map[string]f32`, not `map[string]f64`
+vlib/v/checker/tests/map_init_wrong_type.vv:3:15: error: invalid map value: expected `f32`, not `float literal`
     1 | fn main() {
     2 |    mut a := map[string]f32{}
     3 |    a = { 'x': 12.3 }
-      |        ~~~~~~~~~~~~~
+      |               ~~~~
     4 |    _ = {2:0 3:0 "hi":0}
     5 |    _ = {2:0 3:`@` 4:0}
 vlib/v/checker/tests/map_init_wrong_type.vv:4:17: error: invalid map key: expected `int`, not `string`
diff --git a/vlib/v/tests/map_assign_array_of_interface_test.v b/vlib/v/tests/map_assign_array_of_interface_test.v
new file mode 100644
index 0000000000..de3f89fb85
--- /dev/null
+++ b/vlib/v/tests/map_assign_array_of_interface_test.v
@@ -0,0 +1,40 @@
+interface Animal {
+	say(message string) ?
+}
+
+struct Cat {
+	name string
+}
+
+struct Dog {
+	name string
+}
+
+pub fn (cat Cat) say(message string) ? {
+	println('$message, meow')
+}
+
+pub fn (dog Dog) say(message string) ? {
+	println('$message, wooff')
+}
+
+fn test_map_assign_array_of_interface() {
+	mut owner_and_animals := map[string][]Animal{}
+
+	owner_and_animals = map{
+		'John Doe': [
+			Cat{
+				name: 'Bobby'
+			},
+			Dog{
+				name: 'Hulk'
+			},
+		]
+	}
+	println(owner_and_animals)
+	assert owner_and_animals['John Doe'].len == 2
+	println(owner_and_animals['John Doe'][0])
+	assert '${owner_and_animals['John Doe'][0]}'.contains("name: 'Bobby'")
+	println(owner_and_animals['John Doe'][1])
+	assert '${owner_and_animals['John Doe'][1]}'.contains("name: 'Hulk'")
+}