gopls/internal/undeclaredname: add missing colon when appropriate

gopls/doc/features/diagnostics.md

@@ -213,12 +213,13 @@ func main() {
 }
 ```
 
-The quick fix would be:
+The quick fix would insert a declaration with a default
+value inferring its type from the context:
 
 ```go
 func main() {
   x := 42
-  y :=
+  y := 0
   min(x, y)
 }
 ```

gopls/doc/release/v0.17.0.md

@@ -33,10 +33,10 @@ removed as needed.
 
 The user can invoke this code action by selecting a function name, the keywords
 `func`, `const`, `var`, `type`, or by placing the caret on them without selecting,
-or by selecting a whole declaration or multiple declrations.
+or by selecting a whole declaration or multiple declarations.
 
 In order to avoid ambiguity and surprise about what to extract, some kinds
-of paritial selection of a declration cannot invoke this code action.
+of paritial selection of a declaration cannot invoke this code action.
 
 ## Pull diagnostics
 

gopls/internal/golang/undeclared.go

@@ -16,7 +16,6 @@ import (
 
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/ast/astutil"
-	"golang.org/x/tools/gopls/internal/util/safetoken"
 	"golang.org/x/tools/gopls/internal/util/typesutil"
 	"golang.org/x/tools/internal/analysisinternal"
 	"golang.org/x/tools/internal/typesinternal"
@@ -87,33 +86,118 @@ func CreateUndeclared(fset *token.FileSet, start, end token.Pos, content []byte,
 	if isCallPosition(path) {
 		return newFunctionDeclaration(path, file, pkg, info, fset)
 	}
+	var (
+		firstRef     *ast.Ident // We should insert the new declaration before the first occurrence of the undefined ident.
+		assignTokPos token.Pos
+		funcDecl     = path[len(path)-2].(*ast.FuncDecl) // This is already ensured by [undeclaredFixTitle].
+		parent       = ast.Node(funcDecl)
+	)
+	// Search from enclosing FuncDecl to path[0], since we can not use := syntax outside function.
+	// Adds the missing colon after the first undefined symbol
+	// when it sits in lhs of an AssignStmt.
+	ast.Inspect(funcDecl, func(n ast.Node) bool {
+		if n == nil || firstRef != nil {
+			return false
+		}
+		if n, ok := n.(*ast.Ident); ok && n.Name == ident.Name && info.ObjectOf(n) == nil {
+			firstRef = n
+			// Only consider adding colon at the first occurrence.
+			if pos, ok := replaceableAssign(info, n, parent); ok {
+				assignTokPos = pos
+				return false
+			}
+		}
+		parent = n
+		return true
+	})
+	if assignTokPos.IsValid() {
+		return fset, &analysis.SuggestedFix{
+			TextEdits: []analysis.TextEdit{{
+				Pos:     assignTokPos,
+				End:     assignTokPos,
+				NewText: []byte(":"),
+			}},
+		}, nil
+	}
 
-	// Get the place to insert the new statement.
-	insertBeforeStmt := analysisinternal.StmtToInsertVarBefore(path)
+	// firstRef should never be nil, at least one ident at cursor position should be found,
+	// but be defensive.
+	if firstRef == nil {
+		return nil, nil, fmt.Errorf("no identifier found")
+	}
+	p, _ := astutil.PathEnclosingInterval(file, firstRef.Pos(), firstRef.Pos())
+	insertBeforeStmt := analysisinternal.StmtToInsertVarBefore(p)
 	if insertBeforeStmt == nil {
 		return nil, nil, fmt.Errorf("could not locate insertion point")
 	}
-
-	insertBefore := safetoken.StartPosition(fset, insertBeforeStmt.Pos()).Offset
-
-	// Get the indent to add on the line after the new statement.
-	// Since this will have a parse error, we can not use format.Source().
-	contentBeforeStmt, indent := content[:insertBefore], "\n"
-	if nl := bytes.LastIndex(contentBeforeStmt, []byte("\n")); nl != -1 {
-		indent = string(contentBeforeStmt[nl:])
+	indent, err := calculateIndentation(content, fset.File(file.FileStart), insertBeforeStmt)
+	if err != nil {
+		return nil, nil, err
+	}
+	typs := typesutil.TypesFromContext(info, path, start)
+	if typs == nil {
+		// Default to 0.
+		typs = []types.Type{types.Typ[types.Int]}
+	}
+	assignStmt := &ast.AssignStmt{
+		Lhs: []ast.Expr{ast.NewIdent(ident.Name)},
+		Tok: token.DEFINE,
+		Rhs: []ast.Expr{typesinternal.ZeroExpr(file, pkg, typs[0])},
+	}
+	var buf bytes.Buffer
+	if err := format.Node(&buf, fset, assignStmt); err != nil {
+		return nil, nil, err
 	}
+	newLineIndent := "\n" + indent
+	assignment := strings.ReplaceAll(buf.String(), "\n", newLineIndent) + newLineIndent
 
-	// Create the new local variable statement.
-	newStmt := fmt.Sprintf("%s := %s", ident.Name, indent)
 	return fset, &analysis.SuggestedFix{
-		TextEdits: []analysis.TextEdit{{
-			Pos:     insertBeforeStmt.Pos(),
-			End:     insertBeforeStmt.Pos(),
-			NewText: []byte(newStmt),
-		}},
+		TextEdits: []analysis.TextEdit{
+			{
+				Pos:     insertBeforeStmt.Pos(),
+				End:     insertBeforeStmt.Pos(),
+				NewText: []byte(assignment),
+			},
+		},
 	}, nil
 }
 
+// replaceableAssign returns position of token.ASSIGN if ident meets the following conditions:
+// 1) parent node must be an *ast.AssignStmt with Tok set to token.ASSIGN.
+// 2) ident must not be self assignment.
+//
+// For example, we should not add a colon when
+// a = a + 1
+// ^   ^ cursor here
+func replaceableAssign(info *types.Info, ident *ast.Ident, parent ast.Node) (token.Pos, bool) {
+	var pos token.Pos
+	if assign, ok := parent.(*ast.AssignStmt); ok && assign.Tok == token.ASSIGN {
+		for _, rhs := range assign.Rhs {
+			if referencesIdent(info, rhs, ident) {
+				return pos, false
+			}
+		}
+		return assign.TokPos, true
+	}
+	return pos, false
+}
+
+// referencesIdent checks whether the given undefined ident appears in the given expression.
+func referencesIdent(info *types.Info, expr ast.Expr, ident *ast.Ident) bool {
+	var hasIdent bool
+	ast.Inspect(expr, func(n ast.Node) bool {
+		if n == nil {
+			return false
+		}
+		if i, ok := n.(*ast.Ident); ok && i.Name == ident.Name && info.ObjectOf(i) == nil {
+			hasIdent = true
+			return false
+		}
+		return true
+	})
+	return hasIdent
+}
+
 func newFunctionDeclaration(path []ast.Node, file *ast.File, pkg *types.Package, info *types.Info, fset *token.FileSet) (*token.FileSet, *analysis.SuggestedFix, error) {
 	if len(path) < 3 {
 		return nil, nil, fmt.Errorf("unexpected set of enclosing nodes: %v", path)

gopls/internal/test/marker/testdata/quickfix/undeclared/undeclared_variable.txt

@@ -0,0 +1,108 @@
+Tests of suggested fixes for "undeclared name" diagnostics,
+which are of ("compiler", "error") type.
+
+-- flags --
+-ignore_extra_diags
+
+-- go.mod --
+module example.com
+go 1.12
+
+-- a.go --
+package undeclared_var
+
+func a() {
+	z, _ := 1+y, 11 //@quickfix("y", re"(undeclared name|undefined): y", a)
+	_ = z
+}
+
+-- @a/a.go --
+@@ -4 +4 @@
++	y := 0
+-- b.go --
+package undeclared_var
+
+func b() {
+	if 100 < 90 {
+	} else if 100 > n+2 { //@quickfix("n", re"(undeclared name|undefined): n", b)
+	}
+}
+
+-- @b/b.go --
+@@ -4 +4 @@
++	n := 0
+-- c.go --
+package undeclared_var
+
+func c() {
+	for i < 200 { //@quickfix("i", re"(undeclared name|undefined): i", c)
+	}
+	r() //@diag("r", re"(undeclared name|undefined): r")
+}
+
+-- @c/c.go --
+@@ -4 +4 @@
++	i := 0
+-- add_colon.go --
+package undeclared_var
+
+func addColon() {
+	ac = 1 //@quickfix("ac", re"(undeclared name|undefined): ac", add_colon)
+}
+
+-- @add_colon/add_colon.go --
+@@ -4 +4 @@
+-	ac = 1 //@quickfix("ac", re"(undeclared name|undefined): ac", add_colon)
++	ac := 1 //@quickfix("ac", re"(undeclared name|undefined): ac", add_colon)
+-- add_colon_first.go --
+package undeclared_var
+
+func addColonAtFirstStmt() {
+	ac = 1
+	ac = 2
+	ac = 3
+	b := ac //@quickfix("ac", re"(undeclared name|undefined): ac", add_colon_first)
+}
+
+-- @add_colon_first/add_colon_first.go --
+@@ -4 +4 @@
+-	ac = 1
++	ac := 1
+-- self_assign.go --
+package undeclared_var
+
+func selfAssign() {
+	ac = ac + 1
+	ac = ac + 2 //@quickfix("ac", re"(undeclared name|undefined): ac", lhs)
+	ac = ac + 3 //@quickfix("ac + 3", re"(undeclared name|undefined): ac", rhs)
+}
+
+-- @lhs/self_assign.go --
+@@ -4 +4 @@
++	ac := nil
+-- @rhs/self_assign.go --
+@@ -4 +4 @@
++	ac := 0
+-- correct_type.go --
+package undeclared_var
+import "fmt"
+func selfAssign() {
+	fmt.Printf(ac) //@quickfix("ac", re"(undeclared name|undefined): ac", string)
+}
+-- @string/correct_type.go --
+@@ -4 +4 @@
++	ac := ""
+-- ignore.go --
+package undeclared_var
+import "fmt"
+type Foo struct {
+	bar int
+}
+func selfAssign() {
+	f := Foo{}
+	b = f.bar
+	c := bar //@quickfix("bar", re"(undeclared name|undefined): bar", ignore)
+}
+-- @ignore/ignore.go --
+@@ -9 +9 @@
++	bar := nil

gopls/internal/test/marker/testdata/quickfix/undeclaredfunc.txt → gopls/internal/test/marker/testdata/quickfix/undeclared/undeclaredfunc.txt

@@ -1,8 +1,6 @@
 This test checks the quick fix for "undeclared: f" that declares the
 missing function. See #47558.
 
-TODO(adonovan): infer the result variables from the context (int, in this case).
-
 -- a.go --
 package a