internal/lsp: cache multiple packages depending on parse modes

This change shifts our approach to make sure that a top-level package
only ever imports "trimmed" packages.

Change-Id: I63c35791ef6efad7dac248a9ff877835f46df9ed
Reviewed-on: https://go-review.googlesource.com/c/tools/+/196523
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
diff --git a/internal/lsp/cache/check.go b/internal/lsp/cache/check.go
index 0a6e844..0357a19 100644
--- a/internal/lsp/cache/check.go
+++ b/internal/lsp/cache/check.go
@@ -62,6 +62,19 @@
 	config *packages.Config
 }
 
+func (cph *checkPackageHandle) mode() source.ParseMode {
+	if len(cph.files) == 0 {
+		return -1
+	}
+	mode := cph.files[0].Mode()
+	for _, ph := range cph.files[1:] {
+		if ph.Mode() != mode {
+			return -1
+		}
+	}
+	return mode
+}
+
 // checkPackageData contains the data produced by type-checking a package.
 type checkPackageData struct {
 	memoize.NoCopy
@@ -85,18 +98,13 @@
 		log.Error(ctx, "no ParseGoHandles", err, telemetry.Package.Of(m.id))
 		return nil, err
 	}
-	key := checkPackageKey{
-		id:     string(m.id),
-		files:  hashParseKeys(phs),
-		config: hashConfig(imp.config),
-	}
 	cph := &checkPackageHandle{
 		m:       m,
 		files:   phs,
 		config:  imp.config,
 		imports: make(map[packagePath]*checkPackageHandle),
 	}
-	h := imp.view.session.cache.store.Bind(key, func(ctx context.Context) interface{} {
+	h := imp.view.session.cache.store.Bind(cph.key(), func(ctx context.Context) interface{} {
 		data := &checkPackageData{}
 		data.pkg, data.err = imp.typeCheck(ctx, cph, m)
 		return data
@@ -163,6 +171,14 @@
 	return data.pkg, data.err
 }
 
+func (cph *checkPackageHandle) key() checkPackageKey {
+	return checkPackageKey{
+		id:     string(cph.m.id),
+		files:  hashParseKeys(cph.files),
+		config: hashConfig(cph.config),
+	}
+}
+
 func (imp *importer) parseGoHandles(ctx context.Context, m *metadata) ([]source.ParseGoHandle, error) {
 	phs := make([]source.ParseGoHandle, 0, len(m.files))
 	for _, uri := range m.files {
@@ -343,9 +359,12 @@
 
 	// Set the package even if we failed to parse the file.
 	if gof.cphs == nil {
-		gof.cphs = make(map[packageID]source.CheckPackageHandle)
+		gof.cphs = make(map[packageKey]*checkPackageHandle)
 	}
-	gof.cphs[cph.m.id] = cph
+	gof.cphs[packageKey{
+		id:   cph.m.id,
+		mode: ph.Mode(),
+	}] = cph
 
 	file, _, _, err := ph.Parse(ctx)
 	if err != nil {
diff --git a/internal/lsp/cache/gofile.go b/internal/lsp/cache/gofile.go
index 6c55ef7..c22ae6b 100644
--- a/internal/lsp/cache/gofile.go
+++ b/internal/lsp/cache/gofile.go
@@ -33,15 +33,20 @@
 
 	imports []*ast.ImportSpec
 
-	cphs map[packageID]source.CheckPackageHandle
+	cphs map[packageKey]*checkPackageHandle
 	meta map[packageID]*metadata
 }
 
+type packageKey struct {
+	id   packageID
+	mode source.ParseMode
+}
+
 func (f *goFile) CheckPackageHandles(ctx context.Context) ([]source.CheckPackageHandle, error) {
 	ctx = telemetry.File.With(ctx, f.URI())
 	fh := f.Handle(ctx)
 
-	if f.isDirty(ctx, fh) || f.wrongParseMode(ctx, fh, source.ParseFull) {
+	if f.isDirty(ctx, fh) {
 		if err := f.view.loadParseTypecheck(ctx, f, fh); err != nil {
 			return nil, err
 		}
@@ -50,14 +55,22 @@
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
-	if len(f.cphs) == 0 {
+	var results []source.CheckPackageHandle
+	seenIDs := make(map[string]bool)
+	for _, cph := range f.cphs {
+		if seenIDs[cph.ID()] {
+			continue
+		}
+		if cph.mode() < source.ParseFull {
+			continue
+		}
+		results = append(results, cph)
+		seenIDs[cph.ID()] = true
+	}
+	if len(results) == 0 {
 		return nil, errors.Errorf("no CheckPackageHandles for %s", f.URI())
 	}
-	var cphs []source.CheckPackageHandle
-	for _, cph := range f.cphs {
-		cphs = append(cphs, cph)
-	}
-	return cphs, nil
+	return results, nil
 }
 
 func (f *goFile) GetActiveReverseDeps(ctx context.Context) (files []source.GoFile) {
@@ -119,20 +132,6 @@
 	return result
 }
 
-func (f *goFile) wrongParseMode(ctx context.Context, fh source.FileHandle, mode source.ParseMode) bool {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-
-	for _, cph := range f.cphs {
-		for _, ph := range cph.Files() {
-			if fh.Identity() == ph.File().Identity() {
-				return ph.Mode() < mode
-			}
-		}
-	}
-	return true
-}
-
 // isDirty is true if the file needs to be type-checked.
 // It assumes that the file's view's mutex is held by the caller.
 func (f *goFile) isDirty(ctx context.Context, fh source.FileHandle) bool {
@@ -145,7 +144,7 @@
 	// Note: This must be the first case, otherwise we may not reset the value of f.justOpened.
 	if f.justOpened {
 		f.meta = make(map[packageID]*metadata)
-		f.cphs = make(map[packageID]source.CheckPackageHandle)
+		f.cphs = make(map[packageKey]*checkPackageHandle)
 		f.justOpened = false
 		return true
 	}
@@ -155,9 +154,11 @@
 	if len(f.missingImports) > 0 {
 		return true
 	}
-	for _, cph := range f.cphs {
+	for key, cph := range f.cphs {
+		if key.mode != source.ParseFull {
+			continue
+		}
 		for _, file := range cph.Files() {
-			// There is a type-checked package for the current file handle.
 			if file.File().Identity() == fh.Identity() {
 				return false
 			}
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index bb2477f..e858164 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -57,24 +57,6 @@
 	view.mcache.mu.Lock()
 	defer view.mcache.mu.Unlock()
 
-	var toDelete []packageID
-	f.mu.Lock()
-	for id, cph := range f.cphs {
-		if cph != nil {
-			toDelete = append(toDelete, id)
-		}
-	}
-	f.mu.Unlock()
-
-	// If the AST for this file is trimmed, and we are explicitly type-checking it,
-	// don't ignore function bodies.
-	if f.wrongParseMode(ctx, fh, source.ParseFull) {
-		// Remove the package and all of its reverse dependencies from the cache.
-		for _, id := range toDelete {
-			f.view.remove(ctx, id, map[packageID]struct{}{})
-		}
-	}
-
 	// Get the metadata for the file.
 	meta, err := view.checkMetadata(ctx, f, fh)
 	if err != nil {
diff --git a/internal/lsp/cache/view.go b/internal/lsp/cache/view.go
index c0f4356..1b6bccb 100644
--- a/internal/lsp/cache/view.go
+++ b/internal/lsp/cache/view.go
@@ -2,6 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+// Package cache implements the caching layer for gopls.
 package cache
 
 import (
@@ -331,11 +332,11 @@
 	f.view.mcache.mu.Lock()
 	defer f.view.mcache.mu.Unlock()
 
-	var toDelete []packageID
+	toDelete := make(map[packageID]bool)
 	f.mu.Lock()
-	for id, cph := range f.cphs {
+	for key, cph := range f.cphs {
 		if cph != nil {
-			toDelete = append(toDelete, id)
+			toDelete[key.id] = true
 		}
 	}
 	f.mu.Unlock()
@@ -344,7 +345,7 @@
 	defer f.handleMu.Unlock()
 
 	// Remove the package and all of its reverse dependencies from the cache.
-	for _, id := range toDelete {
+	for id := range toDelete {
 		f.view.remove(ctx, id, map[packageID]struct{}{})
 	}
 
@@ -405,7 +406,15 @@
 			continue
 		}
 		gof.mu.Lock()
-		delete(gof.cphs, id)
+		for _, mode := range []source.ParseMode{
+			source.ParseExported,
+			source.ParseFull,
+		} {
+			delete(gof.cphs, packageKey{
+				id:   id,
+				mode: mode,
+			})
+		}
 		gof.mu.Unlock()
 	}
 	return