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 <matloob@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/cache/session.go b/internal/lsp/cache/session.go
index eb1c85c..cb80f19 100644
--- a/internal/lsp/cache/session.go
+++ b/internal/lsp/cache/session.go
@@ -110,12 +110,13 @@
 		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 @@
 	// 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
 }
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index dcd95e4..8e2b4d6 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -39,6 +39,10 @@
 
 	// 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 @@
 	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 @@
 	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 @@
 		}
 		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
diff --git a/internal/lsp/cmd/references.go b/internal/lsp/cmd/references.go
index 03c1709..8b91356 100644
--- a/internal/lsp/cmd/references.go
+++ b/internal/lsp/cmd/references.go
@@ -8,10 +8,11 @@
 	"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
diff --git a/internal/lsp/cmd/test/references.go b/internal/lsp/cmd/test/references.go
index 55bbf91..9cb5a66 100644
--- a/internal/lsp/cmd/test/references.go
+++ b/internal/lsp/cmd/test/references.go
@@ -6,17 +6,24 @@
 
 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()
diff --git a/internal/lsp/source/references.go b/internal/lsp/source/references.go
index 5fd20f8..9746760 100644
--- a/internal/lsp/source/references.go
+++ b/internal/lsp/source/references.go
@@ -9,6 +9,7 @@
 	"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 @@
 			mappedRange:   rng,
 		}}, references...)
 	}
-	for ident, obj := range info.Uses {
-		if obj == nil || !sameObj(obj, i.Declaration.obj) {
-			continue
+	// 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,
+			})
 		}
-		rng, err := posToMappedRange(ctx, i.Snapshot.View(), i.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,
-		})
 	}
 	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
 }
diff --git a/internal/lsp/testdata/references/other/other.go b/internal/lsp/testdata/references/other/other.go
new file mode 100644
index 0000000..3987d35
--- /dev/null
+++ b/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")
+}
diff --git a/internal/lsp/testdata/references/refs.go b/internal/lsp/testdata/references/refs.go
index a164715..e5a51fd 100644
--- a/internal/lsp/testdata/references/refs.go
+++ b/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) {}
diff --git a/internal/lsp/testdata/summary.txt.golden b/internal/lsp/testdata/summary.txt.golden
index baa0d82..ade9cf3 100644
--- a/internal/lsp/testdata/summary.txt.golden
+++ b/internal/lsp/testdata/summary.txt.golden
@@ -14,7 +14,7 @@
 DefinitionsCount = 38
 TypeDefinitionsCount = 2
 HighlightsCount = 2
-ReferencesCount = 6
+ReferencesCount = 7
 RenamesCount = 20
 PrepareRenamesCount = 8
 SymbolsCount = 1