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) +}