internal/lsp: fix composite literal completion

Fix the following issues:

- We were trying to complete struct literal field names for
  selector expressions (e.g. "Foo{a.B<>}"). Now we only complete field
  names in this case if the expression is an *ast.Ident.
- We weren't including lexical completions in cases where you might be
  completing a field name or a variable name (e.g. "Foo{A<>}").

I refactored composite literal logic to live mostly in one place. Now
enclosingCompositeLiteral computes all the bits of information related
to composite literals. The expected type, completion, and snippet code
make use of those precalculated facts instead of redoing the work.

Change-Id: I29fc808544382c3c77f0bba1843520e04f38e79b
GitHub-Last-Rev: 3489062be342ab0f00325d3b3ae9ce681df7cf2e
GitHub-Pull-Request: golang/tools#96
Reviewed-on: https://go-review.googlesource.com/c/tools/+/176601
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
diff --git a/internal/lsp/source/completion.go b/internal/lsp/source/completion.go
index 37a1ca2..ac37738 100644
--- a/internal/lsp/source/completion.go
+++ b/internal/lsp/source/completion.go
@@ -126,14 +126,30 @@
 	// not a value.
 	preferTypeNames bool
 
-	// enclosingCompositeLiteral is the composite literal enclosing the position.
-	enclosingCompositeLiteral *ast.CompositeLit
+	// enclosingCompositeLiteral contains information about the composite literal
+	// enclosing the position.
+	enclosingCompositeLiteral *compLitInfo
+}
 
-	// enclosingKeyValue is the key value expression enclosing the position.
-	enclosingKeyValue *ast.KeyValueExpr
+type compLitInfo struct {
+	// cl is the *ast.CompositeLit enclosing the position.
+	cl *ast.CompositeLit
 
-	// inCompositeLiteralField is true if we are completing a composite literal field.
-	inCompositeLiteralField bool
+	// clType is the type of cl.
+	clType types.Type
+
+	// kv is the *ast.KeyValueExpr enclosing the position, if any.
+	kv *ast.KeyValueExpr
+
+	// inKey is true if we are certain the position is in the key side
+	// of a key-value pair.
+	inKey bool
+
+	// maybeInFieldName is true if inKey is false and it is possible
+	// we are completing a struct field name. For example,
+	// "SomeStruct{<>}" will be inKey=false, but maybeInFieldName=true
+	// because we _could_ be completing a field name.
+	maybeInFieldName bool
 }
 
 // found adds a candidate completion.
@@ -186,7 +202,7 @@
 		return nil, "", nil
 	}
 
-	lit, kv, inCompositeLiteralField := enclosingCompositeLiteral(path, pos)
+	clInfo := enclosingCompositeLiteral(path, pos, pkg.GetTypesInfo())
 	c := &completer{
 		types:                     pkg.GetTypes(),
 		info:                      pkg.GetTypesInfo(),
@@ -198,30 +214,26 @@
 		seen:                      make(map[types.Object]bool),
 		enclosingFunction:         enclosingFunction(path, pos, pkg.GetTypesInfo()),
 		preferTypeNames:           preferTypeNames(path, pos),
-		enclosingCompositeLiteral: lit,
-		enclosingKeyValue:         kv,
-		inCompositeLiteralField:   inCompositeLiteralField,
+		enclosingCompositeLiteral: clInfo,
+	}
+
+	if ident, ok := path[0].(*ast.Ident); ok {
+		// Set the filter prefix.
+		c.prefix = ident.Name[:pos-ident.Pos()]
 	}
 
 	c.expectedType = expectedType(c)
 
-	// Composite literals are handled entirely separately.
-	if c.enclosingCompositeLiteral != nil {
-		c.expectedType = c.expectedCompositeLiteralType(c.enclosingCompositeLiteral, c.enclosingKeyValue)
-
-		if c.inCompositeLiteralField {
-			if err := c.compositeLiteral(c.enclosingCompositeLiteral, c.enclosingKeyValue); err != nil {
-				return nil, "", err
-			}
-			return c.items, c.prefix, nil
+	// Struct literals are handled entirely separately.
+	if c.wantStructFieldCompletions() {
+		if err := c.structLiteralFieldName(); err != nil {
+			return nil, "", err
 		}
+		return c.items, c.prefix, nil
 	}
 
 	switch n := path[0].(type) {
 	case *ast.Ident:
-		// Set the filter prefix.
-		c.prefix = n.Name[:pos-n.Pos()]
-
 		// Is this the Sel part of a selector?
 		if sel, ok := path[1].(*ast.SelectorExpr); ok && sel.Sel == n {
 			if err := c.selector(sel); err != nil {
@@ -265,9 +277,19 @@
 			return nil, "", err
 		}
 	}
+
 	return c.items, c.prefix, nil
 }
 
+func (c *completer) wantStructFieldCompletions() bool {
+	clInfo := c.enclosingCompositeLiteral
+	if clInfo == nil {
+		return false
+	}
+
+	return clInfo.isStruct() && (clInfo.inKey || clInfo.maybeInFieldName)
+}
+
 // selector finds completions for the specified selector expression.
 func (c *completer) selector(sel *ast.SelectorExpr) error {
 	// Is sel a qualified identifier?
@@ -370,23 +392,19 @@
 	return nil
 }
 
-// compositeLiteral finds completions for field names inside a composite literal.
-func (c *completer) compositeLiteral(lit *ast.CompositeLit, kv *ast.KeyValueExpr) error {
-	switch n := c.path[0].(type) {
-	case *ast.Ident:
-		c.prefix = n.Name[:c.pos-n.Pos()]
-	}
+// structLiteralFieldName finds completions for struct field names inside a struct literal.
+func (c *completer) structLiteralFieldName() error {
+	clInfo := c.enclosingCompositeLiteral
+
 	// Mark fields of the composite literal that have already been set,
 	// except for the current field.
-	hasKeys := kv != nil // true if the composite literal already has key-value pairs
 	addedFields := make(map[*types.Var]bool)
-	for _, el := range lit.Elts {
+	for _, el := range clInfo.cl.Elts {
 		if kvExpr, ok := el.(*ast.KeyValueExpr); ok {
-			if kv == kvExpr {
+			if clInfo.kv == kvExpr {
 				continue
 			}
 
-			hasKeys = true
 			if key, ok := kvExpr.Key.(*ast.Ident); ok {
 				if used, ok := c.info.Uses[key]; ok {
 					if usedVar, ok := used.(*types.Var); ok {
@@ -396,34 +414,36 @@
 			}
 		}
 	}
-	// If the underlying type of the composite literal is a struct,
-	// collect completions for the fields of this struct.
-	if tv, ok := c.info.Types[lit]; ok {
-		switch t := tv.Type.Underlying().(type) {
-		case *types.Struct:
-			var structPkg *types.Package // package that struct is declared in
-			for i := 0; i < t.NumFields(); i++ {
-				field := t.Field(i)
-				if i == 0 {
-					structPkg = field.Pkg()
-				}
-				if !addedFields[field] {
-					c.found(field, highScore)
-				}
+
+	switch t := clInfo.clType.(type) {
+	case *types.Struct:
+		for i := 0; i < t.NumFields(); i++ {
+			field := t.Field(i)
+			if !addedFields[field] {
+				c.found(field, highScore)
 			}
-			// Add lexical completions if the user hasn't typed a key value expression
-			// and if the struct fields are defined in the same package as the user is in.
-			if !hasKeys && structPkg == c.types {
-				return c.lexical()
-			}
-		default:
+		}
+
+		// Add lexical completions if we aren't certain we are in the key part of a
+		// key-value pair.
+		if clInfo.maybeInFieldName {
 			return c.lexical()
 		}
+	default:
+		return c.lexical()
 	}
+
 	return nil
 }
 
-func enclosingCompositeLiteral(path []ast.Node, pos token.Pos) (lit *ast.CompositeLit, kv *ast.KeyValueExpr, ok bool) {
+func (cl *compLitInfo) isStruct() bool {
+	_, ok := cl.clType.(*types.Struct)
+	return ok
+}
+
+// enclosingCompositeLiteral returns information about the composite literal enclosing the
+// position.
+func enclosingCompositeLiteral(path []ast.Node, pos token.Pos, info *types.Info) *compLitInfo {
 	for _, n := range path {
 		switch n := n.(type) {
 		case *ast.CompositeLit:
@@ -433,39 +453,80 @@
 			//
 			// The position is not part of the composite literal unless it falls within the
 			// curly braces (e.g. "foo.Foo<>Struct{}").
-			if n.Lbrace <= pos && pos <= n.Rbrace {
-				lit = n
+			if !(n.Lbrace <= pos && pos <= n.Rbrace) {
+				return nil
+			}
 
-				// If the cursor position is within a key-value expression inside the composite
-				// literal, we try to determine if it is before or after the colon. If it is before
-				// the colon, we return field completions. If the cursor does not belong to any
-				// expression within the composite literal, we show composite literal completions.
-				if expr, isKeyValue := exprAtPos(pos, n.Elts).(*ast.KeyValueExpr); kv == nil && isKeyValue {
-					kv = expr
+			tv, ok := info.Types[n]
+			if !ok {
+				return nil
+			}
 
-					// If the position belongs to a key-value expression and is after the colon,
-					// don't show composite literal completions.
-					ok = pos <= kv.Colon
-				} else if kv == nil {
-					ok = true
+			clInfo := compLitInfo{
+				cl:     n,
+				clType: tv.Type.Underlying(),
+			}
+
+			var (
+				expr    ast.Expr
+				hasKeys bool
+			)
+			for _, el := range n.Elts {
+				// Remember the expression that the position falls in, if any.
+				if el.Pos() <= pos && pos <= el.End() {
+					expr = el
+				}
+
+				if kv, ok := el.(*ast.KeyValueExpr); ok {
+					hasKeys = true
+					// If expr == el then we know the position falls in this expression,
+					// so also record kv as the enclosing *ast.KeyValueExpr.
+					if expr == el {
+						clInfo.kv = kv
+						break
+					}
 				}
 			}
-			return lit, kv, ok
-		case *ast.KeyValueExpr:
-			if kv == nil {
-				kv = n
 
-				// If the position belongs to a key-value expression and is after the colon,
-				// don't show composite literal completions.
-				ok = pos <= kv.Colon
+			if clInfo.kv != nil {
+				// If in a *ast.KeyValueExpr, we know we are in the key if the position
+				// is to the left of the colon (e.g. "Foo{F<>: V}".
+				clInfo.inKey = pos <= clInfo.kv.Colon
+			} else if hasKeys {
+				// If we aren't in a *ast.KeyValueExpr but the composite literal has
+				// other *ast.KeyValueExprs, we must be on the key side of a new
+				// *ast.KeyValueExpr (e.g. "Foo{F: V, <>}").
+				clInfo.inKey = true
+			} else {
+				switch clInfo.clType.(type) {
+				case *types.Struct:
+					if len(n.Elts) == 0 {
+						// If the struct literal is empty, next could be a struct field
+						// name or an expression (e.g. "Foo{<>}" could become "Foo{F:}"
+						// or "Foo{someVar}").
+						clInfo.maybeInFieldName = true
+					} else if len(n.Elts) == 1 {
+						// If there is one expression and the position is in that expression
+						// and the expression is an identifier, we may be writing a field
+						// name or an expression (e.g. "Foo{F<>}").
+						_, clInfo.maybeInFieldName = expr.(*ast.Ident)
+					}
+				case *types.Map:
+					// If we aren't in a *ast.KeyValueExpr we must be adding a new key
+					// to the map.
+					clInfo.inKey = true
+				}
 			}
+
+			return &clInfo
 		default:
 			if breaksExpectedTypeInference(n) {
-				return nil, nil, false
+				return nil
 			}
 		}
 	}
-	return lit, kv, ok
+
+	return nil
 }
 
 // enclosingFunction returns the signature of the function enclosing the given position.
@@ -485,57 +546,60 @@
 	return nil
 }
 
-func (c *completer) expectedCompositeLiteralType(lit *ast.CompositeLit, kv *ast.KeyValueExpr) types.Type {
-	litType, ok := c.info.Types[lit]
-	if !ok {
-		return nil
-	}
-	switch t := litType.Type.Underlying().(type) {
+func (c *completer) expectedCompositeLiteralType() types.Type {
+	clInfo := c.enclosingCompositeLiteral
+	switch t := clInfo.clType.(type) {
 	case *types.Slice:
+		if clInfo.inKey {
+			return types.Typ[types.Int]
+		}
 		return t.Elem()
 	case *types.Array:
+		if clInfo.inKey {
+			return types.Typ[types.Int]
+		}
 		return t.Elem()
 	case *types.Map:
-		if kv == nil || c.pos <= kv.Colon {
+		if clInfo.inKey {
 			return t.Key()
 		}
 		return t.Elem()
 	case *types.Struct:
-		//  If we are in a key-value expression.
-		if kv != nil {
-			// There is no expected type for a struct field name.
-			if c.pos <= kv.Colon {
-				return nil
-			}
-			// Find the type of the struct field whose name matches the key.
-			if key, ok := kv.Key.(*ast.Ident); ok {
+		// If we are completing a key (i.e. field name), there is no expected type.
+		if clInfo.inKey {
+			return nil
+		}
+
+		// If we are in a key-value pair, but not in the key, then we must be on the
+		// value side. The expected type of the value will be determined from the key.
+		if clInfo.kv != nil {
+			if key, ok := clInfo.kv.Key.(*ast.Ident); ok {
 				for i := 0; i < t.NumFields(); i++ {
 					if field := t.Field(i); field.Name() == key.Name {
 						return field.Type()
 					}
 				}
 			}
-			return nil
-		}
-		// We are in a struct literal, but not a specific key-value pair.
-		// If the struct literal doesn't have explicit field names,
-		// we may still be able to suggest an expected type.
-		for _, el := range lit.Elts {
-			if _, ok := el.(*ast.KeyValueExpr); ok {
-				return nil
+		} else {
+			// If we aren't in a key-value pair and aren't in the key, we must be using
+			// implicit field names.
+
+			// The order of the literal fields must match the order in the struct definition.
+			// Find the element that the position belongs to and suggest that field's type.
+			if i := indexExprAtPos(c.pos, clInfo.cl.Elts); i < t.NumFields() {
+				return t.Field(i).Type()
 			}
 		}
-		// The order of the literal fields must match the order in the struct definition.
-		// Find the element that the position belongs to and suggest that field's type.
-		if i := indexExprAtPos(c.pos, lit.Elts); i < t.NumFields() {
-			return t.Field(i).Type()
-		}
 	}
 	return nil
 }
 
 // expectedType returns the expected type for an expression at the query position.
 func expectedType(c *completer) types.Type {
+	if c.enclosingCompositeLiteral != nil {
+		return c.expectedCompositeLiteralType()
+	}
+
 	var (
 		derefCount int // count of deref "*" operators
 		refCount   int // count of reference "&" operators
diff --git a/internal/lsp/source/completion_snippet.go b/internal/lsp/source/completion_snippet.go
index 2ebdbd8..dfbecd2 100644
--- a/internal/lsp/source/completion_snippet.go
+++ b/internal/lsp/source/completion_snippet.go
@@ -13,28 +13,22 @@
 
 // structFieldSnippets calculates the plain and placeholder snippets for struct literal field names.
 func (c *completer) structFieldSnippets(label, detail string) (*snippet.Builder, *snippet.Builder) {
-	if !c.inCompositeLiteralField {
+	clInfo := c.enclosingCompositeLiteral
+
+	if clInfo == nil || !clInfo.isStruct() {
 		return nil, nil
 	}
 
-	lit := c.enclosingCompositeLiteral
-	kv := c.enclosingKeyValue
-
-	// If we aren't in a composite literal or are already in a key-value expression,
-	// we don't want a snippet.
-	if lit == nil || kv != nil {
+	// If we are already in a key-value expression, we don't want a snippet.
+	if clInfo.kv != nil {
 		return nil, nil
 	}
-	// First, confirm that we are actually completing a struct field.
-	if len(lit.Elts) > 0 {
-		i := indexExprAtPos(c.pos, lit.Elts)
-		if i >= len(lit.Elts) {
-			return nil, nil
-		}
-		// If the expression is not an identifier, it is not a struct field name.
-		if _, ok := lit.Elts[i].(*ast.Ident); !ok {
-			return nil, nil
-		}
+
+	// We don't want snippet unless we are completing a field name. maybeInFieldName
+	// means we _might_ not be a struct field name, but this method is only called for
+	// struct fields, so we can ignore that possibility.
+	if !clInfo.inKey && !clInfo.maybeInFieldName {
+		return nil, nil
 	}
 
 	plain, placeholder := &snippet.Builder{}, &snippet.Builder{}
@@ -52,7 +46,7 @@
 
 	// If the cursor position is on a different line from the literal's opening brace,
 	// we are in a multiline literal.
-	if c.view.FileSet().Position(c.pos).Line != c.view.FileSet().Position(lit.Lbrace).Line {
+	if c.view.FileSet().Position(c.pos).Line != c.view.FileSet().Position(clInfo.cl.Lbrace).Line {
 		plain.WriteText(",")
 		placeholder.WriteText(",")
 	}
diff --git a/internal/lsp/testdata/bar/bar.go.in b/internal/lsp/testdata/bar/bar.go.in
index 54f5c64..94d8254 100644
--- a/internal/lsp/testdata/bar/bar.go.in
+++ b/internal/lsp/testdata/bar/bar.go.in
@@ -23,16 +23,16 @@
 	var Valentine int //@item(Valentine, "Valentine", "int", "var")
 
 	_ = foo.StructFoo{
-		Val       //@complete("l", Value)
+		Valu //@complete(" //", Value)
 	}
   	_ = foo.StructFoo{
-		Va        //@complete("a", Value)
+		Va        //@complete("a", Value, Valentine)
 	}
 	_ = foo.StructFoo{
 		Value: 5, //@complete("a", Value)
 	}
 	_ = foo.StructFoo{
-		//@complete("", Value)
+		//@complete("", Value, Valentine, foo, Bar, helper)
 	}
 	_ = foo.StructFoo{
 		Value: Valen //@complete("le", Valentine)
diff --git a/internal/lsp/testdata/complit/complit.go.in b/internal/lsp/testdata/complit/complit.go.in
index 007cb65..e01fdb6 100644
--- a/internal/lsp/testdata/complit/complit.go.in
+++ b/internal/lsp/testdata/complit/complit.go.in
@@ -32,19 +32,35 @@
 		//@complete("", abVar, aaVar, structPosition)
 	}
 
-	_ = position{X: a} //@complete("}", abVar, aaVar)
-	_ = position{a}    //@complete("}", abVar, aaVar)
+	_ = []string{a: ""} //@complete(":", abVar, aaVar)
+	_ = [1]string{a: ""} //@complete(":", abVar, aaVar)
+
+	_ = position{X: a}   //@complete("}", abVar, aaVar)
+	_ = position{a}      //@complete("}", abVar, aaVar)
+	_ = position{a, }      //@complete("}", abVar, aaVar, structPosition)
 
 	_ = []int{a}  //@complete("a", abVar, aaVar, structPosition)
 	_ = [1]int{a} //@complete("a", abVar, aaVar, structPosition)
 
-	var s struct {
+	type myStruct struct {
 		AA int    //@item(fieldAA, "AA", "int", "field")
 		AB string //@item(fieldAB, "AB", "string", "field")
 	}
-	_ = map[int]string{1: "" + s.A} //@complete("}", fieldAB, fieldAA)
+
+	_ = myStruct{
+		AB: a, //@complete(",", aaVar, abVar)
+	}
+
+	var s myStruct
+
+	_ = map[int]string{1: "" + s.A}                                //@complete("}", fieldAB, fieldAA)
 	_ = map[int]string{1: (func(i int) string { return "" })(s.A)} //@complete(")}", fieldAA, fieldAB)
-	_ = map[int]string{1: func() string { s.A }} //@complete(" }", fieldAA, fieldAB)
+	_ = map[int]string{1: func() string { s.A }}                   //@complete(" }", fieldAA, fieldAB)
+
+	_ = position{s.A} //@complete("}", fieldAA, fieldAB)
+
+	var X int //@item(varX, "X", "int", "var")
+	_ = position{X}      //@complete("}", fieldX, varX)
 }
 
 func _() {
diff --git a/internal/lsp/testdata/importedcomplit/imported_complit.go b/internal/lsp/testdata/importedcomplit/imported_complit.go
new file mode 100644
index 0000000..99424ec
--- /dev/null
+++ b/internal/lsp/testdata/importedcomplit/imported_complit.go
@@ -0,0 +1,26 @@
+package importedcomplit
+
+import (
+	"golang.org/x/tools/internal/lsp/foo"
+)
+
+func _() {
+	var V int //@item(icVVar, "V", "int", "var")
+	_ = foo.StructFoo{V} //@complete("}", Value, icVVar)
+}
+
+func _() {
+	var (
+		aa string //@item(icAAVar, "aa", "string", "var")
+		ab int    //@item(icABVar, "ab", "int", "var")
+	)
+
+	_ = foo.StructFoo{a} //@complete("}", abVar, aaVar)
+
+	var s struct {
+		AA string //@item(icFieldAA, "AA", "string", "field")
+		AB int    //@item(icFieldAB, "AB", "int", "field")
+	}
+
+	_ = foo.StructFoo{s.} //@complete("}", icFieldAB, icFieldAA)
+}
diff --git a/internal/lsp/tests/tests.go b/internal/lsp/tests/tests.go
index e099f69..fc5a606 100644
--- a/internal/lsp/tests/tests.go
+++ b/internal/lsp/tests/tests.go
@@ -28,7 +28,7 @@
 // We hardcode the expected number of test cases to ensure that all tests
 // are being executed. If a test is added, this number must be changed.
 const (
-	ExpectedCompletionsCount       = 97
+	ExpectedCompletionsCount       = 106
 	ExpectedDiagnosticsCount       = 17
 	ExpectedFormatCount            = 5
 	ExpectedDefinitionsCount       = 33