|  | 
| 8 | 8 | 	"bytes" | 
| 9 | 9 | 	"context" | 
| 10 | 10 | 	"fmt" | 
|  | 11 | +	"go/ast" | 
| 11 | 12 | 	"go/format" | 
| 12 | 13 | 	"go/parser" | 
| 13 | 14 | 	"go/token" | 
| @@ -51,6 +52,18 @@ func stubMissingCalledFunctionFixer(ctx context.Context, snapshot *cache.Snapsho | 
| 51 | 52 | 	return insertDeclsAfter(ctx, snapshot, pkg.Metadata(), si.Fset, si.After, si.Emit) | 
| 52 | 53 | } | 
| 53 | 54 | 
 | 
|  | 55 | +// stubMissingStructFieldFixer returns a suggested fix to declare the missing | 
|  | 56 | +// field that the user may want to generate based on SelectorExpr | 
|  | 57 | +// at the cursor position. | 
|  | 58 | +func stubMissingStructFieldFixer(ctx context.Context, snapshot *cache.Snapshot, pkg *cache.Package, pgf *parsego.File, start, end token.Pos) (*token.FileSet, *analysis.SuggestedFix, error) { | 
|  | 59 | +	nodes, _ := astutil.PathEnclosingInterval(pgf.File, start, end) | 
|  | 60 | +	fi := GetFieldStubInfo(pkg.FileSet(), pkg.TypesInfo(), nodes) | 
|  | 61 | +	if fi == nil { | 
|  | 62 | +		return nil, nil, fmt.Errorf("invalid type request") | 
|  | 63 | +	} | 
|  | 64 | +	return insertStructField(ctx, snapshot, pkg.Metadata(), fi) | 
|  | 65 | +} | 
|  | 66 | + | 
| 54 | 67 | // An emitter writes new top-level declarations into an existing | 
| 55 | 68 | // file. References to symbols should be qualified using qual, which | 
| 56 | 69 | // respects the local import environment. | 
| @@ -238,3 +251,66 @@ func trimVersionSuffix(path string) string { | 
| 238 | 251 | 	} | 
| 239 | 252 | 	return path | 
| 240 | 253 | } | 
|  | 254 | + | 
|  | 255 | +func insertStructField(ctx context.Context, snapshot *cache.Snapshot, meta *metadata.Package, fieldInfo *StructFieldInfo) (*token.FileSet, *analysis.SuggestedFix, error) { | 
|  | 256 | +	if fieldInfo == nil { | 
|  | 257 | +		return nil, nil, fmt.Errorf("no field info provided") | 
|  | 258 | +	} | 
|  | 259 | + | 
|  | 260 | +	// get the file containing the struct definition using the position | 
|  | 261 | +	declPGF, _, err := parseFull(ctx, snapshot, fieldInfo.Fset, fieldInfo.Named.Obj().Pos()) | 
|  | 262 | +	if err != nil { | 
|  | 263 | +		return nil, nil, fmt.Errorf("failed to parse file declaring struct: %w", err) | 
|  | 264 | +	} | 
|  | 265 | +	if declPGF.Fixed() { | 
|  | 266 | +		return nil, nil, fmt.Errorf("file contains parse errors: %s", declPGF.URI) | 
|  | 267 | +	} | 
|  | 268 | + | 
|  | 269 | +	// find the struct type declaration | 
|  | 270 | +	var structType *ast.StructType | 
|  | 271 | +	ast.Inspect(declPGF.File, func(n ast.Node) bool { | 
|  | 272 | +		if typeSpec, ok := n.(*ast.TypeSpec); ok { | 
|  | 273 | +			if typeSpec.Name.Name == fieldInfo.Named.Obj().Name() { | 
|  | 274 | +				if st, ok := typeSpec.Type.(*ast.StructType); ok { | 
|  | 275 | +					structType = st | 
|  | 276 | +					return false | 
|  | 277 | +				} | 
|  | 278 | +			} | 
|  | 279 | +		} | 
|  | 280 | +		return true | 
|  | 281 | +	}) | 
|  | 282 | + | 
|  | 283 | +	if structType == nil { | 
|  | 284 | +		return nil, nil, fmt.Errorf("could not find struct definition") | 
|  | 285 | +	} | 
|  | 286 | + | 
|  | 287 | +	// find the position to insert the new field (end of struct fields) | 
|  | 288 | +	insertPos := structType.Fields.Closing - 1 | 
|  | 289 | +	if insertPos == structType.Fields.Opening { | 
|  | 290 | +		// struct has no fields yet | 
|  | 291 | +		insertPos = structType.Fields.Closing | 
|  | 292 | +	} | 
|  | 293 | + | 
|  | 294 | +	var buf bytes.Buffer | 
|  | 295 | +	if err := fieldInfo.Emit(&buf, types.RelativeTo(fieldInfo.Named.Obj().Pkg())); err != nil { | 
|  | 296 | +		return nil, nil, err | 
|  | 297 | +	} | 
|  | 298 | + | 
|  | 299 | +	_, err = declPGF.Mapper.PosRange(declPGF.Tok, insertPos, insertPos) | 
|  | 300 | +	if err != nil { | 
|  | 301 | +		return nil, nil, err | 
|  | 302 | +	} | 
|  | 303 | + | 
|  | 304 | +	textEdit := analysis.TextEdit{ | 
|  | 305 | +		Pos:     insertPos, | 
|  | 306 | +		End:     insertPos, | 
|  | 307 | +		NewText: []byte(buf.String()), | 
|  | 308 | +	} | 
|  | 309 | + | 
|  | 310 | +	fix := &analysis.SuggestedFix{ | 
|  | 311 | +		Message:   fmt.Sprintf("Add field %s to struct %s", fieldInfo.Expr.Sel.Name, fieldInfo.Named.Obj().Name()), | 
|  | 312 | +		TextEdits: []analysis.TextEdit{textEdit}, | 
|  | 313 | +	} | 
|  | 314 | + | 
|  | 315 | +	return fieldInfo.Fset, fix, nil | 
|  | 316 | +} | 
0 commit comments