blob: 0d6880014bb1e2c3d707d23148aef1611eccaafd [file] [log] [blame]
// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package source
import (
"bytes"
"fmt"
"go/ast"
"go/format"
"go/parser"
"go/token"
"go/types"
"strings"
"unicode"
"golang.org/x/tools/go/analysis"
"golang.org/x/tools/go/ast/astutil"
"golang.org/x/tools/internal/analysisinternal"
"golang.org/x/tools/internal/span"
)
func extractVariable(fset *token.FileSet, rng span.Range, src []byte, file *ast.File, _ *types.Package, info *types.Info) (*analysis.SuggestedFix, error) {
expr, path, ok, err := CanExtractVariable(rng, file)
if !ok {
return nil, fmt.Errorf("extractVariable: cannot extract %s: %v", fset.Position(rng.Start), err)
}
// Create new AST node for extracted code.
var lhsNames []string
switch expr := expr.(type) {
// TODO: stricter rules for selectorExpr.
case *ast.BasicLit, *ast.CompositeLit, *ast.IndexExpr, *ast.SliceExpr,
*ast.UnaryExpr, *ast.BinaryExpr, *ast.SelectorExpr:
lhsName, _ := generateAvailableIdentifier(expr.Pos(), file, path, info, "x", 0)
lhsNames = append(lhsNames, lhsName)
case *ast.CallExpr:
tup, ok := info.TypeOf(expr).(*types.Tuple)
if !ok {
// If the call expression only has one return value, we can treat it the
// same as our standard extract variable case.
lhsName, _ := generateAvailableIdentifier(expr.Pos(), file, path, info, "x", 0)
lhsNames = append(lhsNames, lhsName)
break
}
idx := 0
for i := 0; i < tup.Len(); i++ {
// Generate a unique variable for each return value.
var lhsName string
lhsName, idx = generateAvailableIdentifier(expr.Pos(), file, path, info, "x", idx)
lhsNames = append(lhsNames, lhsName)
}
default:
return nil, fmt.Errorf("cannot extract %T", expr)
}
insertBeforeStmt := analysisinternal.StmtToInsertVarBefore(path)
if insertBeforeStmt == nil {
return nil, fmt.Errorf("cannot find location to insert extraction")
}
tok := fset.File(expr.Pos())
if tok == nil {
return nil, fmt.Errorf("no file for pos %v", fset.Position(file.Pos()))
}
indent, err := calculateIndentation(src, tok, insertBeforeStmt)
if err != nil {
return nil, err
}
newLineIndent := "\n" + indent
lhs := strings.Join(lhsNames, ", ")
assignStmt := &ast.AssignStmt{
Lhs: []ast.Expr{ast.NewIdent(lhs)},
Tok: token.DEFINE,
Rhs: []ast.Expr{expr},
}
var buf bytes.Buffer
if err := format.Node(&buf, fset, assignStmt); err != nil {
return nil, err
}
assignment := strings.ReplaceAll(buf.String(), "\n", newLineIndent) + newLineIndent
return &analysis.SuggestedFix{
TextEdits: []analysis.TextEdit{
{
Pos: insertBeforeStmt.Pos(),
End: insertBeforeStmt.Pos(),
NewText: []byte(assignment),
},
{
Pos: rng.Start,
End: rng.End,
NewText: []byte(lhs),
},
},
}, nil
}
// CanExtractVariable reports whether the code in the given range can be
// extracted to a variable.
func CanExtractVariable(rng span.Range, file *ast.File) (ast.Expr, []ast.Node, bool, error) {
if rng.Start == rng.End {
return nil, nil, false, fmt.Errorf("start and end are equal")
}
path, _ := astutil.PathEnclosingInterval(file, rng.Start, rng.End)
if len(path) == 0 {
return nil, nil, false, fmt.Errorf("no path enclosing interval")
}
for _, n := range path {
if _, ok := n.(*ast.ImportSpec); ok {
return nil, nil, false, fmt.Errorf("cannot extract variable in an import block")
}
}
node := path[0]
if rng.Start != node.Pos() || rng.End != node.End() {
return nil, nil, false, fmt.Errorf("range does not map to an AST node")
}
expr, ok := node.(ast.Expr)
if !ok {
return nil, nil, false, fmt.Errorf("node is not an expression")
}
switch expr.(type) {
case *ast.BasicLit, *ast.CompositeLit, *ast.IndexExpr, *ast.CallExpr,
*ast.SliceExpr, *ast.UnaryExpr, *ast.BinaryExpr, *ast.SelectorExpr:
return expr, path, true, nil
}
return nil, nil, false, fmt.Errorf("cannot extract an %T to a variable", expr)
}
// Calculate indentation for insertion.
// When inserting lines of code, we must ensure that the lines have consistent
// formatting (i.e. the proper indentation). To do so, we observe the indentation on the
// line of code on which the insertion occurs.
func calculateIndentation(content []byte, tok *token.File, insertBeforeStmt ast.Node) (string, error) {
line := tok.Line(insertBeforeStmt.Pos())
lineOffset, err := Offset(tok, tok.LineStart(line))
if err != nil {
return "", err
}
stmtOffset, err := Offset(tok, insertBeforeStmt.Pos())
if err != nil {
return "", err
}
return string(content[lineOffset:stmtOffset]), nil
}
// generateAvailableIdentifier adjusts the new function name until there are no collisions in scope.
// Possible collisions include other function and variable names. Returns the next index to check for prefix.
func generateAvailableIdentifier(pos token.Pos, file *ast.File, path []ast.Node, info *types.Info, prefix string, idx int) (string, int) {
scopes := CollectScopes(info, path, pos)
return generateIdentifier(idx, prefix, func(name string) bool {
return file.Scope.Lookup(name) != nil || !isValidName(name, scopes)
})
}
func generateIdentifier(idx int, prefix string, hasCollision func(string) bool) (string, int) {
name := prefix
if idx != 0 {
name += fmt.Sprintf("%d", idx)
}
for hasCollision(name) {
idx++
name = fmt.Sprintf("%v%d", prefix, idx)
}
return name, idx + 1
}
// isValidName checks for variable collision in scope.
func isValidName(name string, scopes []*types.Scope) bool {
for _, scope := range scopes {
if scope == nil {
continue
}
if scope.Lookup(name) != nil {
return false
}
}
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
}
// extractMethod refactors the selected block of code into a new method.
func extractMethod(fset *token.FileSet, rng span.Range, src []byte, file *ast.File, pkg *types.Package, info *types.Info) (*analysis.SuggestedFix, error) {
return extractFunctionMethod(fset, rng, src, file, pkg, info, true)
}
// extractFunction refactors the selected block of code into a new function.
func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.File, pkg *types.Package, info *types.Info) (*analysis.SuggestedFix, error) {
return extractFunctionMethod(fset, rng, src, file, pkg, info, false)
}
// extractFunctionMethod refactors the selected block of code into a new function/method.
// 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
// and leading whitespace characters to ensure the range is precisely bounded
// by AST nodes. Next, we determine the variables that will be the parameters
// and return values of the extracted function/method. Lastly, we construct the call
// of the function/method and insert this call as well as the extracted function/method into
// their proper locations.
func extractFunctionMethod(fset *token.FileSet, rng span.Range, src []byte, file *ast.File, pkg *types.Package, info *types.Info, isMethod bool) (*analysis.SuggestedFix, error) {
errorPrefix := "extractFunction"
if isMethod {
errorPrefix = "extractMethod"
}
p, ok, methodOk, err := CanExtractFunction(fset, rng, src, file)
if (!ok && !isMethod) || (!methodOk && isMethod) {
return nil, fmt.Errorf("%s: cannot extract %s: %v", errorPrefix,
fset.Position(rng.Start), err)
}
tok, path, rng, outer, start := p.tok, p.path, p.rng, p.outer, p.start
fileScope := info.Scopes[file]
if fileScope == nil {
return nil, fmt.Errorf("%s: file scope is empty", errorPrefix)
}
pkgScope := fileScope.Parent()
if pkgScope == nil {
return nil, fmt.Errorf("%s: package scope is empty", errorPrefix)
}
// 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 separately because
// non-nested return statements are guaranteed to execute.
var retStmts []*ast.ReturnStmt
var hasNonNestedReturn bool
startParent := findParent(outer, start)
ast.Inspect(outer, func(n ast.Node) bool {
if n == nil {
return false
}
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
}
retStmts = append(retStmts, ret)
return false
})
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
// the appropriate parameters and return values.
variables, err := collectFreeVars(info, file, fileScope, pkgScope, rng, path[0])
if err != nil {
return nil, err
}
var (
receiverUsed bool
receiver *ast.Field
receiverName string
receiverObj types.Object
)
if isMethod {
if outer == nil || outer.Recv == nil || len(outer.Recv.List) == 0 {
return nil, fmt.Errorf("%s: cannot extract need method receiver", errorPrefix)
}
receiver = outer.Recv.List[0]
if len(receiver.Names) == 0 || receiver.Names[0] == nil {
return nil, fmt.Errorf("%s: cannot extract need method receiver name", errorPrefix)
}
recvName := receiver.Names[0]
receiverName = recvName.Name
receiverObj = info.ObjectOf(recvName)
}
var (
params, returns []ast.Expr // used when calling the extracted function
paramTypes, returnTypes []*ast.Field // used in the signature of the extracted function
uninitialized []types.Object // vars we will need to initialize before the call
)
// Avoid duplicates while traversing vars and uninitialized.
seenVars := make(map[types.Object]ast.Expr)
seenUninitialized := make(map[types.Object]struct{})
// Some variables on the left-hand side of our assignment statement may be free. If our
// selection begins in the same scope in which the free variable is defined, we can
// redefine it in our assignment statement. See the following example, where 'b' and
// 'err' (both free variables) can be redefined in the second funcCall() while maintaining
// correctness.
//
//
// Not Redefined:
//
// a, err := funcCall()
// var b int
// b, err = funcCall()
//
// Redefined:
//
// a, err := funcCall()
// b, err := funcCall()
//
// We track the number of free variables that can be redefined to maintain our preference
// of using "x, y, z := fn()" style assignment statements.
var canRedefineCount int
// Each identifier in the selected block must become (1) a parameter to the
// extracted function, (2) a return value of the extracted function, or (3) a local
// variable in the extracted function. Determine the outcome(s) for each variable
// based on whether it is free, altered within the selected block, and used outside
// of the selected block.
for _, v := range variables {
if _, ok := seenVars[v.obj]; ok {
continue
}
if v.obj.Name() == "_" {
// The blank identifier is always a local variable
continue
}
typ := analysisinternal.TypeExpr(fset, file, pkg, v.obj.Type())
if typ == nil {
return nil, fmt.Errorf("nil AST expression for type: %v", v.obj.Name())
}
seenVars[v.obj] = typ
identifier := ast.NewIdent(v.obj.Name())
// An identifier must meet three conditions to become a return value of the
// extracted function. (1) its value must be defined or reassigned within
// the selection (isAssigned), (2) it must be used at least once after the
// selection (isUsed), and (3) its first use after the selection
// cannot be its own reassignment or redefinition (objOverriden).
if v.obj.Parent() == nil {
return nil, fmt.Errorf("parent nil")
}
isUsed, firstUseAfter := objUsed(info, span.NewRange(fset, rng.End, v.obj.Parent().End()), v.obj)
if v.assigned && isUsed && !varOverridden(info, firstUseAfter, v.obj, v.free, outer) {
returnTypes = append(returnTypes, &ast.Field{Type: typ})
returns = append(returns, identifier)
if !v.free {
uninitialized = append(uninitialized, v.obj)
} else if v.obj.Parent().Pos() == startParent.Pos() {
canRedefineCount++
}
}
// An identifier must meet two conditions to become a parameter of the
// extracted function. (1) it must be free (isFree), and (2) its first
// use within the selection cannot be its own definition (isDefined).
if v.free && !v.defined {
// Skip the selector for a method.
if isMethod && v.obj == receiverObj {
receiverUsed = true
continue
}
params = append(params, identifier)
paramTypes = append(paramTypes, &ast.Field{
Names: []*ast.Ident{identifier},
Type: typ,
})
}
}
// 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:
//
// func main() {
// ast.Inspect(node, func(n ast.Node) bool {
// v := 1 // this line extracted
// return true
// })
// }
//
// '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
}
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, err := Offset(tok, rng.Start)
if err != nil {
return nil, err
}
endOffset, err := Offset(tok, rng.End)
if err != nil {
return nil, err
}
selection := src[startOffset:endOffset]
extractedBlock, err := parseBlockStmt(fset, selection)
if err != nil {
return nil, err
}
// 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.
//
// If there is a return that is guaranteed to execute (hasNonNestedReturns=true), then
// we don't need to include this additional condition check and can simply return.
//
// Before:
//
// func _() int {
// a := 1
// b := 2
// **if a == b {
// return a
// }
// return b**
// }
//
// After:
//
// func _() int {
// a := 1
// b := 2
// return x0(a, b)
// }
//
// func x0(a int, b int) int {
// if a == b {
// return a
// }
// return b
// }
var retVars []*returnVariable
var ifReturn *ast.IfStmt
if containsReturnStatement {
if !hasNonNestedReturn {
// 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 accommodate 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, hasNonNestedReturn)
if err != nil {
return nil, err
}
}
// 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.
// This only needs to be done if the selections does not have a non-nested return, otherwise
// it already terminates with a return statement.
hasReturnValues := len(returns)+len(retVars) > 0
if hasReturnValues && !hasNonNestedReturn {
extractedBlock.List = append(extractedBlock.List, &ast.ReturnStmt{
Results: append(returns, getZeroVals(retVars)...),
})
}
// Construct the appropriate call to the extracted function.
// We must meet two conditions to use ":=" instead of '='. (1) there must be at least
// one variable on the lhs that is uninitialized (non-free) prior to the assignment.
// (2) all of the initialized (free) variables on the lhs must be able to be redefined.
sym := token.ASSIGN
canDefineCount := len(uninitialized) + canRedefineCount
canDefine := len(uninitialized)+len(retVars) > 0 && canDefineCount == len(returns)
if canDefine {
sym = token.DEFINE
}
var name, funName string
if isMethod {
name = "newMethod"
// TODO(suzmue): generate a name that does not conflict for "newMethod".
funName = name
} else {
name = "newFunction"
funName, _ = generateAvailableIdentifier(rng.Start, file, path, info, name, 0)
}
extractedFunCall := generateFuncCall(hasNonNestedReturn, hasReturnValues, params,
append(returns, getNames(retVars)...), funName, sym, receiverName)
// Build the extracted function.
newFunc := &ast.FuncDecl{
Name: ast.NewIdent(funName),
Type: &ast.FuncType{
Params: &ast.FieldList{List: paramTypes},
Results: &ast.FieldList{List: append(returnTypes, getDecls(retVars)...)},
},
Body: extractedBlock,
}
if isMethod {
var names []*ast.Ident
if receiverUsed {
names = append(names, ast.NewIdent(receiverName))
}
newFunc.Recv = &ast.FieldList{
List: []*ast.Field{{
Names: names,
Type: receiver.Type,
}},
}
}
// Create variable declarations for any identifiers that need to be initialized prior to
// calling the extracted function. We do not manually initialize variables if every return
// value is uninitialized. We can use := to initialize the variables in this situation.
var declarations []ast.Stmt
if canDefineCount != len(returns) {
declarations = initializeVars(uninitialized, retVars, seenUninitialized, seenVars)
}
var declBuf, replaceBuf, newFuncBuf, ifBuf, commentBuf 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
}
// Find all the comments within the range and print them to be put somewhere.
// TODO(suzmue): print these in the extracted function at the correct place.
for _, cg := range file.Comments {
if cg.Pos().IsValid() && cg.Pos() < rng.End && cg.Pos() >= rng.Start {
for _, c := range cg.List {
fmt.Fprintln(&commentBuf, c.Text)
}
}
}
// We're going to replace the whole enclosing function,
// so preserve the text before and after the selected block.
outerStart, err := Offset(tok, outer.Pos())
if err != nil {
return nil, err
}
outerEnd, err := Offset(tok, outer.End())
if err != nil {
return nil, err
}
before := src[outerStart:startOffset]
after := src[endOffset:outerEnd]
indent, err := calculateIndentation(src, tok, start)
if err != nil {
return nil, err
}
newLineIndent := "\n" + indent
var fullReplacement strings.Builder
fullReplacement.Write(before)
if commentBuf.Len() > 0 {
comments := strings.ReplaceAll(commentBuf.String(), "\n", newLineIndent)
fullReplacement.WriteString(comments)
}
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
return &analysis.SuggestedFix{
TextEdits: []analysis.TextEdit{{
Pos: outer.Pos(),
End: outer.End(),
NewText: []byte(fullReplacement.String()),
}},
}, 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, error) {
offset, err := Offset(tok, rng.Start)
if err != nil {
return span.Range{}, err
}
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)
// Move backwards to find a non-whitespace character.
offset, err = Offset(tok, rng.End)
if err != nil {
return span.Range{}, err
}
for o := offset - 1; 0 <= o && o < len(content); o-- {
if !unicode.IsSpace(rune(content[o])) {
break
}
offset = o
}
rng.End = tok.Pos(offset)
return rng, nil
}
// 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
}
// variable describes the status of a variable within a selection.
type variable struct {
obj types.Object
// free reports whether the variable is a free variable, meaning it should
// be a parameter to the extracted function.
free bool
// assigned reports whether the variable is assigned to in the selection.
assigned bool
// defined reports whether the variable is defined in the selection.
defined bool
}
// 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
// variables will be used as arguments in the extracted function. It also returns a
// list of identifiers that may need to be returned by the extracted function.
// Some of the code in this function has been adapted from tools/cmd/guru/freevars.go.
func collectFreeVars(info *types.Info, file *ast.File, fileScope, pkgScope *types.Scope, rng span.Range, node ast.Node) ([]*variable, error) {
// id returns non-nil if n denotes an object that is referenced by the span
// and defined either within the span or in the lexical environment. The bool
// return value acts as an indicator for where it was defined.
id := func(n *ast.Ident) (types.Object, bool) {
obj := info.Uses[n]
if obj == nil {
return info.Defs[n], false
}
if obj.Name() == "_" {
return nil, false // exclude objects denoting '_'
}
if _, ok := obj.(*types.PkgName); ok {
return nil, false // imported package
}
if !(file.Pos() <= obj.Pos() && obj.Pos() <= file.End()) {
return nil, false // not defined in this file
}
scope := obj.Parent()
if scope == nil {
return nil, false // e.g. interface method, struct field
}
if scope == fileScope || scope == pkgScope {
return nil, false // defined at file or package scope
}
if rng.Start <= obj.Pos() && obj.Pos() <= rng.End {
return obj, false // defined within selection => not free
}
return obj, true
}
// sel returns non-nil if n denotes a selection o.x.y that is referenced by the
// span and defined either within the span or in the lexical environment. The bool
// return value acts as an indicator for where it was defined.
var sel func(n *ast.SelectorExpr) (types.Object, bool)
sel = func(n *ast.SelectorExpr) (types.Object, bool) {
switch x := astutil.Unparen(n.X).(type) {
case *ast.SelectorExpr:
return sel(x)
case *ast.Ident:
return id(x)
}
return nil, false
}
seen := make(map[types.Object]*variable)
firstUseIn := make(map[types.Object]token.Pos)
var vars []types.Object
ast.Inspect(node, func(n ast.Node) bool {
if n == nil {
return false
}
if rng.Start <= n.Pos() && n.End() <= rng.End {
var obj types.Object
var isFree, prune bool
switch n := n.(type) {
case *ast.Ident:
obj, isFree = id(n)
case *ast.SelectorExpr:
obj, isFree = sel(n)
prune = true
}
if obj != nil {
seen[obj] = &variable{
obj: obj,
free: isFree,
}
vars = append(vars, obj)
// Find the first time that the object is used in the selection.
first, ok := firstUseIn[obj]
if !ok || n.Pos() < first {
firstUseIn[obj] = n.Pos()
}
if prune {
return false
}
}
}
return n.Pos() <= rng.End
})
// Find identifiers that are initialized or whose values are altered at some
// point in the selected block. For example, in a selected block from lines 2-4,
// variables x, y, and z are included in assigned. However, in a selected block
// from lines 3-4, only variables y and z are included in assigned.
//
// 1: var a int
// 2: var x int
// 3: y := 3
// 4: z := x + a
//
ast.Inspect(node, func(n ast.Node) bool {
if n == nil {
return false
}
if n.Pos() < rng.Start || n.End() > rng.End {
return n.Pos() <= rng.End
}
switch n := n.(type) {
case *ast.AssignStmt:
for _, assignment := range n.Lhs {
lhs, ok := assignment.(*ast.Ident)
if !ok {
continue
}
obj, _ := id(lhs)
if obj == nil {
continue
}
if _, ok := seen[obj]; !ok {
continue
}
seen[obj].assigned = true
if n.Tok != token.DEFINE {
continue
}
// Find identifiers that are defined prior to being used
// elsewhere in the selection.
// TODO: Include identifiers that are assigned prior to being
// used elsewhere in the selection. Then, change the assignment
// to a definition in the extracted function.
if firstUseIn[obj] != lhs.Pos() {
continue
}
// Ensure that the object is not used in its own re-definition.
// For example:
// var f float64
// f, e := math.Frexp(f)
for _, expr := range n.Rhs {
if referencesObj(info, expr, obj) {
continue
}
if _, ok := seen[obj]; !ok {
continue
}
seen[obj].defined = true
break
}
}
return false
case *ast.DeclStmt:
gen, ok := n.Decl.(*ast.GenDecl)
if !ok {
return false
}
for _, spec := range gen.Specs {
vSpecs, ok := spec.(*ast.ValueSpec)
if !ok {
continue
}
for _, vSpec := range vSpecs.Names {
obj, _ := id(vSpec)
if obj == nil {
continue
}
if _, ok := seen[obj]; !ok {
continue
}
seen[obj].assigned = true
}
}
return false
case *ast.IncDecStmt:
if ident, ok := n.X.(*ast.Ident); !ok {
return false
} else if obj, _ := id(ident); obj == nil {
return false
} else {
if _, ok := seen[obj]; !ok {
return false
}
seen[obj].assigned = true
}
}
return true
})
var variables []*variable
for _, obj := range vars {
v, ok := seen[obj]
if !ok {
return nil, fmt.Errorf("no seen types.Object for %v", obj)
}
variables = append(variables, v)
}
return variables, nil
}
// referencesObj checks whether the given object appears in the given expression.
func referencesObj(info *types.Info, expr ast.Expr, obj types.Object) bool {
var hasObj bool
ast.Inspect(expr, func(n ast.Node) bool {
if n == nil {
return false
}
ident, ok := n.(*ast.Ident)
if !ok {
return true
}
objUse := info.Uses[ident]
if obj == objUse {
hasObj = true
return false
}
return false
})
return hasObj
}
type fnExtractParams struct {
tok *token.File
path []ast.Node
rng span.Range
outer *ast.FuncDecl
start ast.Node
}
// CanExtractFunction reports whether the code in the given range can be
// extracted to a function.
func CanExtractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.File) (*fnExtractParams, bool, bool, error) {
if rng.Start == rng.End {
return nil, false, false, fmt.Errorf("start and end are equal")
}
tok := fset.File(file.Pos())
if tok == nil {
return nil, false, false, fmt.Errorf("no file for pos %v", fset.Position(file.Pos()))
}
var err error
rng, err = adjustRangeForWhitespace(rng, tok, src)
if err != nil {
return nil, false, false, err
}
path, _ := astutil.PathEnclosingInterval(file, rng.Start, rng.End)
if len(path) == 0 {
return nil, false, false, fmt.Errorf("no path enclosing interval")
}
// Node that encloses the selection must be a statement.
// TODO: Support function extraction for an expression.
_, ok := path[0].(ast.Stmt)
if !ok {
return nil, false, false, fmt.Errorf("node is not a statement")
}
// Find the function declaration that encloses the selection.
var outer *ast.FuncDecl
for _, p := range path {
if p, ok := p.(*ast.FuncDecl); ok {
outer = p
break
}
}
if outer == nil {
return nil, false, false, fmt.Errorf("no enclosing function")
}
// 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 false
}
// 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 end == nil && n.End() == rng.End && n.Pos() >= rng.Start {
end = n
}
return n.Pos() <= rng.End
})
if start == nil || end == nil {
return nil, false, false, fmt.Errorf("range does not map to AST nodes")
}
// If the region is a blockStmt, use the first and last nodes in the block
// statement.
// <rng.start>{ ... }<rng.end> => { <rng.start>...<rng.end> }
if blockStmt, ok := start.(*ast.BlockStmt); ok {
if len(blockStmt.List) == 0 {
return nil, false, false, fmt.Errorf("range maps to empty block statement")
}
start, end = blockStmt.List[0], blockStmt.List[len(blockStmt.List)-1]
rng.Start, rng.End = start.Pos(), end.End()
}
return &fnExtractParams{
tok: tok,
path: path,
rng: rng,
outer: outer,
start: start,
}, true, outer.Recv != nil, nil
}
// objUsed checks if the object is used within the range. It returns the first
// occurrence of the object in the range, if it exists.
func objUsed(info *types.Info, rng span.Range, obj types.Object) (bool, *ast.Ident) {
var firstUse *ast.Ident
for id, objUse := range info.Uses {
if obj != objUse {
continue
}
if id.Pos() < rng.Start || id.End() > rng.End {
continue
}
if firstUse == nil || id.Pos() < firstUse.Pos() {
firstUse = id
}
}
return firstUse != nil, firstUse
}
// varOverridden traverses the given AST node until we find the given identifier. Then, we
// examine the occurrence of the given identifier and check for (1) whether the identifier
// is being redefined. If the identifier is free, we also check for (2) whether the identifier
// is being reassigned. We will not include an identifier in the return statement of the
// extracted function if it meets one of the above conditions.
func varOverridden(info *types.Info, firstUse *ast.Ident, obj types.Object, isFree bool, node ast.Node) bool {
var isOverriden bool
ast.Inspect(node, func(n ast.Node) bool {
if n == nil {
return false
}
assignment, ok := n.(*ast.AssignStmt)
if !ok {
return true
}
// A free variable is initialized prior to the selection. We can always reassign
// this variable after the selection because it has already been defined.
// Conversely, a non-free variable is initialized within the selection. Thus, we
// cannot reassign this variable after the selection unless it is initialized and
// returned by the extracted function.
if !isFree && assignment.Tok == token.ASSIGN {
return false
}
for _, assigned := range assignment.Lhs {
ident, ok := assigned.(*ast.Ident)
// Check if we found the first use of the identifier.
if !ok || ident != firstUse {
continue
}
objUse := info.Uses[ident]
if objUse == nil || objUse != obj {
continue
}
// Ensure that the object is not used in its own definition.
// For example:
// var f float64
// f, e := math.Frexp(f)
for _, expr := range assignment.Rhs {
if referencesObj(info, expr, obj) {
return false
}
}
isOverriden = true
return false
}
return false
})
return isOverriden
}
// 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, hasNonNestedReturns bool) ([]*returnVariable, *ast.IfStmt, error) {
var retVars []*returnVariable
var cond *ast.Ident
if !hasNonNestedReturns {
// Generate information for the added bool value.
name, _ := generateAvailableIdentifier(pos, file, path, info, "shouldReturn", 0)
cond = &ast.Ident{Name: name}
retVars = append(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 {
idx := 0
for _, 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")
}
var name string
name, idx = generateAvailableIdentifier(pos, file,
path, info, "returnValue", idx)
retVars = append(retVars, &returnVariable{
name: ast.NewIdent(name),
decl: &ast.Field{Type: expr},
zeroVal: analysisinternal.ZeroValue(
fset, file, pkg, typ),
})
}
}
var ifReturn *ast.IfStmt
if !hasNonNestedReturns {
// 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 false
}
if n, ok := n.(*ast.ReturnStmt); ok {
n.Results = append(zeroVals, n.Results...)
return false
}
return true
})
return nil
}
// generateFuncCall constructs a call expression for the extracted function, described by the
// given parameters and return variables.
func generateFuncCall(hasNonNestedReturn, hasReturnVals bool, params, returns []ast.Expr, name string, token token.Token, selector string) ast.Node {
var replace ast.Node
callExpr := &ast.CallExpr{
Fun: ast.NewIdent(name),
Args: params,
}
if selector != "" {
callExpr = &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: ast.NewIdent(selector),
Sel: ast.NewIdent(name),
},
Args: params,
}
}
if hasReturnVals {
if hasNonNestedReturn {
// Create a return statement that returns the result of the function call.
replace = &ast.ReturnStmt{
Return: 0,
Results: []ast.Expr{callExpr},
}
} else {
// Assign the result of the function call.
replace = &ast.AssignStmt{
Lhs: returns,
Tok: token,
Rhs: []ast.Expr{callExpr},
}
}
} else {
replace = callExpr
}
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 all of the variables in the
// extracted function's return statement are either not defined prior to the extracted block
// or can be safely redefined. However, for example, if z is already defined
// in a different scope, we replace the selected block with:
//
// var x int
// var y string
// x, y, z = fn()
func initializeVars(uninitialized []types.Object, retVars []*returnVariable, seenUninitialized map[types.Object]struct{}, seenVars map[types.Object]ast.Expr) []ast.Stmt {
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})
}
// 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
}
// 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
}