internal/lsp: support return statements in extract function

Previously, users could not extract code that contained a return
statement. Now, users can extract code with return statements, as long
as the statements are nested within an if, case, or other control
flow statement.

Updates golang/go#37170

Change-Id: I2df52d0241517472decabce3666a32392ff257bd
Reviewed-on: https://go-review.googlesource.com/c/tools/+/243650
Run-TryBot: Josh Baum <joshbaum@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

Author: joshbaum

internal/analysisinternal/analysis.go

@@ -281,3 +281,22 @@ func baseIfStmt(path []ast.Node, index int) ast.Stmt {
 	}
 	return stmt.(ast.Stmt)
 }
+
+// WalkASTWithParent walks the AST rooted at n. The semantics are
+// similar to ast.Inspect except it does not call f(nil).
+func WalkASTWithParent(n ast.Node, f func(n ast.Node, parent ast.Node) bool) {
+	var ancestors []ast.Node
+	ast.Inspect(n, func(n ast.Node) (recurse bool) {
+		if n == nil {
+			ancestors = ancestors[:len(ancestors)-1]
+			return false
+		}
+
+		var parent ast.Node
+		if len(ancestors) > 0 {
+			parent = ancestors[len(ancestors)-1]
+		}
+		ancestors = append(ancestors, n)
+		return f(n, parent)
+	})
+}