Detect potentially blocking return statements and correctly resume.

Fixes gopherjs#603.

The fix consists of two parts:

 1. If a function has a deferred call, we assume that call might be
    blocking and therefore return statements in the function may behave
    like a blocking call even though the returned expression may not be
    blocking itself.

 2. When generating code for return statements that were marked as
    blocking we add a resumption point to make sure that before
    resuming the deferred function we re-return the correct value.

compiler/analysis/info.go

@@ -1,9 +1,11 @@
 package analysis
 
 import (
+	"fmt"
 	"go/ast"
 	"go/token"
 	"go/types"
+	"strings"
 
 	"github.com/gopherjs/gopherjs/compiler/astutil"
 	"github.com/gopherjs/gopherjs/compiler/typesutil"
@@ -32,6 +34,19 @@ func (src astPath) copy() astPath {
 	return dst
 }
 
+func (ap astPath) String() string {
+	s := &strings.Builder{}
+	s.WriteString("[")
+	for i, n := range ap {
+		if i > 0 {
+			s.WriteString(", ")
+		}
+		fmt.Fprintf(s, "%T(%p)", n, n)
+	}
+	s.WriteString("]")
+	return s.String()
+}
+
 type Info struct {
 	*types.Info
 	Pkg           *types.Package
@@ -88,6 +103,18 @@ func AnalyzePkg(files []*ast.File, fileSet *token.FileSet, typesInfo *types.Info
 		ast.Walk(info.InitFuncInfo, file)
 	}
 
+	for _, funcInfo := range info.allInfos {
+		if !funcInfo.HasDefer {
+			continue
+		}
+		// Conservatively assume that if a function has a deferred call, it might be
+		// blocking, and therefore all return statements need to be treated as
+		// blocking.
+		for _, returnStmt := range funcInfo.returnStmts {
+			funcInfo.markBlocking(returnStmt)
+		}
+	}
+
 	// Propagate information about blocking calls to the caller functions.
 	for {
 		done := true
@@ -136,6 +163,8 @@ type FuncInfo struct {
 	GotoLabel map[*types.Label]bool
 	// List of continue statements in the function.
 	continueStmts []continueStmt
+	// List of return statements in the function.
+	returnStmts []astPath
 	// List of other functions from the current package this function calls. If
 	// any of them are blocking, this function will become blocking too.
 	localCallees map[*types.Func][]astPath
@@ -239,6 +268,11 @@ func (fi *FuncInfo) Visit(node ast.Node) ast.Visitor {
 			ast.Walk(fi, funcLit.Body)
 		}
 		return fi
+	case *ast.ReturnStmt:
+		// Capture all return statements in the function. They could become blocking
+		// if the function has a blocking deferred call.
+		fi.returnStmts = append(fi.returnStmts, fi.visitorStack.copy())
+		return fi
 	default:
 		return fi
 	}

compiler/statements.go

@@ -328,11 +328,39 @@ func (fc *funcContext) translateStmt(stmt ast.Stmt, label *types.Label) {
 			results = fc.resultNames
 		}
 		rVal := fc.translateResults(results)
-		if len(fc.Flattened) != 0 {
+
+		if len(fc.Flattened) == 0 {
+			// The function is not flattened and we don't have to worry about
+			// resumption. A plain return statement is sufficient.
+			fc.Printf("return%s;", rVal)
+			return
+		}
+		if !fc.Blocking[s] {
+			// The function is flattened, but the return statement is non-blocking
+			// (i.e. doesn't lead to blocking deferred calls). A regular return
+			// is sufficient, but we also make sure to not resume function body.
+			// TODO(nevkontakte): Is this actually necessary? If neither a returned
+			// expression or deferred functions are blocking, we will never have to
+			// resume and $s is irrelevant here.
 			fc.Printf("$s = -1; return%s;", rVal)
 			return
 		}
-		fc.Printf("return%s;", rVal)
+
+		if rVal != "" {
+			// If returned expression is non empty, evaluate and store it in a
+			// variable to avoid double-execution in case a deferred function blocks.
+			rVar := fc.newVariable("$r")
+			fc.Printf("%s =%s;", rVar, rVal)
+			rVal = " " + rVar
+		}
+
+		// If deferred function is blocking, we need to re-execute return statement
+		// upon resumption to make sure the returned value is not lost.
+		// See: https://github.com/gopherjs/gopherjs/issues/603.
+		nextCase := fc.caseCounter
+		fc.caseCounter++
+		fc.Printf("$s = %[1]d; case %[1]d: return%[2]s;", nextCase, rVal)
+		return
 
 	case *ast.DeferStmt:
 		isBuiltin := false

tests/goroutine_test.go

@@ -202,7 +202,7 @@ func withBlockingDeferral() int {
 }
 
 func TestReturnWithBlockingDefer(t *testing.T) {
-	t.Skip("https://github.com/gopherjs/gopherjs/issues/603")
+	// t.Skip("https://github.com/gopherjs/gopherjs/issues/603")
 	counter = 0
 
 	got := withBlockingDeferral()