From eb047b12a76cb36b486d9cfa5888f12f215b74ec Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Sun, 2 Aug 2020 14:42:29 +0300
Subject: [PATCH] compiler: support creating pointers to struct

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 pkg/compiler/codegen.go      | 39 +++++++++++++++++++++++++++++++-----
 pkg/compiler/pointer_test.go | 18 +++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 pkg/compiler/pointer_test.go

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index fc34b7b04..2691ff35e 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -438,7 +438,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 					c.emitStoreVar(t.X.(*ast.Ident).Name, t.Sel.Name)
 					return nil
 				}
-				strct, ok := typ.Underlying().(*types.Struct)
+				strct, ok := c.getStruct(typ)
 				if !ok {
 					c.prog.Err = fmt.Errorf("nested selector assigns not supported yet")
 					return nil
@@ -623,7 +623,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 	case *ast.CompositeLit:
 		switch typ := c.typeOf(n).Underlying().(type) {
 		case *types.Struct:
-			c.convertStruct(n)
+			c.convertStruct(n, false)
 		case *types.Map:
 			c.convertMap(n)
 		default:
@@ -835,7 +835,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 			}
 			return nil
 		}
-		strct, ok := typ.Underlying().(*types.Struct)
+		strct, ok := c.getStruct(typ)
 		if !ok {
 			c.prog.Err = fmt.Errorf("selectors are supported only on structs")
 			return nil
@@ -846,6 +846,19 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		return nil
 
 	case *ast.UnaryExpr:
+		if n.Op == token.AND {
+			// We support only taking address from struct literals.
+			// For identifiers we can't support "taking address" in a general way
+			// because both struct and array are reference types.
+			lit, ok := n.X.(*ast.CompositeLit)
+			if ok {
+				c.convertStruct(lit, true)
+				return nil
+			}
+			c.prog.Err = fmt.Errorf("'&' can be used only with struct literals")
+			return nil
+		}
+
 		ast.Walk(c, n.X)
 		// From https://golang.org/ref/spec#Operators
 		// there can be only following unary operators
@@ -1325,7 +1338,19 @@ func (c *codegen) convertMap(lit *ast.CompositeLit) {
 	}
 }
 
-func (c *codegen) convertStruct(lit *ast.CompositeLit) {
+func (c *codegen) getStruct(typ types.Type) (*types.Struct, bool) {
+	switch t := typ.Underlying().(type) {
+	case *types.Struct:
+		return t, true
+	case *types.Pointer:
+		strct, ok := t.Elem().Underlying().(*types.Struct)
+		return strct, ok
+	default:
+		return nil, false
+	}
+}
+
+func (c *codegen) convertStruct(lit *ast.CompositeLit, ptr bool) {
 	// Create a new structScope to initialize and store
 	// the positions of its variables.
 	strct, ok := c.typeOf(lit).Underlying().(*types.Struct)
@@ -1336,7 +1361,11 @@ func (c *codegen) convertStruct(lit *ast.CompositeLit) {
 
 	emit.Opcode(c.prog.BinWriter, opcode.NOP)
 	emit.Int(c.prog.BinWriter, int64(strct.NumFields()))
-	emit.Opcode(c.prog.BinWriter, opcode.NEWSTRUCT)
+	if ptr {
+		emit.Opcode(c.prog.BinWriter, opcode.NEWARRAY)
+	} else {
+		emit.Opcode(c.prog.BinWriter, opcode.NEWSTRUCT)
+	}
 
 	keyedLit := len(lit.Elts) > 0
 	if keyedLit {
diff --git a/pkg/compiler/pointer_test.go b/pkg/compiler/pointer_test.go
new file mode 100644
index 000000000..090101056
--- /dev/null
+++ b/pkg/compiler/pointer_test.go
@@ -0,0 +1,18 @@
+package compiler_test
+
+import (
+	"math/big"
+	"testing"
+)
+
+func TestAddressOfLiteral(t *testing.T) {
+	src := `package foo
+	type Foo struct { A int	}
+	func Main() int {
+		f := &Foo{}
+		setA(f, 3)
+		return f.A
+	}
+	func setA(s *Foo, a int) { s.A = a }`
+	eval(t, src, big.NewInt(3))
+}