Improving return with defers blocking

compiler/internal/analysis/defer.go

@@ -37,7 +37,7 @@ func newInstDefer(inst typeparams.Instance) *deferStmt {
 func newLitDefer(lit *ast.FuncLit) *deferStmt {
 	return &deferStmt{
 		isBlocking: func(info *Info) (bool, bool) {
-			return info.FuncLitInfo(lit).IsBlocking(), false
+			return info.FuncLitInfo(lit).IsBlocking(), true
 		},
 	}
 }

compiler/internal/analysis/info_test.go

@@ -302,49 +302,131 @@ func TestBlocking_Defers_WithReturns_WithNamedFuncs(t *testing.T) {
 
 		func blockingBody(c chan bool) int {
 			defer blockingPrint(c)
-			return 42
+			return 42 // line 13
 		}
 
 		func blockingArg(c chan bool) int {
 			defer nonBlockingPrint(<-c)
-			return 42
+			return 42 // line 18
 		}
 
 		func notBlocking(c chan bool) int {
 			defer nonBlockingPrint(true)
-			return 42
+			return 42 // line 23
 		}`)
 	bt.assertBlocking(`blockingPrint`)
 	bt.assertNotBlocking(`nonBlockingPrint`)
 
 	bt.assertBlocking(`blockingBody`)
+	bt.assertBlockingReturn(13)
+
 	bt.assertBlocking(`blockingArg`)
+	// The defer is non-blocking so the return is not blocking
+	// even though the function is blocking.
+	bt.assertNotBlockingReturn(18)
 
 	bt.assertNotBlocking(`notBlocking`)
+	bt.assertNotBlockingReturn(23)
 }
 
 func TestBlocking_Defers_WithMultipleReturns(t *testing.T) {
+	bt := newBlockingTest(t,
+		`package test
+
+		func foo(c chan int) bool {
+			defer func() { // line 4
+				if r := recover(); r != nil {
+					println("Error", r)
+				}
+			}()
+
+			if c == nil {
+				return false // line 11
+			}
+
+			defer func(v int) { // line 14
+				println(v)
+			}(<-c)
+
+			value := <-c
+			if value < 0 {
+				return false // line 20
+			}
+			
+			if value > 0 {
+				defer func() { // line 24
+					println(<-c)
+				}()
+
+				return false // line 28
+			}
 
+			return true // line 31
+		}`)
+	bt.assertBlocking(`foo`)
+	bt.assertNotBlockingLit(4)
+	// Early escape from function without blocking defers is not blocking.
+	bt.assertNotBlockingReturn(11)
+	bt.assertNotBlockingLit(14)
+	// Function has had blocking by this point but no blocking defers yet.
+	bt.assertNotBlockingReturn(20)
+	bt.assertBlockingLit(24)
+	// The return is blocking because of a blocking defer.
+	bt.assertBlockingReturn(28)
+	// Technically the return on line 31 is not blocking since the defer that
+	// is blocking can only exit through the return on line 28, but it would be
+	// difficult to determine which defers would only affect certain returns
+	// without doing full control flow analysis.
+	//
+	// TODO(grantnelson-wf): We could fix this at some point by keeping track
+	// of which flow control statements (e.g. if-statements) are terminating
+	// or not. Any defers added in a terminating control flow would not
+	// propagate to returns that are not in that block.
+	//
+	// For now we simply build up the list of defers as we go making
+	// the return on line 31 also blocking.
+	bt.assertBlockingReturn(31)
 }
 
 func TestBlocking_Returns_WithoutDefers(t *testing.T) {
 	bt := newBlockingTest(t,
 		`package test
 
 		func blocking(c chan bool) bool {
-			return <-c
+			return <-c // line 4
+		}
+
+		func blockingBeforeReturn(c chan bool) bool {
+			v := <-c
+			return v // line 9
 		}
 
 		func indirectlyBlocking(c chan bool) bool {
-			return blocking(c)
+			return blocking(c) // line 13
+		}
+
+		func indirectlyBlockingBeforeReturn(c chan bool) bool {
+			v := blocking(c)
+			return v // line 18
 		}
 
 		func notBlocking(c chan bool) bool {
-			return true
+			return true // line 22
 		}`)
 	bt.assertBlocking(`blocking`)
+	bt.assertBlockingReturn(4)
+
+	bt.assertBlocking(`blockingBeforeReturn`)
+	bt.assertNotBlockingReturn(9)
+
 	bt.assertBlocking(`indirectlyBlocking`)
+	bt.assertBlockingReturn(13)
+
+	bt.assertBlocking(`indirectlyBlockingBeforeReturn`)
+	bt.assertNotBlockingReturn(18)
+
 	bt.assertNotBlocking(`notBlocking`)
+	bt.assertNotBlockingReturn(22)
 }
 
 func TestBlocking_FunctionLiteral(t *testing.T) {
@@ -1216,7 +1298,7 @@ func (bt *blockingTest) assertNotBlockingLit(lineNo int) {
 func (bt *blockingTest) isFuncLitBlocking(lineNo int) bool {
 	fnLit := getNodeAtLineNo[*ast.FuncLit](bt.file, bt.f.FileSet, lineNo)
 	if fnLit == nil {
-		bt.f.T.Fatalf(`FuncLit found on line %d not found in the AST.`, lineNo)
+		bt.f.T.Fatalf(`FuncLit on line %d not found in the AST.`, lineNo)
 	}
 	return bt.pkgInfo.FuncLitInfo(fnLit).IsBlocking()
 }
@@ -1261,27 +1343,42 @@ func (bt *blockingTest) assertNotBlockingReturn(lineNo int) {
 }
 
 func (bt *blockingTest) isReturnBlocking(lineNo int) bool {
-	ret := getNodeAtLineNo[*ast.FuncLit](bt.file, bt.f.FileSet, lineNo)
+	ret := getNodeAtLineNo[*ast.ReturnStmt](bt.file, bt.f.FileSet, lineNo)
 	if ret == nil {
-		bt.f.T.Fatalf(`FuncLit found on line %d not found in the AST.`, lineNo)
+		bt.f.T.Fatalf(`ReturnStmt on line %d not found in the AST.`, lineNo)
+	}
+	for _, info := range bt.pkgInfo.allInfos {
+		if blocking, found := info.Blocking[ret]; found {
+			return blocking
+		}
 	}
-	return bt.pkgInfo.(ret).IsBlocking()
+	// If not found in any info.Blocking, then it is not blocking.
+	return false
 }
 
 func getNodeAtLineNo[N ast.Node](file *ast.File, fSet *token.FileSet, lineNo int) N {
 	var node N
-	keepLooking := false
+	keepLooking := true
 	ast.Inspect(file, func(n ast.Node) bool {
+		if n == nil {
+			return false
+		}
 		nodeLine := fSet.Position(n.Pos()).Line
-		if nodeLine > lineNo {
-			keepLooking = false
-		} else if nodeLine == lineNo {
+		switch {
+		case nodeLine < lineNo:
+			// We haven't reached the line yet, so check if we can skip over
+			// this whole node or if we should look inside it.
+			return fSet.Position(n.End()).Line >= lineNo
+		case nodeLine > lineNo:
+			// We went past it without finding it, so stop looking.
+			return false
+		default: // nodeLine == lineNo
 			if n, ok := n.(N); ok {
 				node = n
 				keepLooking = false
 			}
+			return keepLooking
 		}
-		return keepLooking
 	})
 	return node
 }