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>

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)
+	})
+}

internal/lsp/source/extract.go

@@ -37,7 +37,7 @@ func extractVariable(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 	if !ok {
 		return nil, fmt.Errorf("extractVariable: node is not an expression")
 	}
-	name := generateAvailableIdentifier(expr.Pos(), file, path, info)
+	name := generateAvailableIdentifier(node.Pos(), file, path, info, "x", 0)
 
 	// Create new AST node for extracted code.
 	var assignment string
@@ -117,7 +117,19 @@ func calculateIndentation(content []byte, tok *token.File, insertBeforeStmt ast.
 	return string(content[lineOffset:stmtOffset])
 }
 
-// Check for variable collision in scope.
+// generateAvailableIdentifier adjusts the new function name until there are no collisons in scope.
+// Possible collisions include other function and variable names.
+func generateAvailableIdentifier(pos token.Pos, file *ast.File, path []ast.Node, info *types.Info, prefix string, idx int) string {
+	scopes := collectScopes(info, path, pos)
+	name := prefix + fmt.Sprintf("%d", idx)
+	for file.Scope.Lookup(name) != nil || !isValidName(name, scopes) {
+		idx++
+		name = fmt.Sprintf("%v%d", prefix, idx)
+	}
+	return name
+}
+
+// isValidName checks for variable collision in scope.
 func isValidName(name string, scopes []*types.Scope) bool {
 	for _, scope := range scopes {
 		if scope == nil {
@@ -130,6 +142,20 @@ func isValidName(name string, scopes []*types.Scope) bool {
 	return true
 }
 
+// returnVariable keeps track of the information we need to properly introduce a new variable
+// that we will return in the extracted function.
+type returnVariable struct {
+	// name is the identifier that is used on the left-hand side of the call to
+	// the extracted function.
+	name ast.Expr
+	// decl is the declaration of the variable. It is used in the type signature of the
+	// extracted function and for variable declarations.
+	decl *ast.Field
+	// zeroVal is the "zero value" of the type of the variable. It is used in a return
+	// statement in the extracted function.
+	zeroVal ast.Expr
+}
+
 // extractFunction refactors the selected block of code into a new function.
 // It also replaces the selected block of code with a call to the extracted
 // function. First, we manually adjust the selection range. We remove trailing
@@ -148,7 +174,7 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 	if len(path) == 0 {
 		return nil, fmt.Errorf("extractFunction: no path enclosing interval")
 	}
-	// Node that encloses selection must be a statement.
+	// Node that encloses the selection must be a statement.
 	// TODO: Support function extraction for an expression.
 	if _, ok := path[0].(ast.Stmt); !ok {
 		return nil, fmt.Errorf("extractFunction: ast.Node is not a statement")
@@ -161,7 +187,7 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 	if pkgScope == nil {
 		return nil, fmt.Errorf("extractFunction: package scope is empty")
 	}
-	// Find function enclosing the selection.
+	// Find the function declaration that encloses the selection.
 	var outer *ast.FuncDecl
 	for _, p := range path {
 		if p, ok := p.(*ast.FuncDecl); ok {
@@ -172,43 +198,58 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 	if outer == nil {
 		return nil, fmt.Errorf("extractFunction: no enclosing function")
 	}
-	// At the moment, we don't extract selections containing return statements,
-	// as they are more complex and need to be adjusted to maintain correctness.
-	// TODO: Support extracting and rewriting code with return statements.
-	var containsReturn bool
-	ast.Inspect(outer, func(n ast.Node) bool {
-		if n == nil {
-			return true
-		}
-		if rng.Start <= n.Pos() && n.End() <= rng.End {
-			if _, ok := n.(*ast.ReturnStmt); ok {
-				containsReturn = true
-				return false
-			}
-		}
-		return n.Pos() <= rng.End
-	})
-	if containsReturn {
-		return nil, fmt.Errorf("extractFunction: selected block contains return")
-	}
+
 	// Find the nodes at the start and end of the selection.
 	var start, end ast.Node
 	ast.Inspect(outer, func(n ast.Node) bool {
 		if n == nil {
 			return true
 		}
-		if n.Pos() == rng.Start && n.End() <= rng.End {
+		// Do not override 'start' with a node that begins at the same location but is
+		// nested further from 'outer'.
+		if start == nil && n.Pos() == rng.Start && n.End() <= rng.End {
 			start = n
 		}
-		if n.End() == rng.End && n.Pos() >= rng.Start {
+		if end == nil && n.End() == rng.End && n.Pos() >= rng.Start {
 			end = n
 		}
 		return n.Pos() <= rng.End
 	})
 	if start == nil || end == nil {
-		return nil, fmt.Errorf("extractFunction: start or end node is empty")
+		return nil, nil
 	}
 
+	// TODO: Support non-nested return statements.
+	// A return statement is non-nested if its parent node is equal to the parent node
+	// of the first node in the selection. These cases must be handled seperately because
+	// non-nested return statements are guaranteed to execute. Our control flow does not
+	// properly consider these situations yet.
+	var retStmts []*ast.ReturnStmt
+	var hasNonNestedReturn bool
+	startParent := findParent(outer, start)
+	ast.Inspect(outer, func(n ast.Node) bool {
+		if n == nil {
+			return true
+		}
+		if n.Pos() < rng.Start || n.End() > rng.End {
+			return n.Pos() <= rng.End
+		}
+		ret, ok := n.(*ast.ReturnStmt)
+		if !ok {
+			return true
+		}
+		if findParent(outer, n) == startParent {
+			hasNonNestedReturn = true
+			return false
+		}
+		retStmts = append(retStmts, ret)
+		return true
+	})
+	if hasNonNestedReturn {
+		return nil, fmt.Errorf("extractFunction: selected bloc kcontains non-nested return")
+	}
+	containsReturnStatement := len(retStmts) > 0
+
 	// Now that we have determined the correct range for the selection block,
 	// we must determine the signature of the extracted function. We will then replace
 	// the block with an assignment statement that calls the extracted function with
@@ -265,122 +306,184 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 		}
 	}
 
-	// Our preference is to replace the selected block with an "x, y, z := fn()" style
-	// assignment statement. We can use this style when none of the variables in the
-	// extracted function's return statement have already be initialized outside of the
-	// selected block. However, for example, if z is already defined elsewhere, we
-	// replace the selected block with:
+	// Find the function literal that encloses the selection. The enclosing function literal
+	// may not be the enclosing function declaration (i.e. 'outer'). For example, in the
+	// following block:
 	//
-	// var x int
-	// var y string
-	// x, y, z = fn()
+	// func main() {
+	//     ast.Inspect(node, func(n ast.Node) bool {
+	//         v := 1 // this line extracted
+	//         return true
+	//     })
+	// }
 	//
-	var initializations string
-	if len(uninitialized) > 0 && len(uninitialized) != len(returns) {
-		var declarations []ast.Stmt
-		for _, obj := range uninitialized {
-			if _, ok := seenUninitialized[obj]; ok {
-				continue
-			}
-			seenUninitialized[obj] = struct{}{}
-			valSpec := &ast.ValueSpec{
-				Names: []*ast.Ident{ast.NewIdent(obj.Name())},
-				Type:  seenVars[obj],
-			}
-			genDecl := &ast.GenDecl{
-				Tok:   token.VAR,
-				Specs: []ast.Spec{valSpec},
-			}
-			declarations = append(declarations, &ast.DeclStmt{Decl: genDecl})
+	// 'outer' is main(). However, the extracted selection most directly belongs to
+	// the anonymous function literal, the second argument of ast.Inspect(). We use the
+	// enclosing function literal to determine the proper return types for return statements
+	// within the selection. We still need the enclosing function declaration because this is
+	// the top-level declaration. We inspect the top-level declaration to look for variables
+	// as well as for code replacement.
+	enclosing := outer.Type
+	for _, p := range path {
+		if p == enclosing {
+			break
 		}
-		var declBuf bytes.Buffer
-		if err := format.Node(&declBuf, fset, declarations); err != nil {
-			return nil, err
-		}
-		indent := calculateIndentation(src, tok, start)
-		// Add proper indentation to each declaration. Also add formatting to
-		// the line following the last initialization to ensure that subsequent
-		// edits begin at the proper location.
-		initializations = strings.ReplaceAll(declBuf.String(), "\n", "\n"+indent) +
-			"\n" + indent
-	}
-
-	name := generateAvailableIdentifier(start.Pos(), file, path, info)
-	var replace ast.Node
-	if len(returns) > 0 {
-		// If none of the variables on the left-hand side of the function call have
-		// been initialized before the selection, we can use := instead of =.
-		assignTok := token.ASSIGN
-		if len(uninitialized) == len(returns) {
-			assignTok = token.DEFINE
-		}
-		callExpr := &ast.CallExpr{
-			Fun:  ast.NewIdent(name),
-			Args: params,
-		}
-		replace = &ast.AssignStmt{
-			Lhs: returns,
-			Tok: assignTok,
-			Rhs: []ast.Expr{callExpr},
-		}
-	} else {
-		replace = &ast.CallExpr{
-			Fun:  ast.NewIdent(name),
-			Args: params,
+		if fl, ok := p.(*ast.FuncLit); ok {
+			enclosing = fl.Type
+			break
 		}
 	}
 
+	// We put the selection in a constructed file. We can then traverse and edit
+	// the extracted selection without modifying the original AST.
 	startOffset := tok.Offset(rng.Start)
 	endOffset := tok.Offset(rng.End)
 	selection := src[startOffset:endOffset]
-	// Put selection in constructed file to parse and produce block statement. We can
-	// then use the block statement to traverse and edit extracted function without
-	// altering the original file.
-	text := "package main\nfunc _() { " + string(selection) + " }"
-	extract, err := parser.ParseFile(fset, "", text, 0)
+	extractedBlock, err := parseBlockStmt(fset, selection)
 	if err != nil {
 		return nil, err
 	}
-	if len(extract.Decls) == 0 {
-		return nil, fmt.Errorf("parsed file does not contain any declarations")
+
+	// We need to account for return statements in the selected block, as they will complicate
+	// the logical flow of the extracted function. See the following example, where ** denotes
+	// the range to be extracted.
+	//
+	// Before:
+	//
+	// func _() int {
+	//     a := 1
+	//     b := 2
+	//     **if a == b {
+	//         return a
+	//     }**
+	//     ...
+	// }
+	//
+	// After:
+	//
+	// func _() int {
+	//     a := 1
+	//     b := 2
+	//     cond0, ret0 := x0(a, b)
+	//     if cond0 {
+	//         return ret0
+	//     }
+	//     ...
+	// }
+	//
+	// func x0(a int, b int) (bool, int) {
+	//     if a == b {
+	//         return true, a
+	//     }
+	//     return false, 0
+	// }
+	//
+	// We handle returns by adding an additional boolean return value to the extracted function.
+	// This bool reports whether the original function would have returned. Because the
+	// extracted selection contains a return statement, we must also add the types in the
+	// return signature of the enclosing function to the return signature of the
+	// extracted function. We then add an extra if statement checking this boolean value
+	// in the original function. If the condition is met, the original function should
+	// return a value, mimicking the functionality of the original return statement(s)
+	// in the selection.
+
+	var retVars []*returnVariable
+	var ifReturn *ast.IfStmt
+	if containsReturnStatement {
+		// The selected block contained return statements, so we have to modify the
+		// signature of the extracted function as described above. Adjust all of
+		// the return statements in the extracted function to reflect this change in
+		// signature.
+		if err := adjustReturnStatements(returnTypes, seenVars, fset, file,
+			pkg, extractedBlock); err != nil {
+			return nil, err
+		}
+		// Collect the additional return values and types needed to accomodate return
+		// statements in the selection. Update the type signature of the extracted
+		// function and construct the if statement that will be inserted in the enclosing
+		// function.
+		retVars, ifReturn, err = generateReturnInfo(
+			enclosing, pkg, path, file, info, fset, rng.Start)
+		if err != nil {
+			return nil, err
+		}
 	}
-	decl, ok := extract.Decls[0].(*ast.FuncDecl)
-	if !ok {
-		return nil, fmt.Errorf("parsed file does not contain expected function declaration")
+
+	// Add a return statement to the end of the new function. This return statement must include
+	// the values for the types of the original extracted function signature and (if a return
+	// statement is present in the selection) enclosing function signature.
+	hasReturnValues := len(returns)+len(retVars) > 0
+	if hasReturnValues {
+		extractedBlock.List = append(extractedBlock.List, &ast.ReturnStmt{
+			Results: append(returns, getZeroVals(retVars)...)})
 	}
-	// Add return statement to the end of the new function.
-	if len(returns) > 0 {
-		decl.Body.List = append(decl.Body.List,
-			&ast.ReturnStmt{Results: returns},
-		)
+
+	// Construct the appropriate call to the extracted function.
+	funName := generateAvailableIdentifier(rng.Start, file, path, info, "fn", 0)
+	// If none of the variables on the left-hand side of the function call have
+	// been initialized before the selection, we can use ':=' instead of '='.
+	sym := token.ASSIGN
+	if len(uninitialized) == len(returns) {
+		sym = token.DEFINE
 	}
-	funcDecl := &ast.FuncDecl{
-		Name: ast.NewIdent(name),
+	extractedFunCall := generateFuncCall(hasReturnValues, params,
+		append(returns, getNames(retVars)...), funName, sym)
+
+	// Build the extracted function.
+	newFunc := &ast.FuncDecl{
+		Name: ast.NewIdent(funName),
 		Type: &ast.FuncType{
 			Params:  &ast.FieldList{List: paramTypes},
-			Results: &ast.FieldList{List: returnTypes},
+			Results: &ast.FieldList{List: append(returnTypes, getDecls(retVars)...)},
 		},
-		Body: decl.Body,
+		Body: extractedBlock,
 	}
 
-	var replaceBuf, newFuncBuf bytes.Buffer
-	if err := format.Node(&replaceBuf, fset, replace); err != nil {
-		return nil, err
-	}
-	if err := format.Node(&newFuncBuf, fset, funcDecl); err != nil {
+	// Create variable declarations for any identifiers that need to be initialized prior to
+	// calling the extracted function.
+	declarations, err := initializeVars(
+		uninitialized, returns, retVars, seenUninitialized, seenVars)
+	if err != nil {
+		return nil, err
+	}
+
+	var declBuf, replaceBuf, newFuncBuf, ifBuf bytes.Buffer
+	if err := format.Node(&declBuf, fset, declarations); err != nil {
+		return nil, err
+	}
+	if err := format.Node(&replaceBuf, fset, extractedFunCall); err != nil {
+		return nil, err
+	}
+	if ifReturn != nil {
+		if err := format.Node(&ifBuf, fset, ifReturn); err != nil {
+			return nil, err
+		}
+	}
+	if err := format.Node(&newFuncBuf, fset, newFunc); err != nil {
 		return nil, err
 	}
 
-	outerStart := tok.Offset(outer.Pos())
-	outerEnd := tok.Offset(outer.End())
 	// We're going to replace the whole enclosing function,
 	// so preserve the text before and after the selected block.
+	outerStart := tok.Offset(outer.Pos())
+	outerEnd := tok.Offset(outer.End())
 	before := src[outerStart:startOffset]
 	after := src[endOffset:outerEnd]
+	newLineIndent := "\n" + calculateIndentation(src, tok, start)
+
 	var fullReplacement strings.Builder
 	fullReplacement.Write(before)
-	fullReplacement.WriteString(initializations) // add any initializations, if needed
-	fullReplacement.Write(replaceBuf.Bytes())    // call the extracted function
+	if declBuf.Len() > 0 { // add any initializations, if needed
+		initializations := strings.ReplaceAll(declBuf.String(), "\n", newLineIndent) +
+			newLineIndent
+		fullReplacement.WriteString(initializations)
+	}
+	fullReplacement.Write(replaceBuf.Bytes()) // call the extracted function
+	if ifBuf.Len() > 0 {                      // add the if statement below the function call, if needed
+		ifstatement := newLineIndent +
+			strings.ReplaceAll(ifBuf.String(), "\n", newLineIndent)
+		fullReplacement.WriteString(ifstatement)
+	}
 	fullReplacement.Write(after)
 	fullReplacement.WriteString("\n\n")       // add newlines after the enclosing function
 	fullReplacement.Write(newFuncBuf.Bytes()) // insert the extracted function
@@ -396,6 +499,57 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 	}, nil
 }
 
+// adjustRangeForWhitespace adjusts the given range to exclude unnecessary leading or
+// trailing whitespace characters from selection. In the following example, each line
+// of the if statement is indented once. There are also two extra spaces after the
+// closing bracket before the line break.
+//
+// \tif (true) {
+// \t    _ = 1
+// \t}  \n
+//
+// By default, a valid range begins at 'if' and ends at the first whitespace character
+// after the '}'. But, users are likely to highlight full lines rather than adjusting
+// their cursors for whitespace. To support this use case, we must manually adjust the
+// ranges to match the correct AST node. In this particular example, we would adjust
+// rng.Start forward by one byte, and rng.End backwards by two bytes.
+func adjustRangeForWhitespace(rng span.Range, tok *token.File, content []byte) span.Range {
+	offset := tok.Offset(rng.Start)
+	for offset < len(content) {
+		if !unicode.IsSpace(rune(content[offset])) {
+			break
+		}
+		// Move forwards one byte to find a non-whitespace character.
+		offset += 1
+	}
+	rng.Start = tok.Pos(offset)
+
+	offset = tok.Offset(rng.End)
+	for offset-1 >= 0 {
+		if !unicode.IsSpace(rune(content[offset-1])) {
+			break
+		}
+		// Move backwards one byte to find a non-whitespace character.
+		offset -= 1
+	}
+	rng.End = tok.Pos(offset)
+	return rng
+}
+
+// findParent finds the parent AST node of the given target node, if the target is a
+// descendant of the starting node.
+func findParent(start ast.Node, target ast.Node) ast.Node {
+	var parent ast.Node
+	analysisinternal.WalkASTWithParent(start, func(n, p ast.Node) bool {
+		if n == target {
+			parent = p
+			return false
+		}
+		return true
+	})
+	return parent
+}
+
 // collectFreeVars maps each identifier in the given range to whether it is "free."
 // Given a range, a variable in that range is defined as "free" if it is declared
 // outside of the range and neither at the file scope nor package scope. These free
@@ -548,58 +702,7 @@ func canExtractFunction(fset *token.FileSet, rng span.Range, src []byte, file *a
 	return ok
 }
 
-// Adjust new function name until no collisons in scope. Possible collisions include
-// other function and variable names.
-func generateAvailableIdentifier(pos token.Pos, file *ast.File, path []ast.Node, info *types.Info) string {
-	scopes := collectScopes(info, path, pos)
-	var idx int
-	name := "x0"
-	for file.Scope.Lookup(name) != nil || !isValidName(name, scopes) {
-		idx++
-		name = fmt.Sprintf("x%d", idx)
-	}
-	return name
-}
-
-// adjustRangeForWhitespace adjusts the given range to exclude unnecessary leading or
-// trailing whitespace characters from selection. In the following example, each line
-// of the if statement is indented once. There are also two extra spaces after the
-// closing bracket before the line break.
-//
-// \tif (true) {
-// \t    _ = 1
-// \t}  \n
-//
-// By default, a valid range begins at 'if' and ends at the first whitespace character
-// after the '}'. But, users are likely to highlight full lines rather than adjusting
-// their cursors for whitespace. To support this use case, we must manually adjust the
-// ranges to match the correct AST node. In this particular example, we would adjust
-// rng.Start forward by one byte, and rng.End backwards by two bytes.
-func adjustRangeForWhitespace(rng span.Range, tok *token.File, content []byte) span.Range {
-	offset := tok.Offset(rng.Start)
-	for offset < len(content) {
-		if !unicode.IsSpace(rune(content[offset])) {
-			break
-		}
-		// Move forwards one byte to find a non-whitespace character.
-		offset += 1
-	}
-	rng.Start = tok.Pos(offset)
-
-	offset = tok.Offset(rng.End)
-	for offset-1 >= 0 {
-		if !unicode.IsSpace(rune(content[offset-1])) {
-			break
-		}
-		// Move backwards one byte to find a non-whitespace character.
-		offset -= 1
-	}
-	rng.End = tok.Pos(offset)
-	return rng
-}
-
-// objUsed checks if the object is used after the selection but within
-// the scope of the enclosing function.
+// objUsed checks if the object is used between the given positions.
 func objUsed(obj types.Object, info *types.Info, endSel token.Pos, endScope token.Pos) bool {
 	for id, ob := range info.Uses {
 		if obj == ob && endSel < id.Pos() && id.End() <= endScope {
@@ -608,3 +711,208 @@ func objUsed(obj types.Object, info *types.Info, endSel token.Pos, endScope toke
 	}
 	return false
 }
+
+// parseExtraction generates an AST file from the given text. We then return the portion of the
+// file that represents the text.
+func parseBlockStmt(fset *token.FileSet, src []byte) (*ast.BlockStmt, error) {
+	text := "package main\nfunc _() { " + string(src) + " }"
+	extract, err := parser.ParseFile(fset, "", text, 0)
+	if err != nil {
+		return nil, err
+	}
+	if len(extract.Decls) == 0 {
+		return nil, fmt.Errorf("parsed file does not contain any declarations")
+	}
+	decl, ok := extract.Decls[0].(*ast.FuncDecl)
+	if !ok {
+		return nil, fmt.Errorf("parsed file does not contain expected function declaration")
+	}
+	if decl.Body == nil {
+		return nil, fmt.Errorf("extracted function has no body")
+	}
+	return decl.Body, nil
+}
+
+// generateReturnInfo generates the information we need to adjust the return statements and
+// signature of the extracted function. We prepare names, signatures, and "zero values" that
+// represent the new variables. We also use this information to construct the if statement that
+// is inserted below the call to the extracted function.
+func generateReturnInfo(enclosing *ast.FuncType, pkg *types.Package, path []ast.Node, file *ast.File, info *types.Info, fset *token.FileSet, pos token.Pos) ([]*returnVariable, *ast.IfStmt, error) {
+	// Generate information for the added bool value.
+	cond := &ast.Ident{Name: generateAvailableIdentifier(pos, file, path, info, "cond", 0)}
+	retVars := []*returnVariable{
+		{
+			name:    cond,
+			decl:    &ast.Field{Type: ast.NewIdent("bool")},
+			zeroVal: ast.NewIdent("false"),
+		},
+	}
+	// Generate information for the values in the return signature of the enclosing function.
+	if enclosing.Results != nil {
+		for i, field := range enclosing.Results.List {
+			typ := info.TypeOf(field.Type)
+			if typ == nil {
+				return nil, nil, fmt.Errorf(
+					"failed type conversion, AST expression: %T", field.Type)
+			}
+			expr := analysisinternal.TypeExpr(fset, file, pkg, typ)
+			if expr == nil {
+				return nil, nil, fmt.Errorf("nil AST expression")
+			}
+			retVars = append(retVars, &returnVariable{
+				name: ast.NewIdent(generateAvailableIdentifier(pos, file,
+					path, info, "ret", i)),
+				decl: &ast.Field{Type: expr},
+				zeroVal: analysisinternal.ZeroValue(
+					fset, file, pkg, typ),
+			})
+		}
+	}
+	// Create the return statement for the enclosing function. We must exclude the variable
+	// for the condition of the if statement (cond) from the return statement.
+	ifReturn := &ast.IfStmt{
+		Cond: cond,
+		Body: &ast.BlockStmt{
+			List: []ast.Stmt{&ast.ReturnStmt{Results: getNames(retVars)[1:]}},
+		},
+	}
+	return retVars, ifReturn, nil
+}
+
+// adjustReturnStatements adds "zero values" of the given types to each return statement
+// in the given AST node.
+func adjustReturnStatements(returnTypes []*ast.Field, seenVars map[types.Object]ast.Expr, fset *token.FileSet, file *ast.File, pkg *types.Package, extractedBlock *ast.BlockStmt) error {
+	var zeroVals []ast.Expr
+	// Create "zero values" for each type.
+	for _, returnType := range returnTypes {
+		var val ast.Expr
+		for obj, typ := range seenVars {
+			if typ != returnType.Type {
+				continue
+			}
+			val = analysisinternal.ZeroValue(fset, file, pkg, obj.Type())
+			break
+		}
+		if val == nil {
+			return fmt.Errorf(
+				"could not find matching AST expression for %T", returnType.Type)
+		}
+		zeroVals = append(zeroVals, val)
+	}
+	// Add "zero values" to each return statement.
+	// The bool reports whether the enclosing function should return after calling the
+	// extracted function. We set the bool to 'true' because, if these return statements
+	// execute, the extracted function terminates early, and the enclosing function must
+	// return as well.
+	zeroVals = append(zeroVals, ast.NewIdent("true"))
+	ast.Inspect(extractedBlock, func(n ast.Node) bool {
+		if n == nil {
+			return true
+		}
+		if n, ok := n.(*ast.ReturnStmt); ok {
+			n.Results = append(zeroVals, n.Results...)
+			return true
+		}
+		return true
+	})
+	return nil
+}
+
+// generateFuncCall constructs a call expression for the extracted function, described by the
+// given parameters and return variables.
+func generateFuncCall(hasReturnVals bool, params, returns []ast.Expr, name string, token token.Token) ast.Node {
+	var replace ast.Node
+	if hasReturnVals {
+		callExpr := &ast.CallExpr{
+			Fun:  ast.NewIdent(name),
+			Args: params,
+		}
+		replace = &ast.AssignStmt{
+			Lhs: returns,
+			Tok: token,
+			Rhs: []ast.Expr{callExpr},
+		}
+	} else {
+		replace = &ast.CallExpr{
+			Fun:  ast.NewIdent(name),
+			Args: params,
+		}
+	}
+	return replace
+}
+
+// initializeVars creates variable declarations, if needed.
+// Our preference is to replace the selected block with an "x, y, z := fn()" style
+// assignment statement. We can use this style when none of the variables in the
+// extracted function's return statement have already be initialized outside of the
+// selected block. However, for example, if z is already defined elsewhere, we
+// replace the selected block with:
+//
+// var x int
+// var y string
+// x, y, z = fn()
+func initializeVars(uninitialized []types.Object, returns []ast.Expr, retVars []*returnVariable, seenUninitialized map[types.Object]struct{}, seenVars map[types.Object]ast.Expr) ([]ast.Stmt, error) {
+	var declarations []ast.Stmt
+	// We do not manually initialize variables if every return value is unitialized.
+	// We can use := to initialize the variables in this situation.
+	if len(uninitialized) == len(returns) {
+		return declarations, nil
+	}
+	for _, obj := range uninitialized {
+		if _, ok := seenUninitialized[obj]; ok {
+			continue
+		}
+		seenUninitialized[obj] = struct{}{}
+		valSpec := &ast.ValueSpec{
+			Names: []*ast.Ident{ast.NewIdent(obj.Name())},
+			Type:  seenVars[obj],
+		}
+		genDecl := &ast.GenDecl{
+			Tok:   token.VAR,
+			Specs: []ast.Spec{valSpec},
+		}
+		declarations = append(declarations, &ast.DeclStmt{Decl: genDecl})
+	}
+	// Each variable added from a return statement in the selection
+	// must be initialized.
+	for i, retVar := range retVars {
+		n := retVar.name.(*ast.Ident)
+		valSpec := &ast.ValueSpec{
+			Names: []*ast.Ident{n},
+			Type:  retVars[i].decl.Type,
+		}
+		genDecl := &ast.GenDecl{
+			Tok:   token.VAR,
+			Specs: []ast.Spec{valSpec},
+		}
+		declarations = append(declarations, &ast.DeclStmt{Decl: genDecl})
+	}
+	return declarations, nil
+}
+
+// getNames returns the names from the given list of returnVariable.
+func getNames(retVars []*returnVariable) []ast.Expr {
+	var names []ast.Expr
+	for _, retVar := range retVars {
+		names = append(names, retVar.name)
+	}
+	return names
+}
+
+// getZeroVals returns the "zero values" from the given list of returnVariable.
+func getZeroVals(retVars []*returnVariable) []ast.Expr {
+	var zvs []ast.Expr
+	for _, retVar := range retVars {
+		zvs = append(zvs, retVar.zeroVal)
+	}
+	return zvs
+}
+
+// getDecls returns the declarations from the given list of returnVariable.
+func getDecls(retVars []*returnVariable) []*ast.Field {
+	var decls []*ast.Field
+	for _, retVar := range retVars {
+		decls = append(decls, retVar.decl)
+	}
+	return decls
+}

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_args_returns.go.golden

@@ -3,13 +3,13 @@ package extract
 
 func _() {
 	a := 1
-	a = x0(a) //@mark(e1, "2")
+	a = fn0(a) //@mark(e1, "2")
 	//@extractfunc(s1, e1)
 	b := a * 2
 	var _ = 3 + 4
 }
 
-func x0(a int) int {
+func fn0(a int) int {
 	a = 5
 	a = a + 2
 	return a

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_basic.go

@@ -0,0 +1,10 @@
+package extract
+
+func _() bool {
+	x := 1
+	if x == 0 { //@mark(s0, "if")
+		return true
+	} //@mark(e0, "}")
+	return false
+	//@extractfunc(s0, e0)
+}

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_basic.go.golden

@@ -0,0 +1,20 @@
+-- functionextraction_extract_return_basic_5_2 --
+package extract
+
+func _() bool {
+	x := 1
+	cond0, ret0 := x0(x)
+	if cond0 {
+		return ret0
+	} //@mark(e0, "}")
+	return false
+	//@extractfunc(s0, e0)
+}
+
+func x0(x int) (bool, bool) {
+	if x == 0 {
+		return true, true
+	}
+	return false, false
+}
+

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_complex.go

@@ -0,0 +1,17 @@
+package extract
+
+import "fmt"
+
+func _() (int, string, error) {
+	x := 1
+	y := "hello"
+	z := "bye" //@mark(s0, "z")
+	if y == z {
+		return x, y, fmt.Errorf("same")
+	} else {
+		z = "hi"
+		return x, z, nil
+	} //@mark(e0, "}")
+	return x, z, nil
+	//@extractfunc(s0, e0)
+}

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_complex.go.golden

@@ -0,0 +1,27 @@
+-- functionextraction_extract_return_complex_8_2 --
+package extract
+
+import "fmt"
+
+func _() (int, string, error) {
+	x := 1
+	y := "hello"
+	z, cond0, ret0, ret1, ret2 := x0(y, x)
+	if cond0 {
+		return ret0, ret1, ret2
+	} //@mark(e0, "}")
+	return x, z, nil
+	//@extractfunc(s0, e0)
+}
+
+func x0(y string, x int) (string, bool, int, string, error) {
+	z := "bye"
+	if y == z {
+		return "", true, x, y, fmt.Errorf("same")
+	} else {
+		z = "hi"
+		return "", true, x, z, nil
+	}
+	return z, false, 0, "", nil
+}
+

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_func_lit.go

@@ -0,0 +1,13 @@
+package extract
+
+import "go/ast"
+
+func _() {
+	ast.Inspect(ast.NewIdent("a"), func(n ast.Node) bool {
+		if n == nil { //@mark(s0, "if")
+			return true
+		} //@mark(e0, "}")
+		return false
+	})
+	//@extractfunc(s0, e0)
+}

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_func_lit.go.golden

@@ -0,0 +1,23 @@
+-- functionextraction_extract_return_func_lit_7_3 --
+package extract
+
+import "go/ast"
+
+func _() {
+	ast.Inspect(ast.NewIdent("a"), func(n ast.Node) bool {
+		cond0, ret0 := x0(n)
+		if cond0 {
+			return ret0
+		} //@mark(e0, "}")
+		return false
+	})
+	//@extractfunc(s0, e0)
+}
+
+func x0(n ast.Node) (bool, bool) {
+	if n == nil {
+		return true, true
+	}
+	return false, false
+}
+

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_init.go

@@ -0,0 +1,12 @@
+package extract
+
+func _() string {
+	x := 1
+	if x == 0 { //@mark(s0, "if")
+		x = 3
+		return "a"
+	} //@mark(e0, "}")
+	x = 2
+	return "b"
+	//@extractfunc(s0, e0)
+}

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_return_init.go.golden

@@ -0,0 +1,24 @@
+-- functionextraction_extract_return_init_5_2 --
+package extract
+
+func _() string {
+	x := 1
+	var cond0 bool
+	var ret0 string
+	x, cond0, ret0 = fn0(x)
+	if cond0 {
+		return ret0
+	} //@mark(e0, "}")
+	x = 2
+	return "b"
+	//@extractfunc(s0, e0)
+}
+
+func fn0(x int) (int, bool, string) {
+	if x == 0 {
+		x = 3
+		return 0, true, "a"
+	}
+	return x, false, ""
+}
+

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_scope.go

@@ -1,10 +1,10 @@
 package extract
 
 func _() {
-	x0 := 1
-	a := x0 //@extractfunc("a", "x0")
+	fn0 := 1
+	a := fn0 //@extractfunc("a", "fn0")
 }
 
-func x1() int {
+func fn1() int {
 	return 1
 }

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_scope.go.golden

@@ -2,15 +2,15 @@
 package extract
 
 func _() {
-	x0 := 1
-	x2(x0) //@extractfunc("a", "x0")
+	fn0 := 1
+	fn2(fn0) //@extractfunc("a", "fn0")
 }
 
-func x2(x0 int) {
-	a := x0
+func fn2(fn0 int) {
+	a := fn0
 }
 
-func x1() int {
+func fn1() int {
 	return 1
 }
 

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_smart_initialization.go.golden

@@ -4,12 +4,12 @@ package extract
 func _() {
 	var a []int
 	var b int
-	a, b = x0(a)           //@mark(e4, "4")
+	a, b = fn0(a)           //@mark(e4, "4")
 	//@extractfunc(s4, e4)
 	a = append(a, b)
 }
 
-func x0(a []int) ([]int, int) {
+func fn0(a []int) ([]int, int) {
 	a = append(a, 2)
 	b := 4
 	return a, b

internal/lsp/testdata/lsp/primarymod/extract/extract_function/extract_smart_return.go.golden

@@ -4,12 +4,12 @@ package extract
 func _() {
 	var b []int
 	var a int
-	b = x0(a, b) //@mark(e2, ")")
+	b = fn0(a, b) //@mark(e2, ")")
 	b[0] = 1
 	//@extractfunc(s2, e2)
 }
 
-func x0(a int, b []int) []int {
+func fn0(a int, b []int) []int {
 	a = 2
 	b = []int{}
 	b = append(b, a)