From 058958729dacfba4ac2b94be346563a7d636aab1 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 23 Jan 2020 17:06:15 +0300
Subject: [PATCH] compiler: support map literals

---
 pkg/compiler/codegen.go  | 15 +++++++++++++
 pkg/compiler/map_test.go | 46 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 pkg/compiler/map_test.go

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 063bfc922..87ad57291 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -357,6 +357,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 +377,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
@@ -700,6 +704,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.
diff --git a/pkg/compiler/map_test.go b/pkg/compiler/map_test.go
new file mode 100644
index 000000000..0ae70841e
--- /dev/null
+++ b/pkg/compiler/map_test.go
@@ -0,0 +1,46 @@
+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),
+	},
+}
+
+func TestMaps(t *testing.T) {
+	runTestCases(t, mapTestCases)
+}