diff --git a/pkg/compiler/analysis.go b/pkg/compiler/analysis.go index b075d7032..fc2c36067 100644 --- a/pkg/compiler/analysis.go +++ b/pkg/compiler/analysis.go @@ -133,16 +133,13 @@ func (f funcUsage) funcUsed(name string) bool { return ok } -// hasReturnStmt looks if the given FuncDecl has a return statement. -func hasReturnStmt(decl ast.Node) (b bool) { - ast.Inspect(decl, func(node ast.Node) bool { - if _, ok := node.(*ast.ReturnStmt); ok { - b = true - return false - } - return true - }) - return +// lastStmtIsReturn checks if last statement of the declaration was return statement.. +func lastStmtIsReturn(decl *ast.FuncDecl) (b bool) { + if l := len(decl.Body.List); l != 0 { + _, ok := decl.Body.List[l-1].(*ast.ReturnStmt) + return ok + } + return false } func analyzeFuncUsage(pkgs map[*types.Package]*loader.PackageInfo) funcUsage { diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go index 696e96914..da5b085c4 100644 --- a/pkg/compiler/codegen.go +++ b/pkg/compiler/codegen.go @@ -263,8 +263,10 @@ func (c *codegen) convertFuncDecl(file ast.Node, decl *ast.FuncDecl) { ast.Walk(c, decl.Body) - // If this function returns the void (no return stmt) we will cleanup its junk on the stack. - if !hasReturnStmt(decl) { + // If we have reached the end of the function without encountering `return` statement, + // we should clean alt.stack manually. + // This can be the case with void and named-return functions. + if !lastStmtIsReturn(decl) { c.saveSequencePoint(decl.Body) emit.Opcode(c.prog.BinWriter, opcode.FROMALTSTACK) emit.Opcode(c.prog.BinWriter, opcode.DROP) diff --git a/pkg/compiler/function_call_test.go b/pkg/compiler/function_call_test.go index 179e0d136..005707d76 100644 --- a/pkg/compiler/function_call_test.go +++ b/pkg/compiler/function_call_test.go @@ -1,6 +1,7 @@ package compiler_test import ( + "fmt" "math/big" "testing" ) @@ -132,3 +133,28 @@ func TestFunctionWithVoidReturn(t *testing.T) { eval(t, src, big.NewInt(6)) }) } + +func TestFunctionWithVoidReturnBranch(t *testing.T) { + src := ` + package testcase + func Main() int { + x := %t + f(x) + return 2 + } + + func f(x bool) { + if x { + return + } + } + ` + t.Run("ReturnBranch", func(t *testing.T) { + src := fmt.Sprintf(src, true) + eval(t, src, big.NewInt(2)) + }) + t.Run("NoReturn", func(t *testing.T) { + src := fmt.Sprintf(src, false) + eval(t, src, big.NewInt(2)) + }) +}