From b2a3cee4517126165a42a247ba9502b9c61fa289 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 6 May 2020 15:39:25 +0300
Subject: [PATCH] compiler: support using `return` in some branches

When `return` is used in one codepath, but not the other,
we do not clean altstack properly. This commit fixes it.
---
 pkg/compiler/analysis.go           | 17 +++++++----------
 pkg/compiler/codegen.go            |  6 ++++--
 pkg/compiler/function_call_test.go | 26 ++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/pkg/compiler/analysis.go b/pkg/compiler/analysis.go
index 90dd5e318..34134c6d5 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 7f1b33eeb..f4586bc60 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -261,8 +261,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))
+	})
+}