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 <[email protected]>
Reviewed-by: Rebecca Stambler <[email protected]>

internal/lsp/cache/check.go

@@ -92,8 +92,10 @@ func (s *snapshot) buildPackageHandle(ctx context.Context, id packageID, mode so
 				dep.check(ctx)
 			}(dep)
 		}
+
 		data := &packageData{}
 		data.pkg, data.err = typeCheck(ctx, fset, m, mode, goFiles, compiledGoFiles, deps)
+
 		return data
 	})
 	ph.handle = h
@@ -413,6 +415,7 @@ func typeCheck(ctx context.Context, fset *token.FileSet, m *metadata, mode sourc
 			}
 		}
 	}
+
 	return pkg, nil
 }
 

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 {

internal/lsp/source/completion_format.go

@@ -187,15 +187,22 @@ func (c *completer) item(ctx context.Context, cand candidate) (CompletionItem, e
 	if cand.imp != nil && cand.imp.pkg != nil {
 		searchPkg = cand.imp.pkg
 	}
-	file, pkg, err := findPosInPackage(c.snapshot.View(), searchPkg, obj.Pos())
+
+	ph, pkg, err := findPosInPackage(c.snapshot.View(), searchPkg, obj.Pos())
 	if err != nil {
 		return item, nil
 	}
-	ident, err := findIdentifier(ctx, c.snapshot, pkg, file, obj.Pos())
+
+	posToDecl, err := ph.PosToDecl(ctx)
 	if err != nil {
+		return CompletionItem{}, err
+	}
+	decl := posToDecl[obj.Pos()]
+	if decl == nil {
 		return item, nil
 	}
-	hover, err := HoverIdentifier(ctx, ident)
+
+	hover, err := hoverInfo(pkg, obj, decl)
 	if err != nil {
 		event.Error(ctx, "failed to find Hover", err, tag.URI.Of(uri))
 		return item, nil

internal/lsp/source/hover.go

@@ -102,10 +102,7 @@ func HoverIdentifier(ctx context.Context, i *IdentifierInfo) (*HoverInformation,
 		h.SingleLine = objectString(obj, i.qf)
 	}
 	h.ImportPath, h.Link, h.SymbolName = pathLinkAndSymbolName(i)
-	if h.comment != nil {
-		h.FullDocumentation = h.comment.Text()
-		h.Synopsis = doc.Synopsis(h.FullDocumentation)
-	}
+
 	return h, nil
 }
 
@@ -217,13 +214,18 @@ func hover(ctx context.Context, fset *token.FileSet, pkg Package, d Declaration)
 	_, done := event.Start(ctx, "source.hover")
 	defer done()
 
-	obj := d.obj
-	switch node := d.node.(type) {
+	return hoverInfo(pkg, d.obj, d.node)
+}
+
+func hoverInfo(pkg Package, obj types.Object, node ast.Node) (*HoverInformation, error) {
+	var info *HoverInformation
+
+	switch node := node.(type) {
 	case *ast.Ident:
 		// The package declaration.
 		for _, f := range pkg.GetSyntax() {
 			if f.Name == node {
-				return &HoverInformation{comment: f.Doc}, nil
+				info = &HoverInformation{comment: f.Doc}
 			}
 		}
 	case *ast.ImportSpec:
@@ -238,32 +240,47 @@ func hover(ctx context.Context, fset *token.FileSet, pkg Package, d Declaration)
 			var doc *ast.CommentGroup
 			for _, file := range imp.GetSyntax() {
 				if file.Doc != nil {
-					return &HoverInformation{source: obj, comment: doc}, nil
+					info = &HoverInformation{source: obj, comment: doc}
 				}
 			}
 		}
-		return &HoverInformation{source: node}, nil
+		info = &HoverInformation{source: node}
 	case *ast.GenDecl:
 		switch obj := obj.(type) {
 		case *types.TypeName, *types.Var, *types.Const, *types.Func:
-			return formatGenDecl(node, obj, obj.Type())
+			var err error
+			info, err = formatGenDecl(node, obj, obj.Type())
+			if err != nil {
+				return nil, err
+			}
 		}
 	case *ast.TypeSpec:
 		if obj.Parent() == types.Universe {
 			if obj.Name() == "error" {
-				return &HoverInformation{source: node}, nil
+				info = &HoverInformation{source: node}
+			} else {
+				info = &HoverInformation{source: node.Name} // comments not needed for builtins
 			}
-			return &HoverInformation{source: node.Name}, nil // comments not needed for builtins
 		}
 	case *ast.FuncDecl:
 		switch obj.(type) {
 		case *types.Func:
-			return &HoverInformation{source: obj, comment: node.Doc}, nil
+			info = &HoverInformation{source: obj, comment: node.Doc}
 		case *types.Builtin:
-			return &HoverInformation{source: node.Type, comment: node.Doc}, nil
+			info = &HoverInformation{source: node.Type, comment: node.Doc}
 		}
 	}
-	return &HoverInformation{source: obj}, nil
+
+	if info == nil {
+		info = &HoverInformation{source: obj}
+	}
+
+	if info.comment != nil {
+		info.FullDocumentation = info.comment.Text()
+		info.Synopsis = doc.Synopsis(info.FullDocumentation)
+	}
+
+	return info, nil
 }
 
 func formatGenDecl(node *ast.GenDecl, obj types.Object, typ types.Type) (*HoverInformation, error) {
@@ -283,6 +300,7 @@ func formatGenDecl(node *ast.GenDecl, obj types.Object, typ types.Type) (*HoverI
 	if spec == nil {
 		return nil, errors.Errorf("no spec for node %v at position %v", node, obj.Pos())
 	}
+
 	// If we have a field or method.
 	switch obj.(type) {
 	case *types.Var, *types.Const, *types.Func:

internal/lsp/source/identifier.go

@@ -12,7 +12,6 @@ import (
 	"go/types"
 	"strconv"
 
-	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/lsp/protocol"
 	errors "golang.org/x/xerrors"
@@ -203,7 +202,7 @@ func findIdentifier(ctx context.Context, s Snapshot, pkg Package, file *ast.File
 	}
 	result.Declaration.MappedRange = append(result.Declaration.MappedRange, rng)
 
-	if result.Declaration.node, err = objToNode(s.View(), pkg, result.Declaration.obj); err != nil {
+	if result.Declaration.node, err = objToDecl(ctx, view, pkg, result.Declaration.obj); err != nil {
 		return nil, err
 	}
 	typ := pkg.GetTypesInfo().TypeOf(result.ident)
@@ -261,31 +260,18 @@ func hasErrorType(obj types.Object) bool {
 	return types.IsInterface(obj.Type()) && obj.Pkg() == nil && obj.Name() == "error"
 }
 
-func objToNode(v View, pkg Package, obj types.Object) (ast.Decl, error) {
-	declAST, _, err := findPosInPackage(v, pkg, obj.Pos())
+func objToDecl(ctx context.Context, v View, srcPkg Package, obj types.Object) (ast.Decl, error) {
+	ph, _, err := findPosInPackage(v, srcPkg, obj.Pos())
 	if err != nil {
 		return nil, err
 	}
-	path, _ := astutil.PathEnclosingInterval(declAST, obj.Pos(), obj.Pos())
-	if path == nil {
-		return nil, errors.Errorf("no path for object %v", obj.Name())
-	}
-	for _, node := range path {
-		switch node := node.(type) {
-		case *ast.GenDecl:
-			// Type names, fields, and methods.
-			switch obj.(type) {
-			case *types.TypeName, *types.Var, *types.Const, *types.Func:
-				return node, nil
-			}
-		case *ast.FuncDecl:
-			// Function signatures.
-			if _, ok := obj.(*types.Func); ok {
-				return node, nil
-			}
-		}
+
+	posToDecl, err := ph.PosToDecl(ctx)
+	if err != nil {
+		return nil, err
 	}
-	return nil, nil // didn't find a node, but don't fail
+
+	return posToDecl[obj.Pos()], nil
 }
 
 // importSpec handles positions inside of an *ast.ImportSpec.

internal/lsp/source/signature_help.go

@@ -99,7 +99,7 @@ FindCall:
 		comment *ast.CommentGroup
 	)
 	if obj != nil {
-		node, err := objToNode(snapshot.View(), pkg, obj)
+		node, err := objToDecl(ctx, snapshot.View(), pkg, obj)
 		if err != nil {
 			return nil, 0, err
 		}

internal/lsp/source/source_test.go

@@ -858,7 +858,7 @@ func (r *runner) SignatureHelp(t *testing.T, spn span.Span, want *protocol.Signa
 		Signatures:      []protocol.SignatureInformation{*gotSignature},
 		ActiveParameter: float64(gotActiveParameter),
 	}
-	if diff := tests.DiffSignatures(spn, got, want); diff != "" {
+	if diff := tests.DiffSignatures(spn, want, got); diff != "" {
 		t.Error(diff)
 	}
 }