gopls, internal/lsp: support fillstruct for partially-filled structs

This updates fillstruct to work even when the struct is partially
filled. User supplied fields are preserved but comments are blown away.

Preserving comments appears to be very hard with the current ast
library. One possible option is to do manual string shenanigans, but
after exploring that path it seems like A Bad Idea.

Fixes golang/go#39804

Change-Id: Iec0bb93db05d4d726dfa6c77a8139f53b14bcc77
Reviewed-on: https://go-review.googlesource.com/c/tools/+/262018
Run-TryBot: Jean de Klerk <deklerk@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
Trust: Jean de Klerk <deklerk@google.com>
diff --git a/internal/lsp/analysis/fillstruct/fillstruct.go b/internal/lsp/analysis/fillstruct/fillstruct.go
index c760b69..36a63a1 100644
--- a/internal/lsp/analysis/fillstruct/fillstruct.go
+++ b/internal/lsp/analysis/fillstruct/fillstruct.go
@@ -49,11 +49,6 @@
 		}
 		expr := n.(*ast.CompositeLit)
 
-		// TODO: Handle partially-filled structs as well.
-		if len(expr.Elts) != 0 {
-			return
-		}
-
 		var file *ast.File
 		for _, f := range pass.Files {
 			if f.Pos() <= expr.Pos() && expr.Pos() <= f.End() {
@@ -90,10 +85,10 @@
 		if fieldCount == 0 || fieldCount == len(expr.Elts) {
 			return
 		}
+
 		var fillable bool
 		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() != pass.Pkg && !field.Exported() {
 				continue
@@ -137,6 +132,7 @@
 			break
 		}
 	}
+
 	if info == nil {
 		return nil, fmt.Errorf("nil types.Info")
 	}
@@ -161,6 +157,17 @@
 	}
 	fieldCount := obj.NumFields()
 
+	// Check which types have already been filled in. (we only want to fill in
+	// the unfilled types, or else we'll blat user-supplied details)
+	prefilledTypes := map[string]ast.Expr{}
+	for _, e := range expr.Elts {
+		if kv, ok := e.(*ast.KeyValueExpr); ok {
+			if key, ok := kv.Key.(*ast.Ident); ok {
+				prefilledTypes[key.Name] = kv.Value
+			}
+		}
+	}
+
 	// Use a new fileset to build up a token.File for the new composite
 	// literal. We need one line for foo{, one line for }, and one line for
 	// each field we're going to set. format.Node only cares about line
@@ -186,21 +193,6 @@
 		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() {
@@ -214,7 +206,27 @@
 				Name:    obj.Field(i).Name(),
 			},
 			Colon: pos,
-			Value: value,
+		}
+		if expr, ok := prefilledTypes[obj.Field(i).Name()]; ok {
+			kv.Value = expr
+		} else {
+			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
+			}
+
+			kv.Value = value
 		}
 		elts = append(elts, kv)
 		line++
@@ -251,31 +263,51 @@
 	index := bytes.Index(firstLine, trimmed)
 	whitespace := firstLine[:index]
 
-	var newExpr bytes.Buffer
-	if err := format.Node(&newExpr, fakeFset, cl); err != nil {
-		return nil, fmt.Errorf("failed to format %s: %v", cl.Type, err)
+	// First pass through the formatter: turn the expr into a string.
+	var formatBuf bytes.Buffer
+	if err := format.Node(&formatBuf, fakeFset, cl); err != nil {
+		return nil, fmt.Errorf("failed to run first format on:\n%s\ngot err: %v", cl.Type, err)
 	}
-	split = bytes.Split(newExpr.Bytes(), []byte("\n"))
+	sug := indent(formatBuf.Bytes(), whitespace)
+
+	if len(prefilledTypes) > 0 {
+		// Attempt a second pass through the formatter to line up columns.
+		sourced, err := format.Source(sug)
+		if err == nil {
+			sug = indent(sourced, whitespace)
+		}
+	}
+
+	return &analysis.SuggestedFix{
+		TextEdits: []analysis.TextEdit{
+			{
+				Pos:     expr.Pos(),
+				End:     expr.End(),
+				NewText: sug,
+			},
+		},
+	}, nil
+}
+
+// indent works line by line through str, indenting (prefixing) each line with
+// ind.
+func indent(str, ind []byte) []byte {
+	split := bytes.Split(str, []byte("\n"))
 	newText := bytes.NewBuffer(nil)
 	for i, s := range split {
+		if len(s) == 0 {
+			continue
+		}
 		// Don't add the extra indentation to the first line.
 		if i != 0 {
-			newText.Write(whitespace)
+			newText.Write(ind)
 		}
 		newText.Write(s)
 		if i < len(split)-1 {
 			newText.WriteByte('\n')
 		}
 	}
-	return &analysis.SuggestedFix{
-		TextEdits: []analysis.TextEdit{
-			{
-				Pos:     expr.Pos(),
-				End:     expr.End(),
-				NewText: newText.Bytes(),
-			},
-		},
-	}, nil
+	return newText.Bytes()
 }
 
 // populateValue constructs an expression to fill the value of a struct field.
diff --git a/internal/lsp/analysis/fillstruct/testdata/src/a/a.go b/internal/lsp/analysis/fillstruct/testdata/src/a/a.go
index b5df1de..f69fe83 100644
--- a/internal/lsp/analysis/fillstruct/testdata/src/a/a.go
+++ b/internal/lsp/analysis/fillstruct/testdata/src/a/a.go
@@ -27,7 +27,7 @@
 
 var _ = twoArgStruct{} // want ""
 
-var _ = twoArgStruct{
+var _ = twoArgStruct{ // want ""
 	bar: "bar",
 }
 
diff --git a/internal/lsp/cmd/test/suggested_fix.go b/internal/lsp/cmd/test/suggested_fix.go
index 965edd4..bf7fb30 100644
--- a/internal/lsp/cmd/test/suggested_fix.go
+++ b/internal/lsp/cmd/test/suggested_fix.go
@@ -12,7 +12,7 @@
 	"golang.org/x/tools/internal/span"
 )
 
-func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string) {
+func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string, expectedActions int) {
 	uri := spn.URI()
 	filename := uri.Filename()
 	args := []string{"fix", "-a", fmt.Sprintf("%s", spn)}
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index f9a3cb2..81d3168 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -458,7 +458,7 @@
 	}
 }
 
-func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string) {
+func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string, expectedActions int) {
 	uri := spn.URI()
 	view, err := r.server.session.ViewOf(uri)
 	if err != nil {
@@ -509,10 +509,13 @@
 	if err != nil {
 		t.Fatalf("CodeAction %s failed: %v", spn, err)
 	}
-	if len(actions) != 1 {
+	if len(actions) != expectedActions {
 		// Hack: We assume that we only get one code action per range.
-		// TODO(rstambler): Support multiple code actions per test.
-		t.Fatalf("unexpected number of code actions, want 1, got %v", len(actions))
+		var cmds []string
+		for _, a := range actions {
+			cmds = append(cmds, fmt.Sprintf("%s (%s)", a.Command.Command, a.Title))
+		}
+		t.Fatalf("unexpected number of code actions, want %d, got %d: %v", expectedActions, len(actions), cmds)
 	}
 	action := actions[0]
 	var match bool
@@ -529,7 +532,7 @@
 	if cmd := action.Command; cmd != nil {
 		edits, err := commandToEdits(r.ctx, snapshot, fh, rng, action.Command.Command)
 		if err != nil {
-			t.Fatal(err)
+			t.Fatalf("error converting command %q to edits: %v", action.Command.Command, err)
 		}
 		res, err = applyTextDocumentEdits(r, edits)
 		if err != nil {
@@ -565,7 +568,11 @@
 	if !command.Applies(ctx, snapshot, fh, rng) {
 		return nil, fmt.Errorf("cannot apply %v", command.ID())
 	}
-	return command.SuggestedFix(ctx, snapshot, fh, rng)
+	edits, err := command.SuggestedFix(ctx, snapshot, fh, rng)
+	if err != nil {
+		return nil, fmt.Errorf("error calling command.SuggestedFix: %v", err)
+	}
+	return edits, nil
 }
 
 func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span) {
diff --git a/internal/lsp/source/source_test.go b/internal/lsp/source/source_test.go
index 2d79960..169e3c3 100644
--- a/internal/lsp/source/source_test.go
+++ b/internal/lsp/source/source_test.go
@@ -923,8 +923,9 @@
 }
 
 // These are pure LSP features, no source level functionality to be tested.
-func (r *runner) Link(t *testing.T, uri span.URI, wantLinks []tests.Link)         {}
-func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string)  {}
+func (r *runner) Link(t *testing.T, uri span.URI, wantLinks []tests.Link) {}
+func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string, expectedActions int) {
+}
 func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span) {}
 func (r *runner) CodeLens(t *testing.T, uri span.URI, want []protocol.CodeLens)   {}
 
diff --git a/internal/lsp/testdata/fillstruct/fill_struct_partial.go b/internal/lsp/testdata/fillstruct/fill_struct_partial.go
new file mode 100644
index 0000000..97b517d
--- /dev/null
+++ b/internal/lsp/testdata/fillstruct/fill_struct_partial.go
@@ -0,0 +1,24 @@
+package fillstruct
+
+type StructPartialA struct {
+	PrefilledInt int
+	UnfilledInt  int
+	StructPartialB
+}
+
+type StructPartialB struct {
+	PrefilledInt int
+	UnfilledInt  int
+}
+
+func fill() {
+	a := StructPartialA{
+		PrefilledInt: 5,
+	} //@suggestedfix("}", "refactor.rewrite")
+	b := StructPartialB{
+		/* this comment should disappear */
+		PrefilledInt: 7, // This comment should be blown away.
+		/* As should
+		this one */
+	} //@suggestedfix("}", "refactor.rewrite")
+}
diff --git a/internal/lsp/testdata/fillstruct/fill_struct_partial.go.golden b/internal/lsp/testdata/fillstruct/fill_struct_partial.go.golden
new file mode 100644
index 0000000..2d063c1
--- /dev/null
+++ b/internal/lsp/testdata/fillstruct/fill_struct_partial.go.golden
@@ -0,0 +1,52 @@
+-- suggestedfix_fill_struct_partial_17_2 --
+package fillstruct
+
+type StructPartialA struct {
+	PrefilledInt int
+	UnfilledInt  int
+	StructPartialB
+}
+
+type StructPartialB struct {
+	PrefilledInt int
+	UnfilledInt  int
+}
+
+func fill() {
+	a := StructPartialA{
+		PrefilledInt:   5,
+		UnfilledInt:    0,
+		StructPartialB: StructPartialB{},
+	} //@suggestedfix("}", "refactor.rewrite")
+	b := StructPartialB{
+		/* this comment should disappear */
+		PrefilledInt: 7, // This comment should be blown away.
+		/* As should
+		this one */
+	} //@suggestedfix("}", "refactor.rewrite")
+}
+
+-- suggestedfix_fill_struct_partial_23_2 --
+package fillstruct
+
+type StructPartialA struct {
+	PrefilledInt int
+	UnfilledInt  int
+	StructPartialB
+}
+
+type StructPartialB struct {
+	PrefilledInt int
+	UnfilledInt  int
+}
+
+func fill() {
+	a := StructPartialA{
+		PrefilledInt: 5,
+	} //@suggestedfix("}", "refactor.rewrite")
+	b := StructPartialB{
+		PrefilledInt: 7,
+		UnfilledInt:  0,
+	} //@suggestedfix("}", "refactor.rewrite")
+}
+
diff --git a/internal/lsp/testdata/summary.txt.golden b/internal/lsp/testdata/summary.txt.golden
index a6aa546..12324fa 100644
--- a/internal/lsp/testdata/summary.txt.golden
+++ b/internal/lsp/testdata/summary.txt.golden
@@ -13,7 +13,7 @@
 FormatCount = 6
 ImportCount = 8
 SemanticTokenCount = 3
-SuggestedFixCount = 38
+SuggestedFixCount = 40
 FunctionExtractionCount = 12
 DefinitionsCount = 64
 TypeDefinitionsCount = 2
diff --git a/internal/lsp/tests/tests.go b/internal/lsp/tests/tests.go
index 33c8625..7eaf9ac 100644
--- a/internal/lsp/tests/tests.go
+++ b/internal/lsp/tests/tests.go
@@ -134,7 +134,7 @@
 	Format(*testing.T, span.Span)
 	Import(*testing.T, span.Span)
 	SemanticTokens(*testing.T, span.Span)
-	SuggestedFix(*testing.T, span.Span, []string)
+	SuggestedFix(*testing.T, span.Span, []string, int)
 	FunctionExtraction(*testing.T, span.Span, span.Span)
 	Definition(*testing.T, span.Span, Definition)
 	Implementation(*testing.T, span.Span, []span.Span)
@@ -633,7 +633,7 @@
 			}
 			t.Run(SpanName(spn), func(t *testing.T) {
 				t.Helper()
-				tests.SuggestedFix(t, spn, actionKinds)
+				tests.SuggestedFix(t, spn, actionKinds, 1)
 			})
 		}
 	})