internal/lsp: use the memoize package to get *token.Files

This change does not actually use the token handle for GetToken right
now, but implements the approach for memoizing *token.Files.

Change-Id: I75919f4e97abd6893b202c021adecd2c9dbfc2be
Reviewed-on: https://go-review.googlesource.com/c/tools/+/182277
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>

internal/lsp/cache/check.go

@@ -157,41 +157,40 @@ func (imp *importer) typeCheck(id packageID) (*pkg, error) {
 
 func (imp *importer) cachePackage(ctx context.Context, pkg *pkg, meta *metadata) {
 	for _, filename := range pkg.files {
-		// TODO: If a file is in multiple packages, which package do we store?
+		f, err := imp.view.getFile(span.FileURI(filename))
-
-		fURI := span.FileURI(filename)
-		f, err := imp.view.getFile(fURI)
 		if err != nil {
-			imp.view.Session().Logger().Errorf(ctx, "no file: %v", err)
+			imp.view.session.log.Errorf(ctx, "no file: %v", err)
 			continue
 		}
 		gof, ok := f.(*goFile)
 		if !ok {
-			imp.view.Session().Logger().Errorf(ctx, "%v is not a Go file", f.URI())
+			imp.view.session.log.Errorf(ctx, "%v is not a Go file", f.URI())
 			continue
 		}
-
+		// Set the package even if we failed to parse the file.
-		// 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]
+		// Get the *token.File directly from the AST.
-		if fAST == nil {
+		gof.ast = pkg.syntax[filename]
+		if gof.ast == nil {
+			imp.view.session.log.Errorf(ctx, "no AST information for %s", filename)
 			continue
 		}
-
+		if gof.ast.file == nil {
-		if !fAST.file.Pos().IsValid() {
+			imp.view.session.log.Errorf(ctx, "no AST for %s", filename)
-			imp.view.Session().Logger().Errorf(ctx, "invalid position for file %v", fAST.file.Name)
+		}
+		pos := gof.ast.file.Pos()
+		if !pos.IsValid() {
+			imp.view.session.log.Errorf(ctx, "AST for %s has an invalid position", filename)
 			continue
 		}
-		tok := imp.view.Session().Cache().FileSet().File(fAST.file.Pos())
+		tok := imp.view.session.cache.FileSet().File(pos)
 		if tok == nil {
-			imp.view.Session().Logger().Errorf(ctx, "no token.File for %v", fAST.file.Name)
+			imp.view.session.log.Errorf(ctx, "no *token.File for %s", filename)
 			continue
 		}
 		gof.token = tok
-		gof.ast = fAST
+		gof.imports = gof.ast.file.Imports
-		gof.imports = fAST.file.Imports
 	}
 
 	// Set imports of package to correspond to cached packages.

internal/lsp/cache/file.go

@@ -29,10 +29,12 @@ type fileBase struct {
 	uris  []span.URI
 	fname string
 
-	view     *view
+	view *view
+
 	handleMu sync.Mutex
 	handle   source.FileHandle
-	token    *token.File
+
+	token *token.File
 }
 
 func basename(filename string) string {

internal/lsp/cache/load.go

@@ -112,7 +112,7 @@ func (v *view) parseImports(ctx context.Context, f *goFile) bool {
 		return true
 	}
 	// Get file content in case we don't already have it.
-	parsed, _ := v.session.cache.ParseGo(f.Handle(ctx), source.ParseHeader).Parse(ctx)
+	parsed, _ := v.session.cache.ParseGoHandle(f.Handle(ctx), source.ParseHeader).Parse(ctx)
 	if parsed == nil {
 		return true
 	}

internal/lsp/cache/parse.go

@@ -24,6 +24,7 @@ import (
 // Limits the number of parallel parser calls per process.
 var parseLimit = make(chan bool, 20)
 
+// parseKey uniquely identifies a parsed Go file.
 type parseKey struct {
 	file source.FileIdentity
 	mode source.ParseMode
@@ -37,11 +38,12 @@ type parseGoHandle struct {
 
 type parseGoData struct {
 	memoize.NoCopy
+
 	ast *ast.File
 	err error
 }
 
-func (c *cache) ParseGo(fh source.FileHandle, mode source.ParseMode) source.ParseGoHandle {
+func (c *cache) ParseGoHandle(fh source.FileHandle, mode source.ParseMode) source.ParseGoHandle {
 	key := parseKey{
 		file: fh.Identity(),
 		mode: mode,
@@ -104,7 +106,6 @@ func parseGo(ctx context.Context, c *cache, fh source.FileHandle, mode source.Pa
 //
 // 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) (map[string]*astFile, []error, error) {
 	var (
 		wg     sync.WaitGroup
@@ -124,7 +125,7 @@ func (imp *importer) parseFiles(filenames []string, ignoreFuncBodies bool) (map[
 		if ignoreFuncBodies {
 			mode = source.ParseExported
 		}
-		ph := imp.view.session.cache.ParseGo(fh, mode)
+		ph := imp.view.session.cache.ParseGoHandle(fh, mode)
 		// now read and parse in parallel
 		wg.Add(1)
 		go func(i int, filename string) {

internal/lsp/cache/token.go

@@ -0,0 +1,88 @@
+package cache
+
+import (
+	"context"
+	"fmt"
+	"go/token"
+
+	"golang.org/x/tools/internal/lsp/source"
+	"golang.org/x/tools/internal/memoize"
+)
+
+type tokenKey struct {
+	file source.FileIdentity
+}
+
+type tokenHandle struct {
+	handle *memoize.Handle
+	file   source.FileHandle
+}
+
+type tokenData struct {
+	memoize.NoCopy
+
+	tok *token.File
+	err error
+}
+
+func (c *cache) TokenHandle(fh source.FileHandle) source.TokenHandle {
+	key := tokenKey{
+		file: fh.Identity(),
+	}
+	h := c.store.Bind(key, func(ctx context.Context) interface{} {
+		data := &tokenData{}
+		data.tok, data.err = tokenFile(ctx, c, fh)
+		return data
+	})
+	return &tokenHandle{
+		handle: h,
+	}
+}
+
+func (h *tokenHandle) File() source.FileHandle {
+	return h.file
+}
+
+func (h *tokenHandle) Token(ctx context.Context) (*token.File, error) {
+	v := h.handle.Get(ctx)
+	if v == nil {
+		return nil, ctx.Err()
+	}
+	data := v.(*tokenData)
+	return data.tok, data.err
+}
+
+func tokenFile(ctx context.Context, c *cache, fh source.FileHandle) (*token.File, error) {
+	// First, check if we already have a parsed AST for this file's handle.
+	for _, mode := range []source.ParseMode{
+		source.ParseHeader,
+		source.ParseExported,
+		source.ParseFull,
+	} {
+		pk := parseKey{
+			file: fh.Identity(),
+			mode: mode,
+		}
+		pd, ok := c.store.Cached(pk).(*parseGoData)
+		if !ok {
+			continue
+		}
+		if pd.ast == nil {
+			continue
+		}
+		if !pd.ast.Pos().IsValid() {
+			continue
+		}
+		return c.FileSet().File(pd.ast.Pos()), nil
+	}
+	// We have not yet parsed this file.
+	buf, _, err := fh.Read(ctx)
+	if err != nil {
+		return nil, err
+	}
+	tok := c.FileSet().AddFile(fh.Identity().URI.Filename(), -1, len(buf))
+	if tok == nil {
+		return nil, fmt.Errorf("no token.File for %s", fh.Identity().URI)
+	}
+	return tok, nil
+}

internal/lsp/source/view.go

@@ -29,11 +29,12 @@ type FileIdentity struct {
 type FileHandle interface {
 	// FileSystem returns the file system this handle was acquired from.
 	FileSystem() FileSystem
+
 	// Return the Identity for the file.
 	Identity() FileIdentity
-	// Read reads the contents of a file and returns it along with its hash
+
-	// value.
+	// Read reads the contents of a file and returns it along with its hash value.
-	// If the file is not available, retruns a nil slice and an error.
+	// If the file is not available, returns a nil slice and an error.
 	Read(ctx context.Context) ([]byte, string, error)
 }
 
@@ -43,12 +44,23 @@ type FileSystem interface {
 	GetFile(uri span.URI) FileHandle
 }
 
-// ParseGoHandle represents a handle to the ast for a file.
+// TokenHandle represents a handle to the *token.File for a file.
-type ParseGoHandle interface {
+type TokenHandle interface {
-	// File returns a file handle to get the ast for.
+	// File returns a file handle for which to get the *token.File.
 	File() FileHandle
+
+	// Token returns the *token.File for the file.
+	Token(ctx context.Context) (*token.File, error)
+}
+
+// ParseGoHandle represents a handle to the AST for a file.
+type ParseGoHandle interface {
+	// File returns a file handle for which to get the AST.
+	File() FileHandle
+
 	// Mode returns the parse mode of this handle.
 	Mode() ParseMode
+
 	// Parse returns the parsed AST for the file.
 	// If the file is not available, returns nil and an error.
 	Parse(ctx context.Context) (*ast.File, error)
@@ -61,11 +73,13 @@ const (
 	// ParseHeader specifies that the main package declaration and imports are needed.
 	// This is the mode used when attempting to examine the package graph structure.
 	ParseHeader = ParseMode(iota)
+
 	// ParseExported specifies that the public symbols are needed, but things like
 	// private symbols and function bodies are not.
 	// This mode is used for things where a package is being consumed only as a
 	// dependency.
 	ParseExported
+
 	// ParseFull specifies the full AST is needed.
 	// This is used for files of direct interest where the entire contents must
 	// be considered.
@@ -89,8 +103,11 @@ type Cache interface {
 	// FileSet returns the shared fileset used by all files in the system.
 	FileSet() *token.FileSet
 
-	// Parse returns a ParseHandle for the given file handle.
+	// Token returns a TokenHandle for the given file handle.
-	ParseGo(FileHandle, ParseMode) ParseGoHandle
+	TokenHandle(FileHandle) TokenHandle
+
+	// ParseGo returns a ParseGoHandle for the given file handle.
+	ParseGoHandle(FileHandle, ParseMode) ParseGoHandle
 }
 
 // Session represents a single connection from a client.