internal/lsp: refactor find-references and rename

The main goal is to push the package variant logic from internal/lsp
into internal/lsp/source so all users of internal/lsp/source benefit.

"references" and "rename" now have top-level source.References() and
source.Rename() entry points (as opposed to hanging off
source.Identifier()). I expanded objectsAtProtocolPos() to know about
implicit objects (type switch and import spec), and to
handle *ast.ImportSpec generically. This gets rid of special case
handling of *types.PkgName in various places.

The biggest practical benefit, though, is that "references" no longer
needs to compute the objectpath for every types.Object comparison it
does, instead using direct types.Object equality. This speeds up
"references" and "rename" a lot.

Two other notable improvements that fell out of not using
source.Identifier()'s logic:

- Finding references on an embedded field now shows references to the
  field, not the type being embedded.
- Finding references on an imported object now works
  correctly (previously it searched the importing package's dependents
  rather than the imported package's dependents).

Finally, I refactored findIdentifier() to use pathEnclosingObjNode()
instead of astutil.PathEnclosingInterval. Now we only need a single
call to get the path because pathEnclosingObjNode() has the
"try pos || try pos-1" logic built in.

Change-Id: I667be9bed6ad83912404b90257c5c1485b3a7025
Reviewed-on: https://go-review.googlesource.com/c/tools/+/211999
Run-TryBot: Muir Manders <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Rebecca Stambler <[email protected]>

Author: muirdm

internal/lsp/references.go

@@ -27,62 +27,23 @@ func (s *Server) references(ctx context.Context, params *protocol.ReferenceParam
 	if fh.Identity().Kind != source.Go {
 		return nil, nil
 	}
-	phs, err := snapshot.PackageHandles(ctx, fh)
+
+	references, err := source.References(ctx, view.Snapshot(), fh, params.Position, params.Context.IncludeDeclaration)
 	if err != nil {
-		return nil, nil
+		return nil, err
 	}
 
-	// Get the location of each reference to return as the result.
-	var (
-		locations []protocol.Location
-		seen      = make(map[span.Span]bool)
-		lastIdent *source.IdentifierInfo
-	)
-	for _, ph := range phs {
-		ident, err := source.Identifier(ctx, snapshot, fh, params.Position, source.SpecificPackageHandle(ph.ID()))
-		if err != nil {
-			return nil, err
-		}
-
-		lastIdent = ident
-
-		references, err := ident.References(ctx)
+	var locations []protocol.Location
+	for _, ref := range references {
+		refRange, err := ref.Range()
 		if err != nil {
 			return nil, err
 		}
 
-		for _, ref := range references {
-			refSpan, err := ref.Span()
-			if err != nil {
-				return nil, err
-			}
-			if seen[refSpan] {
-				continue // already added this location
-			}
-			seen[refSpan] = true
-			refRange, err := ref.Range()
-			if err != nil {
-				return nil, err
-			}
-			locations = append(locations, protocol.Location{
-				URI:   protocol.NewURI(ref.URI()),
-				Range: refRange,
-			})
-		}
-	}
-
-	// Only add the identifier's declaration if the client requests it.
-	if params.Context.IncludeDeclaration && lastIdent != nil {
-		rng, err := lastIdent.Declaration.Range()
-		if err != nil {
-			return nil, err
-		}
-		locations = append([]protocol.Location{
-			{
-				URI:   protocol.NewURI(lastIdent.Declaration.URI()),
-				Range: rng,
-			},
-		}, locations...)
+		locations = append(locations, protocol.Location{
+			URI:   protocol.NewURI(ref.URI()),
+			Range: refRange,
+		})
 	}
 
 	return locations, nil

internal/lsp/rename.go

@@ -26,14 +26,12 @@ func (s *Server) rename(ctx context.Context, params *protocol.RenameParams) (*pr
 	if fh.Identity().Kind != source.Go {
 		return nil, nil
 	}
-	ident, err := source.Identifier(ctx, snapshot, fh, params.Position, source.WidestPackageHandle)
-	if err != nil {
-		return nil, nil
-	}
-	edits, err := ident.Rename(ctx, params.NewName)
+
+	edits, err := source.Rename(ctx, snapshot, fh, params.Position, params.NewName)
 	if err != nil {
 		return nil, err
 	}
+
 	var docChanges []protocol.TextDocumentEdit
 	for uri, e := range edits {
 		fh, err := snapshot.GetFile(uri)
@@ -61,13 +59,10 @@ func (s *Server) prepareRename(ctx context.Context, params *protocol.PrepareRena
 	if fh.Identity().Kind != source.Go {
 		return nil, nil
 	}
-	ident, err := source.Identifier(ctx, snapshot, fh, params.Position, source.WidestPackageHandle)
-	if err != nil {
-		return nil, nil // ignore errors
-	}
+
 	// Do not return errors here, as it adds clutter.
 	// Returning a nil result means there is not a valid rename.
-	item, err := ident.PrepareRename(ctx)
+	item, err := source.PrepareRename(ctx, snapshot, fh, params.Position)
 	if err != nil {
 		return nil, nil // ignore errors
 	}

internal/lsp/source/identifier.go

@@ -46,17 +46,6 @@ type Declaration struct {
 	obj  types.Object
 }
 
-func (i *IdentifierInfo) DeclarationReferenceInfo() *ReferenceInfo {
-	return &ReferenceInfo{
-		Name:          i.Declaration.obj.Name(),
-		mappedRange:   i.Declaration.mappedRange,
-		obj:           i.Declaration.obj,
-		ident:         i.ident,
-		pkg:           i.pkg,
-		isDeclaration: true,
-	}
-}
-
 // Identifier returns identifier information for a position
 // in a file, accounting for a potentially incomplete selector.
 func Identifier(ctx context.Context, snapshot Snapshot, fh FileHandle, pos protocol.Position, selectPackage PackagePolicy) (*IdentifierInfo, error) {
@@ -84,62 +73,49 @@ func Identifier(ctx context.Context, snapshot Snapshot, fh FileHandle, pos proto
 
 var ErrNoIdentFound = errors.New("no identifier found")
 
-func findIdentifier(ctx context.Context, snapshot Snapshot, pkg Package, file *ast.File, pos token.Pos) (*IdentifierInfo, error) {
-	if result, err := identifier(ctx, snapshot, pkg, file, pos); err != nil || result != nil {
-		return result, err
-	}
-	// If the position is not an identifier but immediately follows
-	// an identifier or selector period (as is common when
-	// requesting a completion), use the path to the preceding node.
-	ident, err := identifier(ctx, snapshot, pkg, file, pos-1)
-	if ident == nil && err == nil {
-		err = ErrNoIdentFound
-	}
-	return ident, err
-}
-
-// identifier checks a single position for a potential identifier.
-func identifier(ctx context.Context, s Snapshot, pkg Package, file *ast.File, pos token.Pos) (*IdentifierInfo, error) {
-	var err error
-
+func findIdentifier(ctx context.Context, s Snapshot, pkg Package, file *ast.File, pos token.Pos) (*IdentifierInfo, error) {
 	// Handle import specs separately, as there is no formal position for a package declaration.
 	if result, err := importSpec(s, pkg, file, pos); result != nil || err != nil {
 		return result, err
 	}
-	path, _ := astutil.PathEnclosingInterval(file, pos, pos)
+	path := pathEnclosingObjNode(file, pos)
 	if path == nil {
-		return nil, errors.Errorf("can't find node enclosing position")
+		return nil, ErrNoIdentFound
 	}
 
 	view := s.View()
+
+	ident, _ := path[0].(*ast.Ident)
+	if ident == nil {
+		return nil, ErrNoIdentFound
+	}
+
 	result := &IdentifierInfo{
 		Snapshot:  s,
 		qf:        qualifier(file, pkg.GetTypes(), pkg.GetTypesInfo()),
 		pkg:       pkg,
-		ident:     searchForIdent(path[0]),
+		ident:     ident,
 		enclosing: searchForEnclosing(pkg, path),
 	}
 
-	// No identifier at the given position.
-	if result.ident == nil {
-		return nil, nil
-	}
 	var wasEmbeddedField bool
 	for _, n := range path[1:] {
 		if field, ok := n.(*ast.Field); ok {
 			wasEmbeddedField = len(field.Names) == 0
 			break
 		}
 	}
+
 	result.Name = result.ident.Name
+	var err error
 	if result.mappedRange, err = posToMappedRange(view, pkg, result.ident.Pos(), result.ident.End()); err != nil {
 		return nil, err
 	}
 	result.Declaration.obj = pkg.GetTypesInfo().ObjectOf(result.ident)
 	if result.Declaration.obj == nil {
 		// If there was no types.Object for the declaration, there might be an implicit local variable
 		// declaration in a type switch.
-		if objs := typeSwitchVar(pkg.GetTypesInfo(), path); len(objs) > 0 {
+		if objs := typeSwitchImplicits(pkg, path); len(objs) > 0 {
 			// There is no types.Object for the declaration of an implicit local variable,
 			// but all of the types.Objects associated with the usages of this variable can be
 			// used to connect it back to the declaration.
@@ -202,18 +178,6 @@ func identifier(ctx context.Context, s Snapshot, pkg Package, file *ast.File, po
 	return result, nil
 }
 
-func searchForIdent(n ast.Node) *ast.Ident {
-	switch node := n.(type) {
-	case *ast.Ident:
-		return node
-	case *ast.SelectorExpr:
-		return node.Sel
-	case *ast.StarExpr:
-		return searchForIdent(node.X)
-	}
-	return nil
-}
-
 func searchForEnclosing(pkg Package, path []ast.Node) types.Type {
 	for _, n := range path {
 		switch n := n.(type) {
@@ -326,31 +290,69 @@ func importSpec(s Snapshot, pkg Package, file *ast.File, pos token.Pos) (*Identi
 	return result, nil
 }
 
-// typeSwitchVar handles the special case of a local variable implicitly defined in a type switch.
-// In such cases, the definition of the implicit variable will not be recorded in the *types.Info.Defs map,
-// but rather in the *types.Info.Implicits map.
-func typeSwitchVar(info *types.Info, path []ast.Node) []types.Object {
-	if len(path) < 3 {
-		return nil
-	}
-	// Check for [Ident AssignStmt TypeSwitchStmt...]
-	if _, ok := path[0].(*ast.Ident); !ok {
+// typeSwitchImplicits returns all the implicit type switch objects
+// that correspond to the leaf *ast.Ident.
+func typeSwitchImplicits(pkg Package, path []ast.Node) []types.Object {
+	ident, _ := path[0].(*ast.Ident)
+	if ident == nil {
 		return nil
 	}
-	if _, ok := path[1].(*ast.AssignStmt); !ok {
-		return nil
+
+	var (
+		ts     *ast.TypeSwitchStmt
+		assign *ast.AssignStmt
+		cc     *ast.CaseClause
+		obj    = pkg.GetTypesInfo().ObjectOf(ident)
+	)
+
+	// Walk our ancestors to determine if our leaf ident refers to a
+	// type switch variable, e.g. the "a" from "switch a := b.(type)".
+Outer:
+	for i := 1; i < len(path); i++ {
+		switch n := path[i].(type) {
+		case *ast.AssignStmt:
+			// Check if ident is the "a" in "a := foo.(type)". The "a" in
+			// this case has no types.Object, so check for ident equality.
+			if len(n.Lhs) == 1 && n.Lhs[0] == ident {
+				assign = n
+			}
+		case *ast.CaseClause:
+			// Check if ident is a use of "a" within a case clause. Each
+			// case clause implicitly maps "a" to a different types.Object,
+			// so check if ident's object is the case clause's implicit
+			// object.
+			if obj != nil && pkg.GetTypesInfo().Implicits[n] == obj {
+				cc = n
+			}
+		case *ast.TypeSwitchStmt:
+			// Look for the type switch that owns our previously found
+			// *ast.AssignStmt or *ast.CaseClause.
+
+			if n.Assign == assign {
+				ts = n
+				break Outer
+			}
+
+			for _, stmt := range n.Body.List {
+				if stmt == cc {
+					ts = n
+					break Outer
+				}
+			}
+		}
 	}
-	sw, ok := path[2].(*ast.TypeSwitchStmt)
-	if !ok {
+
+	if ts == nil {
 		return nil
 	}
 
-	var res []types.Object
-	for _, stmt := range sw.Body.List {
-		obj := info.Implicits[stmt.(*ast.CaseClause)]
-		if obj != nil {
-			res = append(res, obj)
+	// Our leaf ident refers to a type switch variable. Fan out to the
+	// type switch's implicit case clause objects.
+	var objs []types.Object
+	for _, cc := range ts.Body.List {
+		if ccObj := pkg.GetTypesInfo().Implicits[cc]; ccObj != nil {
+			objs = append(objs, ccObj)
 		}
 	}
-	return res
+	return objs
 }