internal/lsp/analysis/fillstruct: add indentation for struct fields

Getting the right indentation for the text in the AST proves to be a
little complicated. The most reasonable approach seems to be writing
out the AST, getting the lines with the struct definition on it,
and trimming whitespace to get the current indent. Then, we add this
indent to the struct fields in the new text.

Change-Id: I1cc3421d95edae61cfb662254ff3fb759b5c487f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/239199
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
diff --git a/internal/lsp/analysis/fillstruct/fillstruct.go b/internal/lsp/analysis/fillstruct/fillstruct.go
index 8fb4585..bbe7343 100644
--- a/internal/lsp/analysis/fillstruct/fillstruct.go
+++ b/internal/lsp/analysis/fillstruct/fillstruct.go
@@ -11,8 +11,12 @@
 	"fmt"
 	"go/ast"
 	"go/format"
+	"go/printer"
 	"go/token"
 	"go/types"
+	"log"
+	"strings"
+	"unicode"
 
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/inspect"
@@ -96,16 +100,29 @@
 			return
 		}
 
-		// Don't mutate the existing token.File. Instead, create a copy that we can use to modify
-		// position information.
-		original := pass.Fset.File(expr.Lbrace)
-		fset := token.NewFileSet()
-		tok := fset.AddFile(original.Name(), -1, original.Size())
+		var name string
+		switch typ := expr.Type.(type) {
+		case *ast.Ident:
+			name = typ.Name
+		case *ast.SelectorExpr:
+			name = fmt.Sprintf("%s.%s", typ.X, typ.Sel.Name)
+		default:
+			log.Printf("anonymous structs are not yet supported: %v (%T)", expr.Type, expr.Type)
+			return
+		}
 
-		pos := token.Pos(1)
+		// 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
+		// numbers, so we don't need to set columns, and each line can be
+		// 1 byte long.
+		fset := token.NewFileSet()
+		tok := fset.AddFile("", -1, fieldCount+2)
+
 		var elts []ast.Expr
 		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
@@ -115,14 +132,18 @@
 			if value == nil {
 				continue
 			}
-			pos = nextLinePos(tok, pos)
+
+			line := i + 2      // account for 1-based lines and the left brace
+			tok.AddLine(i + 1) // add 1 byte per line
+			pos := tok.LineStart(line)
+
 			kv := &ast.KeyValueExpr{
 				Key: &ast.Ident{
 					NamePos: pos,
 					Name:    field.Name(),
 				},
 				Colon: pos,
-				Value: value, // 'value' has no position. fomat.Node corrects for AST nodes with no position.
+				Value: value,
 			}
 			elts = append(elts, kv)
 		}
@@ -132,32 +153,61 @@
 			return
 		}
 
-		cl := ast.CompositeLit{
-			Type:   expr.Type, // Don't adjust the expr.Type's position.
-			Lbrace: token.Pos(1),
+		// Add the final line for the right brace. Offset is the number of
+		// bytes already added plus 1.
+		tok.AddLine(len(elts) + 1)
+
+		cl := &ast.CompositeLit{
+			Type:   expr.Type,
+			Lbrace: tok.LineStart(1),
 			Elts:   elts,
-			Rbrace: nextLinePos(tok, elts[len(elts)-1].Pos()),
+			Rbrace: tok.LineStart(len(elts) + 2),
 		}
 
-		var buf bytes.Buffer
-		if err := format.Node(&buf, fset, &cl); err != nil {
+		// Print the AST to get the the original source code.
+		var b bytes.Buffer
+		if err := printer.Fprint(&b, pass.Fset, file); err != nil {
+			log.Printf("failed to print original file: %s", err)
 			return
 		}
 
-		msg := "Fill struct"
-		if name, ok := expr.Type.(*ast.Ident); ok {
-			msg = fmt.Sprintf("Fill %s", name)
-		}
+		// Find the line on which the composite literal is declared.
+		split := strings.Split(b.String(), "\n")
+		lineNumber := pass.Fset.Position(expr.Type.Pos()).Line
+		line := split[lineNumber-1] // lines are 1-indexed
 
+		// Trim the whitespace from the left of the line, and use the index
+		// to get the amount of whitespace on the left.
+		trimmed := strings.TrimLeftFunc(line, unicode.IsSpace)
+		i := strings.Index(line, trimmed)
+		whitespace := line[:i]
+
+		var newExpr bytes.Buffer
+		if err := format.Node(&newExpr, fset, cl); err != nil {
+			log.Printf("failed to format %s: %v", cl.Type, err)
+			return
+		}
+		split = strings.Split(newExpr.String(), "\n")
+		var newText strings.Builder
+		for i, s := range split {
+			// Don't add the extra indentation to the first line.
+			if i != 0 {
+				newText.WriteString(whitespace)
+			}
+			newText.WriteString(s)
+			if i < len(split)-1 {
+				newText.WriteByte('\n')
+			}
+		}
 		pass.Report(analysis.Diagnostic{
 			Pos: expr.Lbrace,
 			End: expr.Rbrace,
 			SuggestedFixes: []analysis.SuggestedFix{{
-				Message: msg,
+				Message: fmt.Sprintf("Fill %s with default values", name),
 				TextEdits: []analysis.TextEdit{{
 					Pos:     expr.Pos(),
 					End:     expr.End(),
-					NewText: buf.Bytes(),
+					NewText: []byte(newText.String()),
 				}},
 			}},
 		})
@@ -165,14 +215,6 @@
 	return nil, nil
 }
 
-func nextLinePos(tok *token.File, pos token.Pos) token.Pos {
-	line := tok.Line(pos)
-	if line+1 > tok.LineCount() {
-		tok.AddLine(tok.Offset(pos) + 1)
-	}
-	return tok.LineStart(line + 1)
-}
-
 // populateValue constructs an expression to fill the value of a struct field.
 //
 // When the type of a struct field is a basic literal or interface, we return
diff --git a/internal/lsp/analysis/undeclaredname/undeclared.go b/internal/lsp/analysis/undeclaredname/undeclared.go
index cf7c770..84a96d2 100644
--- a/internal/lsp/analysis/undeclaredname/undeclared.go
+++ b/internal/lsp/analysis/undeclaredname/undeclared.go
@@ -10,7 +10,7 @@
 	"bytes"
 	"fmt"
 	"go/ast"
-	"go/format"
+	"go/printer"
 	"strings"
 
 	"golang.org/x/tools/go/analysis"
@@ -89,7 +89,7 @@
 		}
 
 		var buf bytes.Buffer
-		if err := format.Node(&buf, pass.Fset, file); err != nil {
+		if err := printer.Fprint(&buf, pass.Fset, file); err != nil {
 			continue
 		}
 		old := buf.Bytes()
diff --git a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go
index 7fd1597..fccec13 100644
--- a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go
+++ b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go
@@ -20,4 +20,7 @@
 	a := StructA{}  //@suggestedfix("}", "refactor.rewrite")
 	b := StructA2{} //@suggestedfix("}", "refactor.rewrite")
 	c := StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	if true {
+		_ = StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	}
 }
diff --git a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go.golden b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go.golden
index 7fb2752..8d99703 100644
--- a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go.golden
+++ b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct.go.golden
@@ -19,14 +19,17 @@
 
 func fill() {
 	a := StructA{
-	unexportedIntField: 0,
-	ExportedIntField:   0,
-	MapA:               map[int]string{},
-	Array:              []int{},
-	StructB:            StructB{},
-}  //@suggestedfix("}", "refactor.rewrite")
+		unexportedIntField: 0,
+		ExportedIntField:   0,
+		MapA:               map[int]string{},
+		Array:              []int{},
+		StructB:            StructB{},
+	}  //@suggestedfix("}", "refactor.rewrite")
 	b := StructA2{} //@suggestedfix("}", "refactor.rewrite")
 	c := StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	if true {
+		_ = StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	}
 }
 
 -- suggestedfix_fill_struct_21_16 --
@@ -51,9 +54,12 @@
 func fill() {
 	a := StructA{}  //@suggestedfix("}", "refactor.rewrite")
 	b := StructA2{
-	B: &StructB{},
-} //@suggestedfix("}", "refactor.rewrite")
+		B: &StructB{},
+	} //@suggestedfix("}", "refactor.rewrite")
 	c := StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	if true {
+		_ = StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	}
 }
 
 -- suggestedfix_fill_struct_22_16 --
@@ -79,7 +85,40 @@
 	a := StructA{}  //@suggestedfix("}", "refactor.rewrite")
 	b := StructA2{} //@suggestedfix("}", "refactor.rewrite")
 	c := StructA3{
-	B: StructB{},
-} //@suggestedfix("}", "refactor.rewrite")
+		B: StructB{},
+	} //@suggestedfix("}", "refactor.rewrite")
+	if true {
+		_ = StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	}
+}
+
+-- suggestedfix_fill_struct_24_16 --
+package fillstruct
+
+type StructA struct {
+	unexportedIntField int
+	ExportedIntField   int
+	MapA               map[int]string
+	Array              []int
+	StructB
+}
+
+type StructA2 struct {
+	B *StructB
+}
+
+type StructA3 struct {
+	B StructB
+}
+
+func fill() {
+	a := StructA{}  //@suggestedfix("}", "refactor.rewrite")
+	b := StructA2{} //@suggestedfix("}", "refactor.rewrite")
+	c := StructA3{} //@suggestedfix("}", "refactor.rewrite")
+	if true {
+		_ = StructA3{
+			B: StructB{},
+		} //@suggestedfix("}", "refactor.rewrite")
+	}
 }
 
diff --git a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_nested.go.golden b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_nested.go.golden
index 688ba77..30061a5 100644
--- a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_nested.go.golden
+++ b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_nested.go.golden
@@ -12,8 +12,8 @@
 func nested() {
 	c := StructB{
 		StructC: StructC{
-	unexportedInt: 0,
-}, //@suggestedfix("}", "refactor.rewrite")
+			unexportedInt: 0,
+		}, //@suggestedfix("}", "refactor.rewrite")
 	}
 }
 
diff --git a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_package.go.golden b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_package.go.golden
index d15f5a1..d4cefe5 100644
--- a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_package.go.golden
+++ b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_package.go.golden
@@ -7,7 +7,7 @@
 
 func unexported() {
 	a := data.A{
-	ExportedInt: 0,
-} //@suggestedfix("}", "refactor.rewrite")
+		ExportedInt: 0,
+	} //@suggestedfix("}", "refactor.rewrite")
 }
 
diff --git a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_spaces.go.golden b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_spaces.go.golden
index b0e8937..0d75533 100644
--- a/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_spaces.go.golden
+++ b/internal/lsp/testdata/lsp/primarymod/fillstruct/fill_struct_spaces.go.golden
@@ -7,7 +7,7 @@
 
 func spaces() {
 	d := StructD{
-	ExportedIntField: 0,
-} //@suggestedfix("}", "refactor.rewrite")
+		ExportedIntField: 0,
+	} //@suggestedfix("}", "refactor.rewrite")
 }
 
diff --git a/internal/lsp/testdata/lsp/summary.txt.golden b/internal/lsp/testdata/lsp/summary.txt.golden
index b4889cb..9c3986d4 100644
--- a/internal/lsp/testdata/lsp/summary.txt.golden
+++ b/internal/lsp/testdata/lsp/summary.txt.golden
@@ -11,7 +11,7 @@
 FoldingRangesCount = 2
 FormatCount = 6
 ImportCount = 8
-SuggestedFixCount = 12
+SuggestedFixCount = 13
 DefinitionsCount = 53
 TypeDefinitionsCount = 2
 HighlightsCount = 69