@@ -16,7 +16,6 @@ import (
1616
1717 "golang.org/x/tools/go/analysis"
1818 "golang.org/x/tools/go/ast/astutil"
19- "golang.org/x/tools/gopls/internal/util/safetoken"
2019 "golang.org/x/tools/gopls/internal/util/typesutil"
2120 "golang.org/x/tools/internal/analysisinternal"
2221 "golang.org/x/tools/internal/typesinternal"
@@ -87,33 +86,118 @@ func CreateUndeclared(fset *token.FileSet, start, end token.Pos, content []byte,
8786 if isCallPosition (path ) {
8887 return newFunctionDeclaration (path , file , pkg , info , fset )
8988 }
89+ var (
90+ firstRef * ast.Ident // We should insert the new declaration before the first occurrence of the undefined ident.
91+ assignTokPos token.Pos
92+ funcDecl = path [len (path )- 2 ].(* ast.FuncDecl ) // This is already ensured by [undeclaredFixTitle].
93+ parent = ast .Node (funcDecl )
94+ )
95+ // Search from enclosing FuncDecl to path[0], since we can not use := syntax outside function.
96+ // Adds the missing colon after the first undefined symbol
97+ // when it sits in lhs of an AssignStmt.
98+ ast .Inspect (funcDecl , func (n ast.Node ) bool {
99+ if n == nil || firstRef != nil {
100+ return false
101+ }
102+ if n , ok := n .(* ast.Ident ); ok && n .Name == ident .Name && info .ObjectOf (n ) == nil {
103+ firstRef = n
104+ // Only consider adding colon at the first occurrence.
105+ if pos , ok := replaceableAssign (info , n , parent ); ok {
106+ assignTokPos = pos
107+ return false
108+ }
109+ }
110+ parent = n
111+ return true
112+ })
113+ if assignTokPos .IsValid () {
114+ return fset , & analysis.SuggestedFix {
115+ TextEdits : []analysis.TextEdit {{
116+ Pos : assignTokPos ,
117+ End : assignTokPos ,
118+ NewText : []byte (":" ),
119+ }},
120+ }, nil
121+ }
90122
91- // Get the place to insert the new statement.
92- insertBeforeStmt := analysisinternal .StmtToInsertVarBefore (path )
123+ // firstRef should never be nil, at least one ident at cursor position should be found,
124+ // but be defensive.
125+ if firstRef == nil {
126+ return nil , nil , fmt .Errorf ("no identifier found" )
127+ }
128+ p , _ := astutil .PathEnclosingInterval (file , firstRef .Pos (), firstRef .Pos ())
129+ insertBeforeStmt := analysisinternal .StmtToInsertVarBefore (p )
93130 if insertBeforeStmt == nil {
94131 return nil , nil , fmt .Errorf ("could not locate insertion point" )
95132 }
96-
97- insertBefore := safetoken .StartPosition (fset , insertBeforeStmt .Pos ()).Offset
98-
99- // Get the indent to add on the line after the new statement.
100- // Since this will have a parse error, we can not use format.Source().
101- contentBeforeStmt , indent := content [:insertBefore ], "\n "
102- if nl := bytes .LastIndex (contentBeforeStmt , []byte ("\n " )); nl != - 1 {
103- indent = string (contentBeforeStmt [nl :])
133+ indent , err := calculateIndentation (content , fset .File (file .FileStart ), insertBeforeStmt )
134+ if err != nil {
135+ return nil , nil , err
136+ }
137+ typs := typesutil .TypesFromContext (info , path , start )
138+ if typs == nil {
139+ // Default to 0.
140+ typs = []types.Type {types .Typ [types .Int ]}
141+ }
142+ assignStmt := & ast.AssignStmt {
143+ Lhs : []ast.Expr {ast .NewIdent (ident .Name )},
144+ Tok : token .DEFINE ,
145+ Rhs : []ast.Expr {typesinternal .ZeroExpr (file , pkg , typs [0 ])},
146+ }
147+ var buf bytes.Buffer
148+ if err := format .Node (& buf , fset , assignStmt ); err != nil {
149+ return nil , nil , err
104150 }
151+ newLineIndent := "\n " + indent
152+ assignment := strings .ReplaceAll (buf .String (), "\n " , newLineIndent ) + newLineIndent
105153
106- // Create the new local variable statement.
107- newStmt := fmt .Sprintf ("%s := %s" , ident .Name , indent )
108154 return fset , & analysis.SuggestedFix {
109- TextEdits : []analysis.TextEdit {{
110- Pos : insertBeforeStmt .Pos (),
111- End : insertBeforeStmt .Pos (),
112- NewText : []byte (newStmt ),
113- }},
155+ TextEdits : []analysis.TextEdit {
156+ {
157+ Pos : insertBeforeStmt .Pos (),
158+ End : insertBeforeStmt .Pos (),
159+ NewText : []byte (assignment ),
160+ },
161+ },
114162 }, nil
115163}
116164
165+ // replaceableAssign returns position of token.ASSIGN if ident meets the following conditions:
166+ // 1) parent node must be an *ast.AssignStmt with Tok set to token.ASSIGN.
167+ // 2) ident must not be self assignment.
168+ //
169+ // For example, we should not add a colon when
170+ // a = a + 1
171+ // ^ ^ cursor here
172+ func replaceableAssign (info * types.Info , ident * ast.Ident , parent ast.Node ) (token.Pos , bool ) {
173+ var pos token.Pos
174+ if assign , ok := parent .(* ast.AssignStmt ); ok && assign .Tok == token .ASSIGN {
175+ for _ , rhs := range assign .Rhs {
176+ if referencesIdent (info , rhs , ident ) {
177+ return pos , false
178+ }
179+ }
180+ return assign .TokPos , true
181+ }
182+ return pos , false
183+ }
184+
185+ // referencesIdent checks whether the given undefined ident appears in the given expression.
186+ func referencesIdent (info * types.Info , expr ast.Expr , ident * ast.Ident ) bool {
187+ var hasIdent bool
188+ ast .Inspect (expr , func (n ast.Node ) bool {
189+ if n == nil {
190+ return false
191+ }
192+ if i , ok := n .(* ast.Ident ); ok && i .Name == ident .Name && info .ObjectOf (i ) == nil {
193+ hasIdent = true
194+ return false
195+ }
196+ return true
197+ })
198+ return hasIdent
199+ }
200+
117201func newFunctionDeclaration (path []ast.Node , file * ast.File , pkg * types.Package , info * types.Info , fset * token.FileSet ) (* token.FileSet , * analysis.SuggestedFix , error ) {
118202 if len (path ) < 3 {
119203 return nil , nil , fmt .Errorf ("unexpected set of enclosing nodes: %v" , path )
0 commit comments