internal/lsp: enhance fillstruct and fillreturns to fill with variables

In the previous implementation, we always created a default
value for each type in the struct or return statement in fillstruct
and fillreturns, respectively. Now, we try to find a variable in scope
that matches the expected type. If we find multiple matches, we choose
the variable that is named most similarly to the type. If we do not
find a variable that matches, we maintain the previous functionality.

Change-Id: I3acb7e27476afaa71aaff9ffb69445913575e2b6
Reviewed-on: https://go-review.googlesource.com/c/tools/+/245130
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

@@ -14,6 +14,7 @@ import (
 	"strings"
 
 	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/internal/lsp/fuzzy"
 )
 
 var (
@@ -302,3 +303,119 @@ func WalkASTWithParent(n ast.Node, f func(n ast.Node, parent ast.Node) bool) {
 		return f(n, parent)
 	})
 }
+
+// FindMatchingIdents finds all identifiers in 'node' that match any of the given types.
+// 'pos' represents the position at which the identifiers may be inserted. 'pos' must be within
+// the scope of each of identifier we select. Otherwise, we will insert a variable at 'pos' that
+// is unrecognized.
+func FindMatchingIdents(typs []types.Type, node ast.Node, pos token.Pos, info *types.Info, pkg *types.Package) map[types.Type][]*ast.Ident {
+	matches := map[types.Type][]*ast.Ident{}
+	// Initialize matches to contain the variable types we are searching for.
+	for _, typ := range typs {
+		if typ == nil {
+			continue
+		}
+		matches[typ] = []*ast.Ident{}
+	}
+	seen := map[types.Object]struct{}{}
+	ast.Inspect(node, func(n ast.Node) bool {
+		if n == nil {
+			return false
+		}
+		// Prevent circular definitions. If 'pos' is within an assignment statement, do not
+		// allow any identifiers in that assignment statement to be selected. Otherwise,
+		// we could do the following, where 'x' satisfies the type of 'f0':
+		//
+		// x := fakeStruct{f0: x}
+		//
+		assignment, ok := n.(*ast.AssignStmt)
+		if ok && pos > assignment.Pos() && pos <= assignment.End() {
+			return false
+		}
+		if n.End() > pos {
+			return n.Pos() <= pos
+		}
+		ident, ok := n.(*ast.Ident)
+		if !ok || ident.Name == "_" {
+			return true
+		}
+		obj := info.Defs[ident]
+		if obj == nil || obj.Type() == nil {
+			return true
+		}
+		if _, ok := obj.(*types.TypeName); ok {
+			return true
+		}
+		// Prevent duplicates in matches' values.
+		if _, ok = seen[obj]; ok {
+			return true
+		}
+		seen[obj] = struct{}{}
+		// Find the scope for the given position. Then, check whether the object
+		// exists within the scope.
+		innerScope := pkg.Scope().Innermost(pos)
+		if innerScope == nil {
+			return true
+		}
+		_, foundObj := innerScope.LookupParent(ident.Name, pos)
+		if foundObj != obj {
+			return true
+		}
+		// The object must match one of the types that we are searching for.
+		if idents, ok := matches[obj.Type()]; ok {
+			matches[obj.Type()] = append(idents, ast.NewIdent(ident.Name))
+		}
+		// If the object type does not exactly match any of the target types, greedily
+		// find the first target type that the object type can satisfy.
+		for typ := range matches {
+			if obj.Type() == typ {
+				continue
+			}
+			if equivalentTypes(obj.Type(), typ) {
+				matches[typ] = append(matches[typ], ast.NewIdent(ident.Name))
+			}
+		}
+		return true
+	})
+	return matches
+}
+
+func equivalentTypes(want, got types.Type) bool {
+	if want == got || types.Identical(want, got) {
+		return true
+	}
+	// Code segment to help check for untyped equality from (golang/go#32146).
+	if rhs, ok := want.(*types.Basic); ok && rhs.Info()&types.IsUntyped > 0 {
+		if lhs, ok := got.Underlying().(*types.Basic); ok {
+			return rhs.Info()&types.IsConstType == lhs.Info()&types.IsConstType
+		}
+	}
+	return types.AssignableTo(want, got)
+}
+
+// FindBestMatch employs fuzzy matching to evaluate the similarity of each given identifier to the
+// given pattern. We return the identifier whose name is most similar to the pattern.
+func FindBestMatch(pattern string, idents []*ast.Ident) ast.Expr {
+	fuzz := fuzzy.NewMatcher(pattern)
+	var bestFuzz ast.Expr
+	highScore := float32(-1) // minimum score is -1 (no match)
+	for _, ident := range idents {
+		// TODO: Improve scoring algorithm.
+		score := fuzz.Score(ident.Name)
+		if score > highScore {
+			highScore = score
+			bestFuzz = ident
+		} else if score == -1 {
+			// Order matters in the fuzzy matching algorithm. If we find no match
+			// when matching the target to the identifier, try matching the identifier
+			// to the target.
+			revFuzz := fuzzy.NewMatcher(ident.Name)
+			revScore := revFuzz.Score(pattern)
+			if revScore > highScore {
+				highScore = revScore
+				bestFuzz = ident
+			}
+		}
+	}
+	return bestFuzz
+}

internal/lsp/analysis/fillreturns/fillreturns.go

@@ -89,7 +89,7 @@ outer:
 			return nil, nil
 		}
 
-		// Get the function that encloses the ReturnStmt.
+		// Get the function type that encloses the ReturnStmt.
 		var enclosingFunc *ast.FuncType
 		for _, n := range path {
 			switch node := n.(type) {
@@ -106,6 +106,18 @@ outer:
 			continue
 		}
 
+		// Find the function declaration that encloses the ReturnStmt.
+		var outer *ast.FuncDecl
+		for _, p := range path {
+			if p, ok := p.(*ast.FuncDecl); ok {
+				outer = p
+				break
+			}
+		}
+		if outer == nil {
+			return nil, nil
+		}
+
 		// Skip any return statements that contain function calls with multiple return values.
 		for _, expr := range ret.Results {
 			e, ok := expr.(*ast.CallExpr)
@@ -126,13 +138,17 @@ outer:
 		// For each value in the return function declaration, find the leftmost element
 		// in the return statement that has the desired type. If no such element exits,
 		// fill in the missing value with the appropriate "zero" value.
-		for i, result := range enclosingFunc.Results.List {
-			typ := info.TypeOf(result.Type)
-
+		var retTyps []types.Type
+		for _, ret := range enclosingFunc.Results.List {
+			retTyps = append(retTyps, info.TypeOf(ret.Type))
+		}
+		matches :=
+			analysisinternal.FindMatchingIdents(retTyps, file, ret.Pos(), info, pass.Pkg)
+		for i, retTyp := range retTyps {
 			var match ast.Expr
 			var idx int
 			for j, val := range remaining {
-				if !matchingTypes(info.TypeOf(val), typ) {
+				if !matchingTypes(info.TypeOf(val), retTyp) {
 					continue
 				}
 				if !analysisinternal.IsZeroValue(val) {
@@ -149,12 +165,22 @@ outer:
 				fixed[i] = match
 				remaining = append(remaining[:idx], remaining[idx+1:]...)
 			} else {
-				zv := analysisinternal.ZeroValue(pass.Fset, file, pass.Pkg,
-					info.TypeOf(result.Type))
-				if zv == nil {
+				idents, ok := matches[retTyp]
+				if !ok {
+					return nil, fmt.Errorf("invalid return type: %v", retTyp)
+				}
+				// Find the identifer whose name is most similar to the return type.
+				// If we do not find any identifer that matches the pattern,
+				// generate a zero value.
+				value := analysisinternal.FindBestMatch(retTyp.String(), idents)
+				if value == nil {
+					value = analysisinternal.ZeroValue(
+						pass.Fset, file, pass.Pkg, retTyp)
+				}
+				if value == nil {
 					return nil, nil
 				}
-				fixed[i] = zv
+				fixed[i] = value
 			}
 		}
 

internal/lsp/analysis/fillreturns/testdata/src/a/a.go

@@ -6,6 +6,7 @@ package fillreturns
 
 import (
 	"errors"
+	"go/ast"
 	ast2 "go/ast"
 	"io"
 	"net/http"
@@ -120,3 +121,17 @@ func gotTooMany() int {
 	}
 	return 0, 5, false // want "wrong number of return values \\(want 1, got 3\\)"
 }
+
+func fillVars() (int, string, ast.Node, bool, error) {
+	eint := 0
+	s := "a"
+	var t bool
+	if true {
+		err := errors.New("fail")
+		return // want "wrong number of return values \\(want 5, got 0\\)"
+	}
+	n := ast.NewIdent("ident")
+	int := 3
+	var b bool
+	return "" // want "wrong number of return values \\(want 5, got 1\\)"
+}

internal/lsp/analysis/fillreturns/testdata/src/a/a.go.golden

@@ -6,6 +6,7 @@ package fillreturns
 
 import (
 	"errors"
+	"go/ast"
 	ast2 "go/ast"
 	"io"
 	"net/http"
@@ -120,3 +121,17 @@ func gotTooMany() int {
 	}
 	return 5 // want "wrong number of return values \\(want 1, got 3\\)"
 }
+
+func fillVars() (int, string, ast.Node, bool, error) {
+	eint := 0
+	s := "a"
+	var t bool
+	if true {
+		err := errors.New("fail")
+		return eint, s, nil, false, err // want "wrong number of return values \\(want 5, got 0\\)"
+	}
+	n := ast.NewIdent("ident")
+	int := 3
+	var b bool
+	return int, "", n, b, nil // want "wrong number of return values \\(want 5, got 1\\)"
+}

internal/lsp/analysis/fillstruct/fillstruct.go

@@ -114,7 +114,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 			name = "anonymous struct"
 		}
 		pass.Report(analysis.Diagnostic{
-			Message: fmt.Sprintf("Fill %s with default values", name),
+			Message: fmt.Sprintf("Fill %s", name),
 			Pos:     expr.Pos(),
 			End:     expr.End(),
 		})
@@ -172,18 +172,36 @@ func SuggestedFix(fset *token.FileSet, rng span.Range, content []byte, file *ast
 
 	line := 2 // account for 1-based lines and the left brace
 	var elts []ast.Expr
+	var fieldTyps []types.Type
 	for i := 0; i < fieldCount; i++ {
 		field := obj.Field(i)
-
 		// Ignore fields that are not accessible in the current package.
 		if field.Pkg() != nil && field.Pkg() != pkg && !field.Exported() {
+			fieldTyps = append(fieldTyps, nil)
 			continue
 		}
-
-		value := populateValue(fset, file, pkg, field.Type())
-		if value == nil {
+		fieldTyps = append(fieldTyps, field.Type())
+	}
+	matches := analysisinternal.FindMatchingIdents(fieldTyps, file, rng.Start, info, pkg)
+	for i, fieldTyp := range fieldTyps {
+		if fieldTyp == nil {
 			continue
 		}
+		idents, ok := matches[fieldTyp]
+		if !ok {
+			return nil, fmt.Errorf("invalid struct field type: %v", fieldTyp)
+		}
+
+		// Find the identifer whose name is most similar to the name of the field's key.
+		// If we do not find any identifer that matches the pattern, generate a new value.
+		// NOTE: We currently match on the name of the field key rather than the field type.
+		value := analysisinternal.FindBestMatch(obj.Field(i).Name(), idents)
+		if value == nil {
+			value = populateValue(fset, file, pkg, fieldTyp)
+		}
+		if value == nil {
+			return nil, nil
+		}
 
 		tok.AddLine(line - 1) // add 1 byte per line
 		if line > tok.LineCount() {
@@ -194,7 +212,7 @@ func SuggestedFix(fset *token.FileSet, rng span.Range, content []byte, file *ast
 		kv := &ast.KeyValueExpr{
 			Key: &ast.Ident{
 				NamePos: pos,
-				Name:    field.Name(),
+				Name:    obj.Field(i).Name(),
 			},
 			Colon: pos,
 			Value: value,

internal/lsp/testdata/lsp/primarymod/fillstruct/a.go

@@ -96,3 +96,39 @@ var _ = []ast.BasicLit{
 }
 
 var _ = []ast.BasicLit{{}} //@suggestedfix("}", "refactor.rewrite")
+
+type iStruct struct {
+	X int
+}
+
+type sStruct struct {
+	str string
+}
+
+type multiFill struct {
+	num   int
+	strin string
+	arr   []int
+}
+
+type assignStruct struct {
+	n ast.Node
+}
+
+func fill() {
+	var x int
+	var _ = iStruct{} //@suggestedfix("}", "refactor.rewrite")
+
+	var s string
+	var _ = sStruct{} //@suggestedfix("}", "refactor.rewrite")
+
+	var n int
+	_ = []int{}
+	if true {
+		arr := []int{1, 2}
+	}
+	var _ = multiFill{} //@suggestedfix("}", "refactor.rewrite")
+
+	var node *ast.CompositeLit
+	var _ = assignStruct{} //@suggestedfix("}", "refactor.rewrite")
+}

internal/lsp/testdata/lsp/summary.txt.golden

@@ -12,7 +12,7 @@ DiagnosticsCount = 44
 FoldingRangesCount = 2
 FormatCount = 6
 ImportCount = 8
-SuggestedFixCount = 31
+SuggestedFixCount = 35
 FunctionExtractionCount = 11
 DefinitionsCount = 63
 TypeDefinitionsCount = 2