gopls/internal/cache/parsego: add File.Cursor, and use it

This CL adds a Cursor to the parsego.File. Though the primary
motivation is convenience and flexibility, it is expected
to be an optimization: though it is computed eagerly,
it is retained in the parse cache, and is expected to pay
for itself very quickly by allowing us to replace many
whole-File ast.Inspect operations with more targeted traversals.

The CL replaces all ast.Inspect(file) operations with Cursor,
but there remain many more opportunities for using it in
narrower traversals, and in places that need to navigate to
siblings or ancestors.

Also, amend Cursor.FindPos to use the complete range of the
File, as CL 637738 recently did for astutil.NodeContains.

Also, various clean-ups to InlayHint:
- push the traversals down in InlayHint to avoid having
  to scan a slice for every single node we visit;
- simplify the function signature used for each hint
  algorithm.

Change-Id: I64d0c2cae75fd73a4b539ceb81ad9d6f7d80cfb8
Reviewed-on: https://go-review.googlesource.com/c/tools/+/650396
Reviewed-by: Jonathan Amsterdam <[email protected]>
Auto-Submit: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Commit-Queue: Alan Donovan <[email protected]>
Reviewed-by: Robert Findley <[email protected]>

Author: adonovan

gopls/internal/cache/parsego/file.go

@@ -14,6 +14,7 @@ import (
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/safetoken"
+	"golang.org/x/tools/internal/astutil/cursor"
 )
 
 // A File contains the results of parsing a Go file.
@@ -32,6 +33,8 @@ type File struct {
 	// actual content of the file if we have fixed the AST.
 	Src []byte
 
+	Cursor cursor.Cursor // cursor of *ast.File, sans sibling files
+
 	// fixedSrc and fixedAST report on "fixing" that occurred during parsing of
 	// this file.
 	//
@@ -71,6 +74,11 @@ func (pgf *File) PositionPos(p protocol.Position) (token.Pos, error) {
 	return safetoken.Pos(pgf.Tok, offset)
 }
 
+// PosPosition returns a protocol Position for the token.Pos in this file.
+func (pgf *File) PosPosition(pos token.Pos) (protocol.Position, error) {
+	return pgf.Mapper.PosPosition(pgf.Tok, pos)
+}
+
 // PosRange returns a protocol Range for the token.Pos interval in this file.
 func (pgf *File) PosRange(start, end token.Pos) (protocol.Range, error) {
 	return pgf.Mapper.PosRange(pgf.Tok, start, end)

gopls/internal/cache/parsego/parse.go

@@ -23,11 +23,13 @@ import (
 	"reflect"
 	"slices"
 
+	"golang.org/x/tools/go/ast/inspector"
 	"golang.org/x/tools/gopls/internal/label"
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/util/astutil"
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/safetoken"
+	"golang.org/x/tools/internal/astutil/cursor"
 	"golang.org/x/tools/internal/diff"
 	"golang.org/x/tools/internal/event"
 )
@@ -153,6 +155,11 @@ func Parse(ctx context.Context, fset *token.FileSet, uri protocol.DocumentURI, s
 	}
 	assert(file != nil, "nil *ast.File")
 
+	// Provide a cursor for fast and convenient navigation.
+	inspect := inspector.New([]*ast.File{file})
+	curFile, _ := cursor.Root(inspect).FirstChild()
+	_ = curFile.Node().(*ast.File)
+
 	return &File{
 		URI:      uri,
 		Mode:     mode,
@@ -161,6 +168,7 @@ func Parse(ctx context.Context, fset *token.FileSet, uri protocol.DocumentURI, s
 		fixedAST: fixedAST,
 		File:     file,
 		Tok:      tok,
+		Cursor:   curFile,
 		Mapper:   protocol.NewMapper(uri, src),
 		ParseErr: parseErr,
 	}, fixes

gopls/internal/cache/xrefs/xrefs.go

@@ -44,8 +44,8 @@ func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
 	objectpathFor := new(objectpath.Encoder).For
 
 	for fileIndex, pgf := range files {
-		ast.Inspect(pgf.File, func(n ast.Node) bool {
-			switch n := n.(type) {
+		for cur := range pgf.Cursor.Preorder((*ast.Ident)(nil), (*ast.ImportSpec)(nil)) {
+			switch n := cur.Node().(type) {
 			case *ast.Ident:
 				// Report a reference for each identifier that
 				// uses a symbol exported from another package.
@@ -68,7 +68,7 @@ func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
 							if err != nil {
 								// Capitalized but not exported
 								// (e.g. local const/var/type).
-								return true
+								continue
 							}
 							gobObj = &gobObject{Path: path}
 							objects[obj] = gobObj
@@ -91,7 +91,7 @@ func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
 				// string to the imported package.
 				pkgname := info.PkgNameOf(n)
 				if pkgname == nil {
-					return true // missing import
+					continue // missing import
 				}
 				objects := getObjects(pkgname.Imported())
 				gobObj, ok := objects[nil]
@@ -109,8 +109,7 @@ func Index(files []*parsego.File, pkg *types.Package, info *types.Info) []byte {
 					bug.Reportf("out of bounds import spec %+v", n.Path)
 				}
 			}
-			return true
-		})
+		}
 	}
 
 	// Flatten the maps into slices, and sort for determinism.

gopls/internal/golang/folding_range.go

@@ -46,12 +46,84 @@ func FoldingRange(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle,
 	ranges := commentsFoldingRange(pgf)
 
 	// Walk the ast and collect folding ranges.
-	ast.Inspect(pgf.File, func(n ast.Node) bool {
-		if rng, ok := foldingRangeFunc(pgf, n, lineFoldingOnly); ok {
-			ranges = append(ranges, rng)
+	filter := []ast.Node{
+		(*ast.BasicLit)(nil),
+		(*ast.BlockStmt)(nil),
+		(*ast.CallExpr)(nil),
+		(*ast.CaseClause)(nil),
+		(*ast.CommClause)(nil),
+		(*ast.CompositeLit)(nil),
+		(*ast.FieldList)(nil),
+		(*ast.GenDecl)(nil),
+	}
+	for cur := range pgf.Cursor.Preorder(filter...) {
+		// TODO(suzmue): include trailing empty lines before the closing
+		// parenthesis/brace.
+		var kind protocol.FoldingRangeKind
+		// start and end define the range of content to fold away.
+		var start, end token.Pos
+		switch n := cur.Node().(type) {
+		case *ast.BlockStmt:
+			// Fold between positions of or lines between "{" and "}".
+			start, end = getLineFoldingRange(pgf, n.Lbrace, n.Rbrace, lineFoldingOnly)
+
+		case *ast.CaseClause:
+			// Fold from position of ":" to end.
+			start, end = n.Colon+1, n.End()
+
+		case *ast.CommClause:
+			// Fold from position of ":" to end.
+			start, end = n.Colon+1, n.End()
+
+		case *ast.CallExpr:
+			// Fold between positions of or lines between "(" and ")".
+			start, end = getLineFoldingRange(pgf, n.Lparen, n.Rparen, lineFoldingOnly)
+
+		case *ast.FieldList:
+			// Fold between positions of or lines between opening parenthesis/brace and closing parenthesis/brace.
+			start, end = getLineFoldingRange(pgf, n.Opening, n.Closing, lineFoldingOnly)
+
+		case *ast.GenDecl:
+			// If this is an import declaration, set the kind to be protocol.Imports.
+			if n.Tok == token.IMPORT {
+				kind = protocol.Imports
+			}
+			// Fold between positions of or lines between "(" and ")".
+			start, end = getLineFoldingRange(pgf, n.Lparen, n.Rparen, lineFoldingOnly)
+
+		case *ast.BasicLit:
+			// Fold raw string literals from position of "`" to position of "`".
+			if n.Kind == token.STRING && len(n.Value) >= 2 && n.Value[0] == '`' && n.Value[len(n.Value)-1] == '`' {
+				start, end = n.Pos(), n.End()
+			}
+
+		case *ast.CompositeLit:
+			// Fold between positions of or lines between "{" and "}".
+			start, end = getLineFoldingRange(pgf, n.Lbrace, n.Rbrace, lineFoldingOnly)
+
+		default:
+			panic(n)
 		}
-		return true
-	})
+
+		// Check that folding positions are valid.
+		if !start.IsValid() || !end.IsValid() {
+			continue
+		}
+		if start == end {
+			// Nothing to fold.
+			continue
+		}
+		// in line folding mode, do not fold if the start and end lines are the same.
+		if lineFoldingOnly && safetoken.Line(pgf.Tok, start) == safetoken.Line(pgf.Tok, end) {
+			continue
+		}
+		rng, err := pgf.PosRange(start, end)
+		if err != nil {
+			bug.Reportf("failed to create range: %s", err) // can't happen
+			continue
+		}
+		ranges = append(ranges, foldingRange(kind, rng))
+	}
 
 	// Sort by start position.
 	slices.SortFunc(ranges, func(x, y protocol.FoldingRange) int {
@@ -64,66 +136,6 @@ func FoldingRange(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle,
 	return ranges, nil
 }
 
-// foldingRangeFunc calculates the line folding range for ast.Node n
-func foldingRangeFunc(pgf *parsego.File, n ast.Node, lineFoldingOnly bool) (protocol.FoldingRange, bool) {
-	// TODO(suzmue): include trailing empty lines before the closing
-	// parenthesis/brace.
-	var kind protocol.FoldingRangeKind
-	// start and end define the range of content to fold away.
-	var start, end token.Pos
-	switch n := n.(type) {
-	case *ast.BlockStmt:
-		// Fold between positions of or lines between "{" and "}".
-		start, end = getLineFoldingRange(pgf, n.Lbrace, n.Rbrace, lineFoldingOnly)
-	case *ast.CaseClause:
-		// Fold from position of ":" to end.
-		start, end = n.Colon+1, n.End()
-	case *ast.CommClause:
-		// Fold from position of ":" to end.
-		start, end = n.Colon+1, n.End()
-	case *ast.CallExpr:
-		// Fold between positions of or lines between "(" and ")".
-		start, end = getLineFoldingRange(pgf, n.Lparen, n.Rparen, lineFoldingOnly)
-	case *ast.FieldList:
-		// Fold between positions of or lines between opening parenthesis/brace and closing parenthesis/brace.
-		start, end = getLineFoldingRange(pgf, n.Opening, n.Closing, lineFoldingOnly)
-	case *ast.GenDecl:
-		// If this is an import declaration, set the kind to be protocol.Imports.
-		if n.Tok == token.IMPORT {
-			kind = protocol.Imports
-		}
-		// Fold between positions of or lines between "(" and ")".
-		start, end = getLineFoldingRange(pgf, n.Lparen, n.Rparen, lineFoldingOnly)
-	case *ast.BasicLit:
-		// Fold raw string literals from position of "`" to position of "`".
-		if n.Kind == token.STRING && len(n.Value) >= 2 && n.Value[0] == '`' && n.Value[len(n.Value)-1] == '`' {
-			start, end = n.Pos(), n.End()
-		}
-	case *ast.CompositeLit:
-		// Fold between positions of or lines between "{" and "}".
-		start, end = getLineFoldingRange(pgf, n.Lbrace, n.Rbrace, lineFoldingOnly)
-	}
-
-	// Check that folding positions are valid.
-	if !start.IsValid() || !end.IsValid() {
-		return protocol.FoldingRange{}, false
-	}
-	if start == end {
-		// Nothing to fold.
-		return protocol.FoldingRange{}, false
-	}
-	// in line folding mode, do not fold if the start and end lines are the same.
-	if lineFoldingOnly && safetoken.Line(pgf.Tok, start) == safetoken.Line(pgf.Tok, end) {
-		return protocol.FoldingRange{}, false
-	}
-	rng, err := pgf.PosRange(start, end)
-	if err != nil {
-		bug.Reportf("failed to create range: %s", err) // can't happen
-		return protocol.FoldingRange{}, false
-	}
-	return foldingRange(kind, rng), true
-}
-
 // getLineFoldingRange returns the folding range for nodes with parentheses/braces/brackets
 // that potentially can take up multiple lines.
 func getLineFoldingRange(pgf *parsego.File, open, close token.Pos, lineFoldingOnly bool) (token.Pos, token.Pos) {

gopls/internal/golang/implementation.go

@@ -352,17 +352,14 @@ func localImplementations(ctx context.Context, snapshot *cache.Snapshot, pkg *ca
 	var locs []protocol.Location
 	var methodLocs []methodsets.Location
 	for _, pgf := range pkg.CompiledGoFiles() {
-		ast.Inspect(pgf.File, func(n ast.Node) bool {
-			spec, ok := n.(*ast.TypeSpec)
-			if !ok {
-				return true // not a type declaration
-			}
+		for cur := range pgf.Cursor.Preorder((*ast.TypeSpec)(nil)) {
+			spec := cur.Node().(*ast.TypeSpec)
 			def := pkg.TypesInfo().Defs[spec.Name]
 			if def == nil {
-				return true // "can't happen" for types
+				continue // "can't happen" for types
 			}
 			if def.(*types.TypeName).IsAlias() {
-				return true // skip type aliases to avoid duplicate reporting
+				continue // skip type aliases to avoid duplicate reporting
 			}
 			candidateType := methodsets.EnsurePointer(def.Type())
 
@@ -373,20 +370,20 @@ func localImplementations(ctx context.Context, snapshot *cache.Snapshot, pkg *ca
 			// TODO(adonovan): UX: report I/I pairs too?
 			// The same question appears in the global algorithm (methodsets).
 			if !concreteImplementsIntf(&msets, candidateType, queryType) {
-				return true // not assignable
+				continue // not assignable
 			}
 
 			// Ignore types with empty method sets.
 			// (No point reporting that every type satisfies 'any'.)
 			mset := msets.MethodSet(candidateType)
 			if mset.Len() == 0 {
-				return true
+				continue
 			}
 
 			if method == nil {
 				// Found matching type.
 				locs = append(locs, mustLocation(pgf, spec.Name))
-				return true
+				continue
 			}
 
 			// Find corresponding method.
@@ -407,8 +404,7 @@ func localImplementations(ctx context.Context, snapshot *cache.Snapshot, pkg *ca
 					break
 				}
 			}
-			return true
-		})
+		}
 	}
 
 	// Finally convert method positions to protocol form by reading the files.