compiler: count locals number properly

pkg/compiler/analysis.go

@@ -55,14 +55,14 @@ func (c *codegen) traverseGlobals() (int, int, int) {
 				switch n := node.(type) {
 				case *ast.FuncDecl:
 					if isInitFunc(n) {
-						c, _ := countLocals(n)
-						if c > initLocals {
-							initLocals = c
+						num, _ := c.countLocals(n)
+						if num > initLocals {
+							initLocals = num
 						}
 					} else if isDeployFunc(n) {
-						c, _ := countLocals(n)
-						if c > deployLocals {
-							deployLocals = c
+						num, _ := c.countLocals(n)
+						if num > deployLocals {
+							deployLocals = num
 						}
 					}
 					return !hasDefer

pkg/compiler/codegen.go

@@ -330,7 +330,7 @@ func (c *codegen) convertInitFuncs(f *ast.File, pkg *types.Package, seenBefore b
 		case *ast.FuncDecl:
 			if isInitFunc(n) {
 				if seenBefore {
-					cnt, _ := countLocals(n)
+					cnt, _ := c.countLocals(n)
 					c.clearSlots(cnt)
 					seenBefore = true
 				}
@@ -362,7 +362,7 @@ func (c *codegen) convertDeployFuncs() {
 			case *ast.FuncDecl:
 				if isDeployFunc(n) {
 					if seenBefore {
-						cnt, _ := countLocals(n)
+						cnt, _ := c.countLocals(n)
 						c.clearSlots(cnt)
 					}
 					c.convertFuncDecl(f, n, pkg)
@@ -408,7 +408,7 @@ func (c *codegen) convertFuncDecl(file ast.Node, decl *ast.FuncDecl, pkg *types.
 	// All globals copied into the scope of the function need to be added
 	// to the stack size of the function.
 	if !isInit && !isDeploy {
-		sizeLoc := f.countLocals()
+		sizeLoc := c.countLocalsWithDefer(f)
 		if sizeLoc > 255 {
 			c.prog.Err = errors.New("maximum of 255 local variables is allowed")
 		}
@@ -416,7 +416,6 @@ func (c *codegen) convertFuncDecl(file ast.Node, decl *ast.FuncDecl, pkg *types.
 		if sizeArg > 255 {
 			c.prog.Err = errors.New("maximum of 255 local variables is allowed")
 		}
-		sizeLoc = 255 // FIXME count locals including inline variables
 		if sizeLoc != 0 || sizeArg != 0 {
 			emit.Instruction(c.prog.BinWriter, opcode.INITSLOT, []byte{byte(sizeLoc), byte(sizeArg)})
 		}

pkg/compiler/func_scope.go

@@ -102,11 +102,47 @@ func (c *funcScope) analyzeVoidCalls(node ast.Node) bool {
 	return true
 }
 
-func countLocals(decl *ast.FuncDecl) (int, bool) {
+func (c *codegen) countLocals(decl *ast.FuncDecl) (int, bool) {
+	return c.countLocalsInline(decl, nil, nil)
+}
+
+func (c *codegen) countLocalsInline(decl *ast.FuncDecl, pkg *types.Package, f *funcScope) (int, bool) {
+	oldMap := c.importMap
+	if pkg != nil {
+		c.fillImportMap(f.file, pkg)
+	}
+
 	size := 0
 	hasDefer := false
 	ast.Inspect(decl, func(n ast.Node) bool {
 		switch n := n.(type) {
+		case *ast.CallExpr:
+			var name string
+			switch fun := n.Fun.(type) {
+			case *ast.Ident:
+				var pkgName string
+				if pkg != nil {
+					pkgName = pkg.Path()
+				}
+				name = c.getIdentName(pkgName, fun.Name)
+			case *ast.SelectorExpr:
+				name, _ = c.getFuncNameFromSelector(fun)
+			default:
+				return false
+			}
+			if inner, ok := c.funcs[name]; ok && canInline(name) {
+				for i := range n.Args {
+					switch n.Args[i].(type) {
+					case *ast.Ident:
+					case *ast.BasicLit:
+					default:
+						size++
+					}
+				}
+				innerSz, _ := c.countLocalsInline(inner.decl, inner.pkg, inner)
+				size += innerSz
+			}
+			return false
 		case *ast.FuncType:
 			num := n.Results.NumFields()
 			if num != 0 && len(n.Results.List[0].Names) != 0 {
@@ -119,7 +155,11 @@ func countLocals(decl *ast.FuncDecl) (int, bool) {
 		case *ast.DeferStmt:
 			hasDefer = true
 			return false
-		case *ast.ReturnStmt, *ast.IfStmt:
+		case *ast.ReturnStmt:
+			if pkg == nil {
+				size++
+			}
+		case *ast.IfStmt:
 			size++
 		// This handles the inline GenDecl like "var x = 2"
 		case *ast.ValueSpec:
@@ -136,13 +176,16 @@ func countLocals(decl *ast.FuncDecl) (int, bool) {
 		}
 		return true
 	})
+	if pkg != nil {
+		c.importMap = oldMap
+	}
 	return size, hasDefer
 }
 
-func (c *funcScope) countLocals() int {
-	size, hasDefer := countLocals(c.decl)
+func (c *codegen) countLocalsWithDefer(f *funcScope) int {
+	size, hasDefer := c.countLocals(f.decl)
 	if hasDefer {
-		c.finallyProcessedIndex = size
+		f.finallyProcessedIndex = size
 		size++
 	}
 	return size
@@ -156,12 +199,6 @@ func (c *funcScope) countArgs() int {
 	return n
 }
 
-func (c *funcScope) stackSize() int64 {
-	size := c.countLocals()
-	numArgs := c.countArgs()
-	return int64(size + numArgs)
-}
-
 // newVariable creates a new local variable or argument in the scope of the function.
 func (c *funcScope) newVariable(t varType, name string) int {
 	return c.vars.newVariable(t, name)