internal/lsp/cache: move PosTo{Decl,Field} out of cache

Before, these methods of the Source interface used to use
a cache of buildASTCache, which built a Pos-keyed map for
a whole file, but the necessary algorithm is essentially
a binary search which is plenty fast enough to avoid the
need for cache.

This change implements that algorithm and moves
both methods out of the interface into a single function,
source.FindDeclAndField.

--

I measured the duration of all calls to astCacheData (before)
and FindDeclAndField (after) occurring within this command:

  $ go test -bench=TestBenchmarkConfiguredCompletion -v ./gopls/internal/regtest/bench -completion_workdir=$HOME/w/kubernetes -completion_file=../kubernetes/pkg/generated/openapi/zz_generated.openapi.go -completion_regexp=Get

(The numbers reported by this benchmark are problematic,
which is why I measured call times directly; see
https://github.com/golang/go/issues/53798.)

Results:
before (n=4727) max =  21ms, 90% = 4.4ms, median = 19us
after  (n=6282) max = 2.3ms, 90% = 25us,  median = 14us

The increased number of calls to the function after the change
is due to a longstanding bug in the benchmark: each iteration of
the b.N loop doesn't do a fixed amount of work, it does as much
as it can in 10s. Thus making the code faster simply causes
the benchmark to spend the same amount of time on other parts of
the program--such as the loop that calls FindDeclAndField.

See https://go-review.googlesource.com/c/tools/+/221021 for
background on the previous implementation.

Change-Id: I745ecc4e65378fbe97f456228cafba84105b7e49
Reviewed-on: https://go-review.googlesource.com/c/tools/+/416880
Auto-Submit: Alan Donovan <adonovan@google.com>
Run-TryBot: Alan Donovan <adonovan@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>

internal/lsp/cache/parse.go

@@ -26,7 +26,6 @@ import (
 	"golang.org/x/tools/internal/lsp/safetoken"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/memoize"
-	"golang.org/x/tools/internal/span"
 )
 
 // parseKey uniquely identifies a parsed Go file.
@@ -107,171 +106,6 @@ type parseGoResult struct {
 	err    error
 }
 
-type astCacheKey struct {
-	pkg packageHandleKey
-	uri span.URI
-}
-
-func (s *snapshot) astCacheData(ctx context.Context, spkg source.Package, pos token.Pos) (*astCacheData, error) {
-	pkg := spkg.(*pkg)
-	pkgHandle := s.getPackage(pkg.m.ID, pkg.mode)
-	if pkgHandle == nil {
-		return nil, fmt.Errorf("could not reconstruct package handle for %v", pkg.m.ID)
-	}
-	tok := s.FileSet().File(pos)
-	if tok == nil {
-		return nil, fmt.Errorf("no file for pos %v", pos)
-	}
-	pgf, err := pkg.File(span.URIFromPath(tok.Name()))
-	if err != nil {
-		return nil, err
-	}
-
-	// TODO(adonovan): opt: is it necessary to cache this operation?
-	//
-	// I expect the main benefit of CL 221021, which introduced it,
-	// was the replacement of PathEnclosingInterval, whose
-	// traversal is allocation-intensive, by buildASTCache.
-	//
-	// When run on the largest file in k8s, buildASTCache took
-	// ~6ms, but I expect most of that cost could be eliminated by
-	// using a stripped-down version of PathEnclosingInterval that
-	// cares only about syntax trees and not tokens. A stateless
-	// utility function that is cheap enough to call for each Pos
-	// would be a nice simplification.
-	//
-	// (The basic approach would be to use ast.Inspect, compare
-	// each node with the search Pos, and bail out as soon
-	// as a match is found. The pre-order hook would return false
-	// to avoid descending into any tree whose End is before
-	// the search Pos.)
-	//
-	// A representative benchmark would help.
-	astHandle, release := s.store.Handle(astCacheKey{pkgHandle.key, pgf.URI}, func(ctx context.Context, arg interface{}) interface{} {
-		return buildASTCache(pgf)
-	})
-	defer release()
-
-	d, err := s.awaitHandle(ctx, astHandle)
-	if err != nil {
-		return nil, err
-	}
-	return d.(*astCacheData), nil
-}
-
-func (s *snapshot) PosToDecl(ctx context.Context, spkg source.Package, pos token.Pos) (ast.Decl, error) {
-	data, err := s.astCacheData(ctx, spkg, pos)
-	if err != nil {
-		return nil, err
-	}
-	return data.posToDecl[pos], nil
-}
-
-func (s *snapshot) PosToField(ctx context.Context, spkg source.Package, pos token.Pos) (*ast.Field, error) {
-	data, err := s.astCacheData(ctx, spkg, pos)
-	if err != nil {
-		return nil, err
-	}
-	return data.posToField[pos], nil
-}
-
-// An astCacheData maps object positions to syntax nodes for a single Go file.
-type astCacheData struct {
-	// Maps the position of each name declared by a func/var/const/type
-	// Decl to the Decl node.  Also maps the name and type of each field
-	// (broadly defined) to its innermost enclosing Decl.
-	posToDecl map[token.Pos]ast.Decl
-
-	// Maps the position of the Name and Type of each field
-	// (broadly defined) to the Field node.
-	posToField map[token.Pos]*ast.Field
-}
-
-// buildASTCache builds caches to aid in quickly going from the typed
-// world to the syntactic world.
-func buildASTCache(pgf *source.ParsedGoFile) *astCacheData {
-	var (
-		// path contains all ancestors, including n.
-		path []ast.Node
-		// decls contains all ancestors that are decls.
-		decls []ast.Decl
-	)
-
-	data := &astCacheData{
-		posToDecl:  make(map[token.Pos]ast.Decl),
-		posToField: make(map[token.Pos]*ast.Field),
-	}
-
-	ast.Inspect(pgf.File, 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
-}
-
 // parseGoImpl parses the Go source file whose content is provided by fh.
 func parseGoImpl(ctx context.Context, fset *token.FileSet, fh source.FileHandle, mode source.ParseMode) (*source.ParsedGoFile, error) {
 	ctx, done := event.Start(ctx, "cache.parseGo", tag.File.Of(fh.URI().Filename()))

internal/lsp/source/completion/format.go

@@ -242,10 +242,7 @@ Suffixes:
 		return item, nil
 	}
 
-	decl, err := c.snapshot.PosToDecl(ctx, pkg, obj.Pos())
+	decl, _ := source.FindDeclAndField(pkg.GetSyntax(), obj.Pos()) // may be nil
-	if err != nil {
-		return CompletionItem{}, err
-	}
 	hover, err := source.FindHoverContext(ctx, c.snapshot, pkg, obj, decl, nil)
 	if err != nil {
 		event.Error(ctx, "failed to find Hover", err, tag.URI.Of(uri))

internal/lsp/source/hover.go

@@ -610,11 +610,7 @@ func FindHoverContext(ctx context.Context, s Snapshot, pkg Package, obj types.Ob
 				break
 			}
 
-			field, err := s.PosToField(ctx, pkg, obj.Pos())
+			_, field := FindDeclAndField(pkg.GetSyntax(), obj.Pos())
-			if err != nil {
-				return nil, err
-			}
-
 			if field != nil {
 				comment := field.Doc
 				if comment.Text() == "" {
@@ -876,3 +872,99 @@ func anyNonEmpty(x []string) bool {
 	}
 	return false
 }
+
+// FindDeclAndField returns the var/func/type/const Decl that declares
+// the identifier at pos, searching the given list of file syntax
+// trees. If pos is the position of an ast.Field or one of its Names
+// or Ellipsis.Elt, the field is returned, along with the innermost
+// enclosing Decl, which could be only loosely related---consider:
+//
+//    var decl = f(  func(field int) {}  )
+//
+// It returns (nil, nil) if no Field or Decl is found at pos.
+func FindDeclAndField(files []*ast.File, pos token.Pos) (decl ast.Decl, field *ast.Field) {
+	// panic(nil) breaks off the traversal and
+	// causes the function to return normally.
+	defer func() {
+		if x := recover(); x != nil {
+			panic(x)
+		}
+	}()
+
+	// Visit the files in search of the node at pos.
+	var stack []ast.Node
+	for _, file := range files {
+		ast.Inspect(file, func(n ast.Node) bool {
+			if n != nil {
+				stack = append(stack, n) // push
+			} else {
+				stack = stack[:len(stack)-1] // pop
+				return false
+			}
+
+			// Skip subtrees (incl. files) that don't contain the search point.
+			if !(n.Pos() <= pos && pos < n.End()) {
+				return false
+			}
+
+			switch n := n.(type) {
+			case *ast.Field:
+				checkField := func(f ast.Node) {
+					if f.Pos() == pos {
+						field = n
+						for i := len(stack) - 1; i >= 0; i-- {
+							if d, ok := stack[i].(ast.Decl); ok {
+								decl = d // innermost enclosing decl
+								break
+							}
+						}
+						panic(nil) // found
+					}
+				}
+
+				// Check *ast.Field itself. This handles embedded
+				// fields which have no associated *ast.Ident name.
+				checkField(n)
+
+				// Check each field name since you can have
+				// multiple names for the same type expression.
+				for _, name := range n.Names {
+					checkField(name)
+				}
+
+				// Also check "X" in "...X". This makes it easy
+				// to format variadic signature params properly.
+				if ell, ok := n.Type.(*ast.Ellipsis); ok && ell.Elt != nil {
+					checkField(ell.Elt)
+				}
+
+			case *ast.FuncDecl:
+				if n.Name.Pos() == pos {
+					decl = n
+					panic(nil) // found
+				}
+
+			case *ast.GenDecl:
+				for _, spec := range n.Specs {
+					switch spec := spec.(type) {
+					case *ast.TypeSpec:
+						if spec.Name.Pos() == pos {
+							decl = n
+							panic(nil) // found
+						}
+					case *ast.ValueSpec:
+						for _, id := range spec.Names {
+							if id.Pos() == pos {
+								decl = n
+								panic(nil) // found
+							}
+						}
+					}
+				}
+			}
+			return true
+		})
+	}
+
+	return nil, nil
+}

internal/lsp/source/identifier.go

@@ -292,9 +292,8 @@ func findIdentifier(ctx context.Context, snapshot Snapshot, pkg Package, pgf *Pa
 	if err != nil {
 		return nil, err
 	}
-	if result.Declaration.node, err = snapshot.PosToDecl(ctx, declPkg, result.Declaration.obj.Pos()); err != nil {
+	result.Declaration.node, _ = FindDeclAndField(declPkg.GetSyntax(), result.Declaration.obj.Pos()) // may be nil
-		return nil, err
+
-	}
 	// Ensure that we have the full declaration, in case the declaration was
 	// parsed in ParseExported and therefore could be missing information.
 	if result.Declaration.fullDecl, err = fullNode(snapshot, result.Declaration.obj, declPkg); err != nil {

internal/lsp/source/signature_help.go

@@ -98,10 +98,7 @@ FindCall:
 		if err != nil {
 			return nil, 0, err
 		}
-		node, err := snapshot.PosToDecl(ctx, declPkg, obj.Pos())
+		node, _ := FindDeclAndField(declPkg.GetSyntax(), obj.Pos()) // may be nil
-		if err != nil {
-			return nil, 0, err
-		}
 		d, err := FindHoverContext(ctx, snapshot, pkg, obj, node, nil)
 		if err != nil {
 			return nil, 0, err

internal/lsp/source/types_format.go

@@ -259,8 +259,8 @@ func FormatVarType(ctx context.Context, snapshot Snapshot, srcpkg Package, obj *
 		return types.TypeString(obj.Type(), qf)
 	}
 
-	field, err := snapshot.PosToField(ctx, pkg, obj.Pos())
+	_, field := FindDeclAndField(pkg.GetSyntax(), obj.Pos())
-	if err != nil || field == nil {
+	if field == nil {
 		return types.TypeString(obj.Type(), qf)
 	}
 	expr := field.Type

internal/lsp/source/view.go

@@ -79,21 +79,6 @@ type Snapshot interface {
 	// If the file is not available, returns nil and an error.
 	ParseGo(ctx context.Context, fh FileHandle, mode ParseMode) (*ParsedGoFile, error)
 
-	// PosToField is a cache of *ast.Fields by token.Pos. This allows us
-	// to quickly find corresponding *ast.Field node given a *types.Var.
-	// We must refer to the AST to render type aliases properly when
-	// formatting signatures and other types.
-	// May return (nil, nil) if the file didn't declare an object at that position.
-	// TODO(adonovan): seems like a bug?
-	PosToField(ctx context.Context, pkg Package, pos token.Pos) (*ast.Field, error)
-
-	// PosToDecl maps certain objects' positions to their surrounding
-	// ast.Decl. This mapping is used when building the documentation
-	// string for the objects.
-	// May return (nil, nil) if the file didn't declare an object at that position.
-	// TODO(adonovan): seems like a bug?
-	PosToDecl(ctx context.Context, pkg Package, pos token.Pos) (ast.Decl, error)
-
 	// DiagnosePackage returns basic diagnostics, including list, parse, and type errors
 	// for pkg, grouped by file.
 	DiagnosePackage(ctx context.Context, pkg Package) (map[span.URI][]*Diagnostic, error)