From a0c0a26af391dfb6814dc0ca97a568a836f3f81f Mon Sep 17 00:00:00 2001
From: Muir Manders <muir@mnd.rs>
Date: Sat, 8 Jun 2019 20:11:23 -0700
Subject: [PATCH] internal/lsp: fix errors when adding new file to existing
 package

Previously when you added a new file to an existing package, the new
file would get stuck with the "no package for file" error until you
saved the file and then made changed a different file in the
package. There were two changes required to fix the errors:

First, we need to invalidate the package cache when a new file is
added to a package so that the package will actually re-parse and
re-type check. We now notice if file names changed when updating a
package's metadata and invalidate the package cache accordingly.

Second, when dealing with overlay (unsaved) files, we need to map
the *goFile to the package even if we fail to parse the
file (e.g. the new file fails to parse when it is empty). If we don't
map it to the package, the package won't get refreshed as the file is
changed.

Fixes golang/go#32341
---
 internal/lsp/cache/check.go | 34 ++++++++++++++++++++++------------
 internal/lsp/cache/load.go  | 28 ++++++++++++++++++++++++++++
 internal/lsp/cache/parse.go | 16 +++++++---------
 internal/lsp/cache/pkg.go   |  8 ++++----
 internal/lsp/cache/view.go  |  8 ++++++++
 5 files changed, 69 insertions(+), 25 deletions(-)

diff --git a/internal/lsp/cache/check.go b/internal/lsp/cache/check.go
index c295c298a0d..a0829e7f616 100644
--- a/internal/lsp/cache/check.go
+++ b/internal/lsp/cache/check.go
@@ -160,18 +160,10 @@ func (imp *importer) typeCheck(pkgPath string) (*pkg, error) {
 }
 
 func (imp *importer) cachePackage(ctx context.Context, pkg *pkg, meta *metadata) {
-	for _, fAST := range pkg.syntax {
+	for _, filename := range pkg.files {
 		// TODO: If a file is in multiple packages, which package do we store?
-		if !fAST.file.Pos().IsValid() {
-			imp.view.Session().Logger().Errorf(ctx, "invalid position for file %v", fAST.file.Name)
-			continue
-		}
-		tok := imp.view.Session().Cache().FileSet().File(fAST.file.Pos())
-		if tok == nil {
-			imp.view.Session().Logger().Errorf(ctx, "no token.File for %v", fAST.file.Name)
-			continue
-		}
-		fURI := span.FileURI(tok.Name())
+
+		fURI := span.FileURI(filename)
 		f, err := imp.view.getFile(fURI)
 		if err != nil {
 			imp.view.Session().Logger().Errorf(ctx, "no file: %v", err)
@@ -182,10 +174,28 @@ func (imp *importer) cachePackage(ctx context.Context, pkg *pkg, meta *metadata)
 			imp.view.Session().Logger().Errorf(ctx, "%v is not a Go file", f.URI())
 			continue
 		}
+
+		// Set the package even if we failed to parse the file otherwise
+		// future updates to this file won't refresh the package.
+		gof.pkg = pkg
+
+		fAST := pkg.syntax[filename]
+		if fAST == nil {
+			continue
+		}
+
+		if !fAST.file.Pos().IsValid() {
+			imp.view.Session().Logger().Errorf(ctx, "invalid position for file %v", fAST.file.Name)
+			continue
+		}
+		tok := imp.view.Session().Cache().FileSet().File(fAST.file.Pos())
+		if tok == nil {
+			imp.view.Session().Logger().Errorf(ctx, "no token.File for %v", fAST.file.Name)
+			continue
+		}
 		gof.token = tok
 		gof.ast = fAST
 		gof.imports = fAST.file.Imports
-		gof.pkg = pkg
 	}
 
 	// Set imports of package to correspond to cached packages.
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index e8b6890f27c..46bdfa9518d 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -142,6 +142,13 @@ func (v *view) parseImports(ctx context.Context, f *goFile) bool {
 
 func (v *view) link(ctx context.Context, pkgPath string, pkg *packages.Package, parent *metadata) *metadata {
 	m, ok := v.mcache.packages[pkgPath]
+
+	// If a file was added or deleted we need to invalidate the package cache
+	// so relevant packages get parsed and type checked again.
+	if ok && !filenamesIdentical(m.files, pkg.CompiledGoFiles) {
+		v.invalidatePackage(pkgPath)
+	}
+
 	if !ok {
 		m = &metadata{
 			pkgPath:        pkgPath,
@@ -189,3 +196,24 @@ func (v *view) link(ctx context.Context, pkgPath string, pkg *packages.Package,
 	}
 	return m
 }
+
+// filenamesIdentical reports whether two sets of file names are identical.
+func filenamesIdentical(oldFiles, newFiles []string) bool {
+	if len(oldFiles) != len(newFiles) {
+		return false
+	}
+
+	oldByName := make(map[string]struct{}, len(oldFiles))
+	for _, filename := range oldFiles {
+		oldByName[filename] = struct{}{}
+	}
+
+	for _, newFilename := range newFiles {
+		if _, found := oldByName[newFilename]; !found {
+			return false
+		}
+		delete(oldByName, newFilename)
+	}
+
+	return len(oldByName) == 0
+}
diff --git a/internal/lsp/cache/parse.go b/internal/lsp/cache/parse.go
index 113743a4b50..ad1859903de 100644
--- a/internal/lsp/cache/parse.go
+++ b/internal/lsp/cache/parse.go
@@ -34,7 +34,7 @@ var ioLimit = make(chan bool, 20)
 // Because files are scanned in parallel, the token.Pos
 // positions of the resulting ast.Files are not ordered.
 //
-func (imp *importer) parseFiles(filenames []string, ignoreFuncBodies bool) ([]*astFile, []error, error) {
+func (imp *importer) parseFiles(filenames []string, ignoreFuncBodies bool) (map[string]*astFile, []error, error) {
 	var (
 		wg       sync.WaitGroup
 		mu       sync.Mutex
@@ -121,17 +121,15 @@ func (imp *importer) parseFiles(filenames []string, ignoreFuncBodies bool) ([]*a
 		return nil, nil, fatalErr
 	}
 
-	// Eliminate nils, preserving order.
-	var o int
-	for _, f := range parsed {
+	parsedByFilename := make(map[string]*astFile)
+
+	for i, f := range parsed {
 		if f != nil {
-			parsed[o] = f
-			o++
+			parsedByFilename[filenames[i]] = f
 		}
 	}
-	parsed = parsed[:o]
 
-	o = 0
+	var o int
 	for _, err := range errors {
 		if err != nil {
 			errors[o] = err
@@ -140,7 +138,7 @@ func (imp *importer) parseFiles(filenames []string, ignoreFuncBodies bool) ([]*a
 	}
 	errors = errors[:o]
 
-	return parsed, errors, nil
+	return parsedByFilename, errors, nil
 }
 
 // sameFile returns true if x and y have the same basename and denote
diff --git a/internal/lsp/cache/pkg.go b/internal/lsp/cache/pkg.go
index caf1ee5a413..d4a04f6eafe 100644
--- a/internal/lsp/cache/pkg.go
+++ b/internal/lsp/cache/pkg.go
@@ -20,7 +20,7 @@ import (
 type pkg struct {
 	id, pkgPath string
 	files       []string
-	syntax      []*astFile
+	syntax      map[string]*astFile
 	errors      []packages.Error
 	imports     map[string]*pkg
 	types       *types.Package
@@ -137,9 +137,9 @@ func (pkg *pkg) GetFilenames() []string {
 }
 
 func (pkg *pkg) GetSyntax() []*ast.File {
-	syntax := make([]*ast.File, len(pkg.syntax))
-	for i := range pkg.syntax {
-		syntax[i] = pkg.syntax[i].file
+	syntax := make([]*ast.File, 0, len(pkg.syntax))
+	for _, astFile := range pkg.syntax {
+		syntax = append(syntax, astFile.file)
 	}
 	return syntax
 }
diff --git a/internal/lsp/cache/view.go b/internal/lsp/cache/view.go
index 7d81fb05130..b127856908b 100644
--- a/internal/lsp/cache/view.go
+++ b/internal/lsp/cache/view.go
@@ -252,6 +252,14 @@ func (f *goFile) invalidateAST() {
 	}
 }
 
+// invalidatePackage removes the specified package and dependents from the
+// package cache.
+func (v *view) invalidatePackage(pkgPath string) {
+	v.pcache.mu.Lock()
+	defer v.pcache.mu.Unlock()
+	v.remove(pkgPath, make(map[string]struct{}))
+}
+
 // remove invalidates a package and its reverse dependencies in the view's
 // package cache. It is assumed that the caller has locked both the mutexes
 // of both the mcache and the pcache.