From 1ee4acbdbc7c66bdfd2a43a190075df1dc569a7d Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 29 Jun 2020 16:45:27 +0300
Subject: [PATCH] compiler: manage variables in a separate varScope struct

Abstract var scope management from the funcScope.
---
 pkg/compiler/codegen.go    | 10 +++---
 pkg/compiler/func_scope.go | 19 ++----------
 pkg/compiler/vars.go       | 63 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 20 deletions(-)
 create mode 100644 pkg/compiler/vars.go

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 2f9e6db3e..97ea09176 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -166,10 +166,9 @@ func (c *codegen) emitStoreStructField(i int) {
 // according to current scope.
 func (c *codegen) getVarIndex(name string) (varType, int) {
 	if c.scope != nil {
-		if i, ok := c.scope.arguments[name]; ok {
-			return varArgument, i
-		} else if i, ok := c.scope.locals[name]; ok {
-			return varLocal, i
+		vt, val := c.scope.vars.getVarIndex(name)
+		if val >= 0 {
+			return vt, val
 		}
 	}
 	if i, ok := c.globals[name]; ok {
@@ -311,6 +310,9 @@ func (c *codegen) convertFuncDecl(file ast.Node, decl *ast.FuncDecl) {
 		emit.Instruction(c.prog.BinWriter, opcode.INITSLOT, []byte{byte(sizeLoc), byte(sizeArg)})
 	}
 
+	f.vars.newScope()
+	defer f.vars.dropScope()
+
 	// We need to handle methods, which in Go, is just syntactic sugar.
 	// The method receiver will be passed in as first argument.
 	// We check if this declaration has a receiver and load it into scope.
diff --git a/pkg/compiler/func_scope.go b/pkg/compiler/func_scope.go
index b66c8e1dd..2202c120d 100644
--- a/pkg/compiler/func_scope.go
+++ b/pkg/compiler/func_scope.go
@@ -31,8 +31,7 @@ type funcScope struct {
 	variables []string
 
 	// Local variables
-	locals    map[string]int
-	arguments map[string]int
+	vars varScope
 
 	// voidCalls are basically functions that return their value
 	// into nothing. The stack has their return value but there
@@ -55,8 +54,7 @@ func newFuncScope(decl *ast.FuncDecl, label uint16) *funcScope {
 		name:      name,
 		decl:      decl,
 		label:     label,
-		locals:    map[string]int{},
-		arguments: map[string]int{},
+		vars:      newVarScope(),
 		voidCalls: map[*ast.CallExpr]bool{},
 		variables: []string{},
 		i:         -1,
@@ -139,18 +137,7 @@ func (c *funcScope) stackSize() int64 {
 
 // newVariable creates a new local variable or argument in the scope of the function.
 func (c *funcScope) newVariable(t varType, name string) int {
-	var n int
-	switch t {
-	case varLocal:
-		n = len(c.locals)
-		c.locals[name] = n
-	case varArgument:
-		n = len(c.arguments)
-		c.arguments[name] = n
-	default:
-		panic("invalid type")
-	}
-	return n
+	return c.vars.newVariable(t, name)
 }
 
 // newLocal creates a new local variable into the scope of the function.
diff --git a/pkg/compiler/vars.go b/pkg/compiler/vars.go
new file mode 100644
index 000000000..006a7556f
--- /dev/null
+++ b/pkg/compiler/vars.go
@@ -0,0 +1,63 @@
+package compiler
+
+type varScope struct {
+	localsCnt int
+	argCnt    int
+	arguments map[string]int
+	locals    []map[string]int
+}
+
+func newVarScope() varScope {
+	return varScope{
+		arguments: make(map[string]int),
+	}
+}
+
+func (c *varScope) newScope() {
+	c.locals = append(c.locals, map[string]int{})
+}
+
+func (c *varScope) dropScope() {
+	c.locals = c.locals[:len(c.locals)-1]
+}
+
+func (c *varScope) getVarIndex(name string) (varType, int) {
+	for i := len(c.locals) - 1; i >= 0; i-- {
+		if i, ok := c.locals[i][name]; ok {
+			return varLocal, i
+		}
+	}
+	if i, ok := c.arguments[name]; ok {
+		return varArgument, i
+	}
+	return 0, -1
+}
+
+// newVariable creates a new local variable or argument in the scope of the function.
+func (c *varScope) newVariable(t varType, name string) int {
+	var n int
+	switch t {
+	case varLocal:
+		return c.newLocal(name)
+	case varArgument:
+		_, ok := c.arguments[name]
+		if ok {
+			panic("argument is already allocated")
+		}
+		n = len(c.arguments)
+		c.arguments[name] = n
+	default:
+		panic("invalid type")
+	}
+	return n
+}
+
+// newLocal creates a new local variable in the current scope.
+func (c *varScope) newLocal(name string) int {
+	idx := len(c.locals) - 1
+	m := c.locals[idx]
+	m[name] = c.localsCnt
+	c.localsCnt++
+	c.locals[idx] = m
+	return c.localsCnt - 1
+}