Merge pull request #620 from nspcc-dev/feature/map

compiler: support map literals

pkg/compiler/codegen.go

@@ -61,30 +61,34 @@ func (c *codegen) emitLoadConst(t types.TypeAndValue) {
 	}
 	switch typ := t.Type.Underlying().(type) {
 	case *types.Basic:
-		switch typ.Kind() {
-		case types.Int, types.UntypedInt, types.Uint:
-			val, _ := constant.Int64Val(t.Value)
-			emitInt(c.prog.BinWriter, val)
-		case types.String, types.UntypedString:
-			val := constant.StringVal(t.Value)
-			emitString(c.prog.BinWriter, val)
-		case types.Bool, types.UntypedBool:
-			val := constant.BoolVal(t.Value)
-			emitBool(c.prog.BinWriter, val)
-		case types.Byte:
-			val, _ := constant.Int64Val(t.Value)
-			b := byte(val)
-			emitBytes(c.prog.BinWriter, []byte{b})
-		default:
-			c.prog.Err = fmt.Errorf("compiler doesn't know how to convert this basic type: %v", t)
-			return
-		}
+		c.convertBasicType(t, typ)
 	default:
 		c.prog.Err = fmt.Errorf("compiler doesn't know how to convert this constant: %v", t)
 		return
 	}
 }
 
+func (c *codegen) convertBasicType(t types.TypeAndValue, typ *types.Basic) {
+	switch typ.Kind() {
+	case types.Int, types.UntypedInt, types.Uint:
+		val, _ := constant.Int64Val(t.Value)
+		emitInt(c.prog.BinWriter, val)
+	case types.String, types.UntypedString:
+		val := constant.StringVal(t.Value)
+		emitString(c.prog.BinWriter, val)
+	case types.Bool, types.UntypedBool:
+		val := constant.BoolVal(t.Value)
+		emitBool(c.prog.BinWriter, val)
+	case types.Byte:
+		val, _ := constant.Int64Val(t.Value)
+		b := byte(val)
+		emitBytes(c.prog.BinWriter, []byte{b})
+	default:
+		c.prog.Err = fmt.Errorf("compiler doesn't know how to convert this basic type: %v", t)
+		return
+	}
+}
+
 func (c *codegen) emitLoadLocal(name string) {
 	pos := c.scope.loadLocal(name)
 	if pos < 0 {
@@ -357,6 +361,8 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 			typ = c.typeInfo.ObjectOf(t).Type().Underlying()
 		case *ast.SelectorExpr:
 			typ = c.typeInfo.ObjectOf(t.Sel).Type().Underlying()
+		case *ast.MapType:
+			typ = c.typeInfo.TypeOf(t)
 		default:
 			ln := len(n.Elts)
 			// ByteArrays needs a different approach than normal arrays.
@@ -375,6 +381,8 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		switch typ.(type) {
 		case *types.Struct:
 			c.convertStruct(n)
+		case *types.Map:
+			c.convertMap(n)
 		}
 
 		return nil
@@ -570,12 +578,19 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		switch n.Index.(type) {
 		case *ast.BasicLit:
 			t := c.typeInfo.Types[n.Index]
-			val, _ := constant.Int64Val(t.Value)
-			c.emitLoadField(int(val))
+			switch typ := t.Type.Underlying().(type) {
+			case *types.Basic:
+				c.convertBasicType(t, typ)
+			default:
+				c.prog.Err = fmt.Errorf("compiler can't use following type as an index: %T", typ)
+				return nil
+			}
 		default:
 			ast.Walk(c, n.Index)
-			emitOpcode(c.prog.BinWriter, opcode.PICKITEM) // just pickitem here
 		}
+
+		emitOpcode(c.prog.BinWriter, opcode.PICKITEM) // just pickitem here
+
 		return nil
 
 	case *ast.ForStmt:
@@ -700,6 +715,17 @@ func (c *codegen) convertByteArray(lit *ast.CompositeLit) {
 	emitBytes(c.prog.BinWriter, buf)
 }
 
+func (c *codegen) convertMap(lit *ast.CompositeLit) {
+	emitOpcode(c.prog.BinWriter, opcode.NEWMAP)
+	for i := range lit.Elts {
+		elem := lit.Elts[i].(*ast.KeyValueExpr)
+		emitOpcode(c.prog.BinWriter, opcode.DUP)
+		ast.Walk(c, elem.Key)
+		ast.Walk(c, elem.Value)
+		emitOpcode(c.prog.BinWriter, opcode.SETITEM)
+	}
+}
+
 func (c *codegen) convertStruct(lit *ast.CompositeLit) {
 	// Create a new structScope to initialize and store
 	// the positions of its variables.

pkg/compiler/map_test.go

@@ -0,0 +1,62 @@
+package compiler_test
+
+import (
+	"math/big"
+	"testing"
+)
+
+var mapTestCases = []testCase{
+	{
+		"map composite literal",
+		`
+		package foo
+		func Main() int {
+			t := map[int]int{
+				1: 6,
+				2: 9,
+			}
+
+			age := t[2]
+			return age
+		}
+		`,
+		big.NewInt(9),
+	},
+	{
+		"nested map",
+		`
+		package foo
+		func Main() int {
+		t := map[int]map[int]int{
+			1: map[int]int{2: 5, 3: 1},
+			2: nil,
+			5: map[int]int{3: 4, 7: 2},
+		}
+
+		x := t[5][3]
+		return x
+	}
+	`,
+		big.NewInt(4),
+	},
+	{
+		"map with string index",
+		`
+		package foo
+		func Main() string {
+			t := map[string]string{
+				"name": "Valera",
+				"age": "33",
+			}
+
+			name := t["name"]
+			return name
+		}
+		`,
+		[]byte("Valera"),
+	},
+}
+
+func TestMaps(t *testing.T) {
+	runTestCases(t, mapTestCases)
+}