internal/lsp/source: speed up completion candidate formatting

Completion could be slow due to calls to astutil.PathEnclosingInterval
for every candidate during formatting. There were two reasons we
called PEI:

1. To properly render type alias names, we must refer to the AST
   because the alias name is not available in the typed world.
   Previously we would call PEI to find the *type.Var's
   corresponding *ast.Field, but now we have a PosToField cache that
   lets us jump straight from the types.Object's token.Pos to the
   corresponding *ast.Field.

2. To display an object's documentation we must refer to the AST. We
   need the object's declaring node and any containing ast.Decl. We
   now maintain a special PosToDecl cache so we can avoid the PEI call
   in this case as well.

We can't use a single cache for both because the *ast.Field's position
is present in both caches (but points to different nodes). The caches
are memoized to defer generation until they are needed and to save
work creating them if the *ast.Files haven't changed.

These changes speed up completing the fields of
github.com/aws/aws-sdk-go/service/ec2 from 18.5s to 45ms on my laptop.

Fixes golang/go#37450.

Change-Id: I25cc5ea39551db728a2348f346342ebebeddd049
Reviewed-on: https://go-review.googlesource.com/c/tools/+/221021
Run-TryBot: Muir Manders <muir@mnd.rs>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

1 files changed, 139 insertions, 3 deletions

diff --git a/internal/lsp/cache/parse.go b/internal/lsp/cache/parse.go
index a3b1427a8..75d8888cf 100644
--- a/internal/lsp/cache/parse.go
+++ b/internal/lsp/cache/parse.go
@@ -28,10 +28,15 @@ type parseKey struct {
 	mode source.ParseMode
 }
 
+// astCacheKey is similar to parseKey, but is a distinct type because
+// it is used to key a different value within the same map.
+type astCacheKey parseKey
+
 type parseGoHandle struct {
-	handle *memoize.Handle
-	file   source.FileHandle
-	mode   source.ParseMode
+	handle         *memoize.Handle
+	file           source.FileHandle
+	mode           source.ParseMode
+	astCacheHandle *memoize.Handle
 }
 
 type parseGoData struct {
@@ -63,10 +68,14 @@ func (c *Cache) parseGoHandle(ctx context.Context, fh source.FileHandle, mode so
 	h := c.store.Bind(key, func(ctx context.Context) interface{} {
 		return parseGo(ctx, fset, fh, mode)
 	})
+
 	return &parseGoHandle{
 		handle: h,
 		file:   fh,
 		mode:   mode,
+		astCacheHandle: c.store.Bind(astCacheKey(key), func(ctx context.Context) interface{} {
+			return buildASTCache(ctx, h)
+		}),
 	}
 }
 
@@ -111,6 +120,133 @@ func (pgh *parseGoHandle) Cached() (*ast.File, []byte, *protocol.ColumnMapper, e
 	return data.ast, data.src, data.mapper, data.parseError, data.err
 }
 
+func (pgh *parseGoHandle) PosToDecl(ctx context.Context) (map[token.Pos]ast.Decl, error) {
+	v, err := pgh.astCacheHandle.Get(ctx)
+	if err != nil || v == nil {
+		return nil, err
+	}
+
+	data := v.(*astCacheData)
+	if data.err != nil {
+		return nil, data.err
+	}
+
+	return data.posToDecl, nil
+}
+
+func (pgh *parseGoHandle) PosToField(ctx context.Context) (map[token.Pos]*ast.Field, error) {
+	v, err := pgh.astCacheHandle.Get(ctx)
+	if err != nil || v == nil {
+		return nil, err
+	}
+
+	data := v.(*astCacheData)
+	if data.err != nil {
+		return nil, data.err
+	}
+
+	return data.posToField, nil
+}
+
+type astCacheData struct {
+	memoize.NoCopy
+
+	err error
+
+	posToDecl  map[token.Pos]ast.Decl
+	posToField map[token.Pos]*ast.Field
+}
+
+// buildASTCache builds caches to aid in quickly going from the typed
+// world to the syntactic world.
+func buildASTCache(ctx context.Context, parseHandle *memoize.Handle) *astCacheData {
+	var (
+		// path contains all ancestors, including n.
+		path []ast.Node
+		// decls contains all ancestors that are decls.
+		decls []ast.Decl
+	)
+
+	v, err := parseHandle.Get(ctx)
+	if err != nil || v == nil || v.(*parseGoData).ast == nil {
+		return &astCacheData{err: err}
+	}
+
+	data := &astCacheData{
+		posToDecl:  make(map[token.Pos]ast.Decl),
+		posToField: make(map[token.Pos]*ast.Field),
+	}
+
+	ast.Inspect(v.(*parseGoData).ast, func(n ast.Node) bool {
+		if n == nil {
+			lastP := path[len(path)-1]
+			path = path[:len(path)-1]
+			if len(decls) > 0 && decls[len(decls)-1] == lastP {
+				decls = decls[:len(decls)-1]
+			}
+			return false
+		}
+
+		path = append(path, n)
+
+		switch n := n.(type) {
+		case *ast.Field:
+			addField := func(f ast.Node) {
+				if f.Pos().IsValid() {
+					data.posToField[f.Pos()] = n
+					if len(decls) > 0 {
+						data.posToDecl[f.Pos()] = decls[len(decls)-1]
+					}
+				}
+			}
+
+			// Add mapping for *ast.Field itself. This handles embedded
+			// fields which have no associated *ast.Ident name.
+			addField(n)
+
+			// Add mapping for each field name since you can have
+			// multiple names for the same type expression.
+			for _, name := range n.Names {
+				addField(name)
+			}
+
+			// Also map "X" in "...X" to the containing *ast.Field. This
+			// makes it easy to format variadic signature params
+			// properly.
+			if elips, ok := n.Type.(*ast.Ellipsis); ok && elips.Elt != nil {
+				addField(elips.Elt)
+			}
+		case *ast.FuncDecl:
+			decls = append(decls, n)
+
+			if n.Name != nil && n.Name.Pos().IsValid() {
+				data.posToDecl[n.Name.Pos()] = n
+			}
+		case *ast.GenDecl:
+			decls = append(decls, n)
+
+			for _, spec := range n.Specs {
+				switch spec := spec.(type) {
+				case *ast.TypeSpec:
+					if spec.Name != nil && spec.Name.Pos().IsValid() {
+						data.posToDecl[spec.Name.Pos()] = n
+					}
+				case *ast.ValueSpec:
+					for _, id := range spec.Names {
+						if id != nil && id.Pos().IsValid() {
+							data.posToDecl[id.Pos()] = n
+						}
+					}
+				}
+			}
+		}
+
+		return true
+	})
+
+	return data
+}
+
 func hashParseKeys(pghs []*parseGoHandle) string {
 	b := bytes.NewBuffer(nil)
 	for _, pgh := range pghs {