internal/lsp/source: add support for references in the same workspace

When looking for references, look in the entire workspace rather than
the same package. This makes the references query more expensive because
it needs to look at every package in the workspace, but hopefully
it shouln't be user-noticable. This can be made more efficient by only
checking packages that are transitive reverse dependencies. I don't think a
mechanism to get all transitive reverse dependencies exists yet.

One of the references test have been changed: it looked up references
of the builtin int type, but now there are so many refererences that
the test too slow and doesn't make sense any more. Instead look up
references of the type "i" in that file.

Change-Id: I93b3bd3795386f06ce488e76e6c7c8c1b1074e22
Reviewed-on: https://go-review.googlesource.com/c/tools/+/206883
Run-TryBot: Michael Matloob <[email protected]>
Reviewed-by: Rebecca Stambler <[email protected]>

Author: matloob

internal/lsp/cache/session.go

@@ -110,12 +110,13 @@ func (s *session) createView(ctx context.Context, name string, folder span.URI,
 		filesByURI:    make(map[span.URI]viewFile),
 		filesByBase:   make(map[string][]viewFile),
 		snapshot: &snapshot{
-			packages:   make(map[packageKey]*checkPackageHandle),
-			ids:        make(map[span.URI][]packageID),
-			metadata:   make(map[packageID]*metadata),
-			files:      make(map[span.URI]source.FileHandle),
-			importedBy: make(map[packageID][]packageID),
-			actions:    make(map[actionKey]*actionHandle),
+			packages:          make(map[packageKey]*checkPackageHandle),
+			ids:               make(map[span.URI][]packageID),
+			metadata:          make(map[packageID]*metadata),
+			files:             make(map[span.URI]source.FileHandle),
+			importedBy:        make(map[packageID][]packageID),
+			actions:           make(map[actionKey]*actionHandle),
+			workspacePackages: make(map[packageID]bool),
 		},
 		ignoredURIs: make(map[span.URI]struct{}),
 		builtin:     &builtinPkg{},
@@ -146,12 +147,10 @@ func (s *session) createView(ctx context.Context, name string, folder span.URI,
 	// Prepare CheckPackageHandles for every package that's been loaded.
 	// (*snapshot).CheckPackageHandle makes the assumption that every package that's
 	// been loaded has an existing checkPackageHandle.
-	for _, m := range m {
-		_, err := v.snapshot.checkPackageHandle(ctx, m.id, source.ParseFull)
-		if err != nil {
-			return nil, err
-		}
+	if err := v.snapshot.checkWorkspacePackages(ctx, m); err != nil {
+		return nil, err
 	}
+
 	debug.AddView(debugView{v})
 	return v, loadErr
 }

internal/lsp/cache/snapshot.go

@@ -39,6 +39,10 @@ type snapshot struct {
 
 	// actions maps an actionkey to its actionHandle.
 	actions map[actionKey]*actionHandle
+
+	// workspacePackages contains the workspace's packages, which are loaded
+	// when the view is created.
+	workspacePackages map[packageID]bool
 }
 
 type packageKey struct {
@@ -99,16 +103,50 @@ func (s *snapshot) getPackages(uri source.FileURI, m source.ParseMode) (cphs []s
 	return cphs
 }
 
+// checkWorkspacePackages checks the initial set of packages loaded when
+// the view is created. This is needed because
+// (*snapshot).CheckPackageHandle makes the assumption that every package that's
+// been loaded has an existing checkPackageHandle.
+func (s *snapshot) checkWorkspacePackages(ctx context.Context, m []*metadata) error {
+	for _, m := range m {
+		_, err := s.checkPackageHandle(ctx, m.id, source.ParseFull)
+		if err != nil {
+			return err
+		}
+		s.workspacePackages[m.id] = true
+	}
+	return nil
+}
+
 func (s *snapshot) KnownPackages(ctx context.Context) []source.Package {
 	// TODO(matloob): This function exists because KnownImportPaths can't
 	// determine the import paths of all packages. Remove this function
 	// if KnownImportPaths gains that ability. That could happen if
 	// go list or go packages provide that information.
+	pkgIDs := make(map[packageID]bool)
 	s.mu.Lock()
-	defer s.mu.Unlock()
+	for _, m := range s.metadata {
+		pkgIDs[m.id] = true
+	}
+	// Add in all the workspacePackages in case the've been invalidated
+	// in the metadata since their initial load.
+	for id := range s.workspacePackages {
+		pkgIDs[id] = true
+	}
+	s.mu.Unlock()
 
 	var results []source.Package
-	for _, cph := range s.packages {
+	for pkgID := range pkgIDs {
+		mode := source.ParseExported
+		if s.workspacePackages[pkgID] {
+			// Any package in our workspace should be loaded with ParseFull.
+			mode = source.ParseFull
+		}
+		cph, err := s.checkPackageHandle(ctx, pkgID, mode)
+		if err != nil {
+			log.Error(ctx, fmt.Sprintf("cph.Check of %v", cph.m.pkgPath), err)
+			continue
+		}
 		// Check the package now if it's not checked yet.
 		// TODO(matloob): is this too slow?
 		pkg, err := cph.check(ctx)
@@ -279,14 +317,15 @@ func (s *snapshot) clone(ctx context.Context, withoutURI *span.URI, withoutTypes
 	defer s.mu.Unlock()
 
 	result := &snapshot{
-		id:         s.id + 1,
-		view:       s.view,
-		ids:        make(map[span.URI][]packageID),
-		importedBy: make(map[packageID][]packageID),
-		metadata:   make(map[packageID]*metadata),
-		packages:   make(map[packageKey]*checkPackageHandle),
-		actions:    make(map[actionKey]*actionHandle),
-		files:      make(map[span.URI]source.FileHandle),
+		id:                s.id + 1,
+		view:              s.view,
+		ids:               make(map[span.URI][]packageID),
+		importedBy:        make(map[packageID][]packageID),
+		metadata:          make(map[packageID]*metadata),
+		packages:          make(map[packageKey]*checkPackageHandle),
+		actions:           make(map[actionKey]*actionHandle),
+		files:             make(map[span.URI]source.FileHandle),
+		workspacePackages: make(map[packageID]bool),
 	}
 	// Copy all of the FileHandles except for the one that was invalidated.
 	for k, v := range s.files {
@@ -344,6 +383,10 @@ func (s *snapshot) clone(ctx context.Context, withoutURI *span.URI, withoutTypes
 		}
 		result.metadata[k] = v
 	}
+	// Copy the set of initally loaded packages
+	for k, v := range s.workspacePackages {
+		result.workspacePackages[k] = v
+	}
 	// Don't bother copying the importedBy graph,
 	// as it changes each time we update metadata.
 	return result

internal/lsp/cmd/references.go

@@ -8,10 +8,11 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"sort"
+
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/tool"
-	"sort"
 )
 
 // references implements the references verb for gopls

internal/lsp/cmd/test/references.go

@@ -6,17 +6,24 @@ package cmdtest
 
 import (
 	"fmt"
+	"sort"
+	"testing"
+
 	"golang.org/x/tools/internal/lsp/cmd"
 	"golang.org/x/tools/internal/tool"
-	"testing"
 
 	"golang.org/x/tools/internal/span"
 )
 
 func (r *runner) References(t *testing.T, spn span.Span, itemList []span.Span) {
-	var expect string
+	var itemStrings []string
 	for _, i := range itemList {
-		expect += fmt.Sprintln(i)
+		itemStrings = append(itemStrings, fmt.Sprint(i))
+	}
+	sort.Strings(itemStrings)
+	var expect string
+	for _, i := range itemStrings {
+		expect += i + "\n"
 	}
 
 	uri := spn.URI()

internal/lsp/source/references.go

@@ -9,6 +9,7 @@ import (
 	"go/ast"
 	"go/types"
 
+	"golang.org/x/tools/go/types/objectpath"
 	"golang.org/x/tools/internal/telemetry/trace"
 	errors "golang.org/x/xerrors"
 )
@@ -69,29 +70,47 @@ func (i *IdentifierInfo) References(ctx context.Context) ([]*ReferenceInfo, erro
 			mappedRange:   rng,
 		}}, references...)
 	}
-	for ident, obj := range info.Uses {
-		if obj == nil || !sameObj(obj, i.Declaration.obj) {
-			continue
-		}
-		rng, err := posToMappedRange(ctx, i.Snapshot.View(), i.pkg, ident.Pos(), ident.End())
-		if err != nil {
-			return nil, err
+	// TODO(matloob): This only needs to look into reverse-dependencies.
+	// Avoid checking types of other packages.
+	for _, pkg := range i.Snapshot.KnownPackages(ctx) {
+		for ident, obj := range pkg.GetTypesInfo().Uses {
+			if obj == nil || !(sameObj(obj, i.Declaration.obj)) {
+				continue
+			}
+			rng, err := posToMappedRange(ctx, i.Snapshot.View(), pkg, ident.Pos(), ident.End())
+			if err != nil {
+				return nil, err
+			}
+			references = append(references, &ReferenceInfo{
+				Name:        ident.Name,
+				ident:       ident,
+				pkg:         i.pkg,
+				obj:         obj,
+				mappedRange: rng,
+			})
 		}
-		references = append(references, &ReferenceInfo{
-			Name:        ident.Name,
-			ident:       ident,
-			pkg:         i.pkg,
-			obj:         obj,
-			mappedRange: rng,
-		})
 	}
 	return references, nil
 }
 
 // sameObj returns true if obj is the same as declObj.
-// Objects are the same if they have the some Pos and Name.
+// Objects are the same if either they have they have objectpaths
+// and their objectpath and package are the same; or if they don't
+// have object paths and they have the same Pos and Name.
 func sameObj(obj, declObj types.Object) bool {
 	// TODO(suzmue): support the case where an identifier may have two different
 	// declaration positions.
-	return obj.Pos() == declObj.Pos() && obj.Name() == declObj.Name()
+	if obj.Pkg() == nil || declObj.Pkg() == nil {
+		if obj.Pkg() != declObj.Pkg() {
+			return false
+		}
+	} else if obj.Pkg().Path() != declObj.Pkg().Path() {
+		return false
+	}
+	objPath, operr := objectpath.For(obj)
+	declObjPath, doperr := objectpath.For(declObj)
+	if operr != nil || doperr != nil {
+		return obj.Pos() == declObj.Pos() && obj.Name() == declObj.Name()
+	}
+	return objPath == declObjPath
 }

internal/lsp/testdata/references/other/other.go

@@ -0,0 +1,11 @@
+package other
+
+import (
+	references "golang.org/x/tools/internal/lsp/references"
+)
+
+func _() {
+	references.Q = "hello" //@mark(assignExpQ, "Q")
+	bob := func(_ string) {}
+	bob(references.Q) //@mark(bobExpQ, "Q")
+}

internal/lsp/testdata/references/refs.go

@@ -1,17 +1,19 @@
 package refs
 
-type i int //@mark(typeInt, "int"),refs("int", typeInt, argInt, returnInt)
+type i int //@mark(typeI, "i"),refs("i", typeI, argI, returnI)
 
-func _(_ int) []bool { //@mark(argInt, "int")
+func _(_ i) []bool { //@mark(argI, "i")
 	return nil
 }
 
-func _(_ string) int { //@mark(returnInt, "int")
+func _(_ []byte) i { //@mark(returnI, "i")
 	return 0
 }
 
 var q string //@mark(declQ, "q"),refs("q", declQ, assignQ, bobQ)
 
+var Q string //@mark(declExpQ, "Q"), refs("Q", declExpQ, assignExpQ, bobExpQ)
+
 func _() {
 	q = "hello" //@mark(assignQ, "q")
 	bob := func(_ string) {}

internal/lsp/testdata/summary.txt.golden

@@ -14,7 +14,7 @@ SuggestedFixCount = 1
 DefinitionsCount = 38
 TypeDefinitionsCount = 2
 HighlightsCount = 2
-ReferencesCount = 6
+ReferencesCount = 7
 RenamesCount = 20
 PrepareRenamesCount = 8
 SymbolsCount = 1