From 5f30384a38fbaf470202e33139000174892ad371 Mon Sep 17 00:00:00 2001
From: Denis <candyboobers@gmail.com>
Date: Tue, 3 Dec 2024 22:04:24 +0100
Subject: [PATCH] x/tools/gopls: implement struct field generation quickfix

---
 gopls/doc/features/diagnostics.md             |  29 +++
 gopls/doc/release/v0.17.0.md                  |  10 +-
 gopls/internal/golang/codeaction.go           |  83 ++++++++-
 gopls/internal/golang/fix.go                  |   2 +
 gopls/internal/golang/stub.go                 |  76 ++++++++
 .../marker/testdata/quickfix/struct_field.txt | 132 ++++++++++++++
 gopls/internal/util/typesutil/typesutil.go    | 167 ++++++++++++++++++
 7 files changed, 496 insertions(+), 3 deletions(-)
 create mode 100644 gopls/internal/test/marker/testdata/quickfix/struct_field.txt

diff --git a/gopls/doc/features/diagnostics.md b/gopls/doc/features/diagnostics.md
index 21015bcaa35..4391a85f72a 100644
--- a/gopls/doc/features/diagnostics.md
+++ b/gopls/doc/features/diagnostics.md
@@ -248,6 +248,35 @@ func doSomething(i int) string {
   panic("unimplemented")
 }
 ```
+
+### `StubMissingStructField`: Declare missing field T.f
+
+When you attempt to access a field on a type that does not have the field,
+the compiler will report an error such as "type X has no field or method Y".
+In this scenario, gopls now offers a quick fix to generate a stub declaration of
+the missing field, inferring its type from the accessing type or assigning a designated value.
+
+Consider the following code where `Foo` does not have a field `bar`:
+
+```go
+type Foo struct{}
+
+func main() {
+  var s string
+  f := Foo{}
+  s = f.bar // error: f.bar undefined (type Foo has no field or method bar)
+}
+```
+
+Gopls will offer a quick fix, "Declare missing field Foo.bar".
+When invoked, it creates the following declaration:
+
+```go
+type Foo struct{
+  bar string
+}
+```
+
 <!--
 
 dorky details and deletia:
diff --git a/gopls/doc/release/v0.17.0.md b/gopls/doc/release/v0.17.0.md
index dcd2ef06502..4e5b3fa46b0 100644
--- a/gopls/doc/release/v0.17.0.md
+++ b/gopls/doc/release/v0.17.0.md
@@ -113,4 +113,12 @@ into account its signature, including input parameters and results.
 Since this feature is implemented by the server (gopls), it is compatible with
 all LSP-compliant editors. VS Code users may continue to use the client-side
 `Go: Generate Unit Tests For file/function/package` command which utilizes the
-[gotests](https://github.com/cweill/gotests) tool.
\ No newline at end of file
+[gotests](https://github.com/cweill/gotests) tool.
+
+## Generate missing struct field from access
+When you attempt to access a field on a type that does not have the field,
+the compiler will report an error like “type X has no field or method Y”.
+Gopls now offers a new code action, “Declare missing field of T.f”,
+where T is the concrete type and f is the undefined field.
+The stub field's signature is inferred
+from the context of the access.
diff --git a/gopls/internal/golang/codeaction.go b/gopls/internal/golang/codeaction.go
index 27ed29faf06..7196e83effc 100644
--- a/gopls/internal/golang/codeaction.go
+++ b/gopls/internal/golang/codeaction.go
@@ -5,6 +5,7 @@
 package golang
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -322,14 +323,23 @@ func quickFix(ctx context.Context, req *codeActionsRequest) error {
 			}
 
 		// "type X has no field or method Y" compiler error.
-		// Offer a "Declare missing method T.f" code action.
-		// See [stubMissingCalledFunctionFixer] for command implementation.
 		case strings.Contains(msg, "has no field or method"):
 			path, _ := astutil.PathEnclosingInterval(req.pgf.File, start, end)
+
+			// Offer a "Declare missing method T.f" code action if a CallStubInfo found.
+			// See [stubMissingCalledFunctionFixer] for command implementation.
 			si := stubmethods.GetCallStubInfo(req.pkg.FileSet(), info, path, start)
 			if si != nil {
 				msg := fmt.Sprintf("Declare missing method %s.%s", si.Receiver.Obj().Name(), si.MethodName)
 				req.addApplyFixAction(msg, fixMissingCalledFunction, req.loc)
+			} else {
+				// Offer a "Declare missing field T.f" code action.
+				// See [stubMissingStructFieldFixer] for command implementation.
+				fi := GetFieldStubInfo(req.pkg.FileSet(), info, path)
+				if fi != nil {
+					msg := fmt.Sprintf("Declare missing struct field %s.%s", fi.Named.Obj().Name(), fi.Expr.Sel.Name)
+					req.addApplyFixAction(msg, fixMissingStructField, req.loc)
+				}
 			}
 
 		// "undeclared name: x" or "undefined: x" compiler error.
@@ -348,6 +358,75 @@ func quickFix(ctx context.Context, req *codeActionsRequest) error {
 	return nil
 }
 
+func GetFieldStubInfo(fset *token.FileSet, info *types.Info, path []ast.Node) *StructFieldInfo {
+	for _, node := range path {
+		n, ok := node.(*ast.SelectorExpr)
+		if !ok {
+			continue
+		}
+		tv, ok := info.Types[n.X]
+		if !ok {
+			break
+		}
+
+		named, ok := tv.Type.(*types.Named)
+		if !ok {
+			break
+		}
+
+		structType, ok := named.Underlying().(*types.Struct)
+		if !ok {
+			break
+		}
+
+		return &StructFieldInfo{
+			Fset:   fset,
+			Expr:   n,
+			Struct: structType,
+			Named:  named,
+			Info:   info,
+			Path:   path,
+		}
+	}
+
+	return nil
+}
+
+type StructFieldInfo struct {
+	Fset   *token.FileSet
+	Expr   *ast.SelectorExpr
+	Struct *types.Struct
+	Named  *types.Named
+	Info   *types.Info
+	Path   []ast.Node
+}
+
+// Emit writes to out the missing field based on type info.
+func (si *StructFieldInfo) Emit(out *bytes.Buffer, qual types.Qualifier) error {
+	if si.Expr == nil || si.Expr.Sel == nil {
+		return fmt.Errorf("invalid selector expression")
+	}
+
+	// Get types from context at the selector expression position
+	typesFromContext := typesutil.TypesFromContext(si.Info, si.Path, si.Expr.Pos())
+
+	// Default to interface{} if we couldn't determine the type from context
+	var fieldType types.Type
+	if len(typesFromContext) > 0 && typesFromContext[0] != nil {
+		fieldType = typesFromContext[0]
+	} else {
+		// Create a new interface{} type
+		fieldType = types.NewInterfaceType(nil, nil)
+	}
+
+	tpl := "\n\t%s %s"
+	if si.Struct.NumFields() == 0 {
+		tpl += "\n"
+	}
+	fmt.Fprintf(out, tpl, si.Expr.Sel.Name, types.TypeString(fieldType, qual))
+	return nil
+}
+
 // allImportsFixesResult is the result of a lazy call to allImportsFixes.
 // It implements the codeActionsRequest lazyInit interface.
 type allImportsFixesResult struct {
diff --git a/gopls/internal/golang/fix.go b/gopls/internal/golang/fix.go
index 68110206add..bfd5cb364c9 100644
--- a/gopls/internal/golang/fix.go
+++ b/gopls/internal/golang/fix.go
@@ -68,6 +68,7 @@ const (
 	fixCreateUndeclared        = "create_undeclared"
 	fixMissingInterfaceMethods = "stub_missing_interface_method"
 	fixMissingCalledFunction   = "stub_missing_called_function"
+	fixMissingStructField      = "stub_missing_struct_field"
 )
 
 // ApplyFix applies the specified kind of suggested fix to the given
@@ -113,6 +114,7 @@ func ApplyFix(ctx context.Context, fix string, snapshot *cache.Snapshot, fh file
 		fixCreateUndeclared:        singleFile(CreateUndeclared),
 		fixMissingInterfaceMethods: stubMissingInterfaceMethodsFixer,
 		fixMissingCalledFunction:   stubMissingCalledFunctionFixer,
+		fixMissingStructField:      stubMissingStructFieldFixer,
 	}
 	fixer, ok := fixers[fix]
 	if !ok {
diff --git a/gopls/internal/golang/stub.go b/gopls/internal/golang/stub.go
index 036c1f959eb..41e06b64dd9 100644
--- a/gopls/internal/golang/stub.go
+++ b/gopls/internal/golang/stub.go
@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"go/ast"
 	"go/format"
 	"go/parser"
 	"go/token"
@@ -51,6 +52,18 @@ func stubMissingCalledFunctionFixer(ctx context.Context, snapshot *cache.Snapsho
 	return insertDeclsAfter(ctx, snapshot, pkg.Metadata(), si.Fset, si.After, si.Emit)
 }
 
+// stubMissingStructFieldFixer returns a suggested fix to declare the missing
+// field that the user may want to generate based on SelectorExpr
+// at the cursor position.
+func stubMissingStructFieldFixer(ctx context.Context, snapshot *cache.Snapshot, pkg *cache.Package, pgf *parsego.File, start, end token.Pos) (*token.FileSet, *analysis.SuggestedFix, error) {
+	nodes, _ := astutil.PathEnclosingInterval(pgf.File, start, end)
+	fi := GetFieldStubInfo(pkg.FileSet(), pkg.TypesInfo(), nodes)
+	if fi == nil {
+		return nil, nil, fmt.Errorf("invalid type request")
+	}
+	return insertStructField(ctx, snapshot, pkg.Metadata(), fi)
+}
+
 // An emitter writes new top-level declarations into an existing
 // file. References to symbols should be qualified using qual, which
 // respects the local import environment.
@@ -238,3 +251,66 @@ func trimVersionSuffix(path string) string {
 	}
 	return path
 }
+
+func insertStructField(ctx context.Context, snapshot *cache.Snapshot, meta *metadata.Package, fieldInfo *StructFieldInfo) (*token.FileSet, *analysis.SuggestedFix, error) {
+	if fieldInfo == nil {
+		return nil, nil, fmt.Errorf("no field info provided")
+	}
+
+	// get the file containing the struct definition using the position
+	declPGF, _, err := parseFull(ctx, snapshot, fieldInfo.Fset, fieldInfo.Named.Obj().Pos())
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to parse file declaring struct: %w", err)
+	}
+	if declPGF.Fixed() {
+		return nil, nil, fmt.Errorf("file contains parse errors: %s", declPGF.URI)
+	}
+
+	// find the struct type declaration
+	var structType *ast.StructType
+	ast.Inspect(declPGF.File, func(n ast.Node) bool {
+		if typeSpec, ok := n.(*ast.TypeSpec); ok {
+			if typeSpec.Name.Name == fieldInfo.Named.Obj().Name() {
+				if st, ok := typeSpec.Type.(*ast.StructType); ok {
+					structType = st
+					return false
+				}
+			}
+		}
+		return true
+	})
+
+	if structType == nil {
+		return nil, nil, fmt.Errorf("could not find struct definition")
+	}
+
+	// find the position to insert the new field (end of struct fields)
+	insertPos := structType.Fields.Closing - 1
+	if insertPos == structType.Fields.Opening {
+		// struct has no fields yet
+		insertPos = structType.Fields.Closing
+	}
+
+	var buf bytes.Buffer
+	if err := fieldInfo.Emit(&buf, types.RelativeTo(fieldInfo.Named.Obj().Pkg())); err != nil {
+		return nil, nil, err
+	}
+
+	_, err = declPGF.Mapper.PosRange(declPGF.Tok, insertPos, insertPos)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	textEdit := analysis.TextEdit{
+		Pos:     insertPos,
+		End:     insertPos,
+		NewText: []byte(buf.String()),
+	}
+
+	fix := &analysis.SuggestedFix{
+		Message:   fmt.Sprintf("Add field %s to struct %s", fieldInfo.Expr.Sel.Name, fieldInfo.Named.Obj().Name()),
+		TextEdits: []analysis.TextEdit{textEdit},
+	}
+
+	return fieldInfo.Fset, fix, nil
+}
diff --git a/gopls/internal/test/marker/testdata/quickfix/struct_field.txt b/gopls/internal/test/marker/testdata/quickfix/struct_field.txt
new file mode 100644
index 00000000000..972eeed54a2
--- /dev/null
+++ b/gopls/internal/test/marker/testdata/quickfix/struct_field.txt
@@ -0,0 +1,132 @@
+This test checks the 'Declare missing struct field' quick fix.
+
+-- go.mod --
+module module
+
+-- package.go --
+package field
+
+import "module/another"
+import alias "module/second"
+
+type Struct struct{}
+
+type AnotherStruct struct {
+	Chan <-chan Struct
+}
+
+func do() {
+	a := AnotherStruct{}
+
+	s := Struct{}
+	s.intField = 42 //@quickfix("intField", re"has no field or method", value_field)
+
+	var mp map[string]string = map[string]string{"key": "value4"}
+	s.mapField = mp //@quickfix("mapField", re"has no field or method", var_field)
+
+	s.chanField = a.Chan //@quickfix("chanField", re"has no field or method", another_struct_field)
+	s.sliceField = make([]map[string]Struct, 0) //@quickfix("sliceField", re"has no field or method", make_field)
+	s.sliceIntField = []int{1, 2} //@quickfix("sliceIntField", re"has no field or method", slice_int_field)
+	s.another = another.Another{} //@quickfix("another", re"has no field or method", another_package)
+	s.alias = alias.Second{} //@quickfix("alias", re"has no field or method", alias)
+	var al alias.Second
+	s.implicitAlias = al //@quickfix("implicitAlias", re"has no field or method", implicit_alias)
+	s.imported = alias.Second{}.Imported //@quickfix("imported", re"has no field or method", auto_import)
+	s.newField = new(Struct) //@quickfix("newField", re"has no field or method", new_field)
+	s.pointerField = &Struct{} //@quickfix("pointerField", re"has no field or method", pointer)
+	var p *Struct
+	s.derefedField = *p //@quickfix("derefedField", re"has no field or method", deref)
+
+	a.properlyFormattedField = 42 //@quickfix("properlyFormattedField", re"has no field or method", formatted)
+}
+-- another/another.go --
+package another
+
+type Another struct {}
+-- second/second.go --
+package second
+
+import "module/imported"
+
+type Second struct{
+	Imported imported.Imported
+}
+-- imported/imported.go --
+package imported
+
+type Imported struct{}
+-- @value_field/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	intField int
++}
+-- @var_field/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	mapField map[string]string
++}
+-- @another_struct_field/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	chanField <-chan Struct
++}
+-- @slice_int_field/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	sliceIntField []int
++}
+-- @make_field/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	sliceField []map[string]Struct
++}
+-- @another_package/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	another another.Another
++}
+-- @alias/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	alias alias.Second
++}
+-- @implicit_alias/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	implicitAlias alias.Second
++}
+-- @auto_import/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	imported imported.Imported
++}
+-- @new_field/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	newField *Struct
++}
+-- @pointer/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	pointerField *Struct
++}
+-- @deref/package.go --
+@@ -6 +6,3 @@
+-type Struct struct{}
++type Struct struct{
++	derefedField Struct
++}
+-- @formatted/package.go --
+@@ -10 +10 @@
++	properlyFormattedField int
diff --git a/gopls/internal/util/typesutil/typesutil.go b/gopls/internal/util/typesutil/typesutil.go
index 98f5605200e..059778058d6 100644
--- a/gopls/internal/util/typesutil/typesutil.go
+++ b/gopls/internal/util/typesutil/typesutil.go
@@ -9,6 +9,7 @@ import (
 	"go/ast"
 	"go/token"
 	"go/types"
+	"strconv"
 	"strings"
 )
 
@@ -209,6 +210,120 @@ func TypesFromContext(info *types.Info, path []ast.Node, pos token.Pos) []types.
 			t := info.TypeOf(parent.X)
 			typs = append(typs, validType(t))
 		}
+	case *ast.SelectorExpr:
+		for _, n := range path {
+			assignExpr, ok := n.(*ast.AssignStmt)
+			if ok {
+				for _, rh := range assignExpr.Rhs {
+					// basic types
+					basicLit, ok := rh.(*ast.BasicLit)
+					if ok {
+						switch basicLit.Kind {
+						case token.INT:
+							typs = append(typs, types.Typ[types.Int])
+						case token.FLOAT:
+							typs = append(typs, types.Typ[types.Float64])
+						case token.IMAG:
+							typs = append(typs, types.Typ[types.Complex128])
+						case token.STRING:
+							typs = append(typs, types.Typ[types.String])
+						case token.CHAR:
+							typs = append(typs, types.Typ[types.Rune])
+						}
+						break
+					}
+					callExpr, ok := rh.(*ast.CallExpr)
+					if ok {
+						if ident, ok := callExpr.Fun.(*ast.Ident); ok && ident.Name == "make" && len(callExpr.Args) > 0 {
+							arg := callExpr.Args[0]
+							composite, ok := arg.(*ast.CompositeLit)
+							if ok {
+								t := typeFromExpr(info, path, composite)
+								typs = append(typs, t)
+								break
+							}
+							if t := info.TypeOf(arg); t != nil {
+								typs = append(typs, validType(t))
+							}
+						}
+						if ident, ok := callExpr.Fun.(*ast.Ident); ok && ident.Name == "new" && len(callExpr.Args) > 0 {
+							arg := callExpr.Args[0]
+							composite, ok := arg.(*ast.CompositeLit)
+							if ok {
+								t := typeFromExpr(info, path, composite)
+								t = types.NewPointer(t)
+								typs = append(typs, t)
+								break
+							}
+							if t := info.TypeOf(arg); t != nil {
+								if !containsInvalid(t) {
+									t = types.Default(t)
+									t = types.NewPointer(t)
+								} else {
+									t = anyType
+								}
+								typs = append(typs, t)
+							}
+						}
+						break
+					}
+					// a variable
+					ident, ok := rh.(*ast.Ident)
+					if ok {
+						if t := typeFromExpr(info, path, ident); t != nil {
+							typs = append(typs, t)
+						}
+						break
+					}
+
+					selectorExpr, ok := rh.(*ast.SelectorExpr)
+					if ok {
+						if t := typeFromExpr(info, path, selectorExpr.Sel); t != nil {
+							typs = append(typs, t)
+						}
+						break
+					}
+					// composite
+					composite, ok := rh.(*ast.CompositeLit)
+					if ok {
+						t := typeFromExpr(info, path, composite)
+						typs = append(typs, t)
+						break
+					}
+					// a pointer
+					un, ok := rh.(*ast.UnaryExpr)
+					if ok && un.Op == token.AND {
+						composite, ok := un.X.(*ast.CompositeLit)
+						if !ok {
+							break
+						}
+						if t := info.TypeOf(composite); t != nil {
+							if !containsInvalid(t) {
+								t = types.Default(t)
+								t = types.NewPointer(t)
+							} else {
+								t = anyType
+							}
+							typs = append(typs, t)
+						}
+					}
+					starExpr, ok := rh.(*ast.StarExpr)
+					if ok {
+						ident, ok := starExpr.X.(*ast.Ident)
+						if ok {
+							if t := typeFromExpr(info, path, ident); t != nil {
+								if pointer, ok := t.(*types.Pointer); ok {
+									t = pointer.Elem()
+								}
+								typs = append(typs, t)
+							}
+							break
+						}
+					}
+				}
+			}
+		}
+
 	default:
 		// TODO: support other kinds of "holes" as the need arises.
 	}
@@ -257,3 +372,55 @@ func EnclosingSignature(path []ast.Node, info *types.Info) *types.Signature {
 	}
 	return nil
 }
+
+func typeFromExpr(info *types.Info, path []ast.Node, expr ast.Expr) types.Type {
+	t := info.TypeOf(expr)
+	if t == nil {
+		return nil
+	}
+
+	if !containsInvalid(t) {
+		t = types.Default(t)
+		if named, ok := t.(*types.Named); ok {
+			if pkg := named.Obj().Pkg(); pkg != nil {
+				// find the file in the path that contains this assignment
+				var file *ast.File
+				for _, n := range path {
+					if f, ok := n.(*ast.File); ok {
+						file = f
+						break
+					}
+				}
+
+				if file != nil {
+					// look for any import spec that imports this package
+					var pkgName string
+					for _, imp := range file.Imports {
+						if path, _ := strconv.Unquote(imp.Path.Value); path == pkg.Path() {
+							// use the alias if specified, otherwise use package name
+							if imp.Name != nil {
+								pkgName = imp.Name.Name
+							} else {
+								pkgName = pkg.Name()
+							}
+							break
+						}
+					}
+					// fallback to package name if no import found
+					if pkgName == "" {
+						pkgName = pkg.Name()
+					}
+
+					// create new package with the correct name (either alias or original)
+					newPkg := types.NewPackage(pkgName, pkgName)
+					newName := types.NewTypeName(named.Obj().Pos(), newPkg, named.Obj().Name(), nil)
+					t = types.NewNamed(newName, named.Underlying(), nil)
+				}
+			}
+			return t
+		}
+	} else {
+		t = types.Universe.Lookup("any").Type()
+	}
+	return t
+}