go/types/objectpath: memoize scope lookup in objectpath.Encoder

Profiling revealed that scope lookup itself was costly, so memoize the
objects themselves, not just their names.

Also update BenchmarkCompletionFollowingEdit to allow it to be run on
multiple repos, and add a test case for kubernetes.

For golang/go#53992

Change-Id: I17b1f39d8c356e8628610a544306686573a813a7
Reviewed-on: https://go-review.googlesource.com/c/tools/+/502976
gopls-CI: kokoro <[email protected]>
Run-TryBot: Robert Findley <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
TryBot-Bypass: Robert Findley <[email protected]>

go/types/objectpath/objectpath.go

@@ -121,8 +121,8 @@ func For(obj types.Object) (Path, error) {
 // An Encoder amortizes the cost of encoding the paths of multiple objects.
 // The zero value of an Encoder is ready to use.
 type Encoder struct {
-	scopeNamesMemo   map[*types.Scope][]string      // memoization of Scope.Names()
-	namedMethodsMemo map[*types.Named][]*types.Func // memoization of namedMethods()
+	scopeMemo        map[*types.Scope][]types.Object // memoization of scopeObjects
+	namedMethodsMemo map[*types.Named][]*types.Func  // memoization of namedMethods()
 }
 
 // For returns the path to an object relative to its package,
@@ -255,15 +255,14 @@ func (enc *Encoder) For(obj types.Object) (Path, error) {
 	// the best paths because non-types may
 	// refer to types, but not the reverse.
 	empty := make([]byte, 0, 48) // initial space
-	names := enc.scopeNames(scope)
-	for _, name := range names {
-		o := scope.Lookup(name)
+	objs := enc.scopeObjects(scope)
+	for _, o := range objs {
 		tname, ok := o.(*types.TypeName)
 		if !ok {
 			continue // handle non-types in second pass
 		}
 
-		path := append(empty, name...)
+		path := append(empty, o.Name()...)
 		path = append(path, opType)
 
 		T := o.Type()
@@ -289,9 +288,8 @@ func (enc *Encoder) For(obj types.Object) (Path, error) {
 
 	// Then inspect everything else:
 	// non-types, and declared methods of defined types.
-	for _, name := range names {
-		o := scope.Lookup(name)
-		path := append(empty, name...)
+	for _, o := range objs {
+		path := append(empty, o.Name()...)
 		if _, ok := o.(*types.TypeName); !ok {
 			if o.Exported() {
 				// exported non-type (const, var, func)
@@ -746,17 +744,22 @@ func (enc *Encoder) namedMethods(named *types.Named) []*types.Func {
 	return methods
 }
 
-// scopeNames is a memoization of scope.Names. Callers must not modify the result.
-func (enc *Encoder) scopeNames(scope *types.Scope) []string {
-	m := enc.scopeNamesMemo
+// scopeObjects is a memoization of scope objects.
+// Callers must not modify the result.
+func (enc *Encoder) scopeObjects(scope *types.Scope) []types.Object {
+	m := enc.scopeMemo
 	if m == nil {
-		m = make(map[*types.Scope][]string)
-		enc.scopeNamesMemo = m
+		m = make(map[*types.Scope][]types.Object)
+		enc.scopeMemo = m
 	}
-	names, ok := m[scope]
+	objs, ok := m[scope]
 	if !ok {
-		names = scope.Names() // allocates and sorts
-		m[scope] = names
+		names := scope.Names() // allocates and sorts
+		objs = make([]types.Object, len(names))
+		for i, name := range names {
+			objs[i] = scope.Lookup(name)
+		}
+		m[scope] = objs
 	}
-	return names
+	return objs
 }

gopls/internal/regtest/bench/completion_test.go

@@ -6,6 +6,7 @@ package bench
 
 import (
 	"fmt"
+	"sync/atomic"
 	"testing"
 
 	"golang.org/x/tools/gopls/internal/lsp/fake"
@@ -145,31 +146,80 @@ func (c *completer) _() {
 // Edits force type-checked packages to be invalidated, so we want to measure
 // how long it takes before completion results are available.
 func BenchmarkCompletionFollowingEdit(b *testing.B) {
-	file := "internal/lsp/source/completion/completion2.go"
-	fileContent := `
+	tests := []struct {
+		repo           string
+		file           string // repo-relative file to create
+		content        string // file content
+		locationRegexp string // regexp for completion
+	}{
+		{
+			"tools",
+			"internal/lsp/source/completion/completion2.go",
+			`
 package completion
 
 func (c *completer) _() {
 	c.inference.kindMatches(c.)
-	// __MAGIC_STRING_1
 }
-`
-	setup := func(env *Env) {
-		env.CreateBuffer(file, fileContent)
+`,
+			`func \(c \*completer\) _\(\) {\n\tc\.inference\.kindMatches\((c)`,
+		},
+		{
+			"kubernetes",
+			"pkg/kubelet/kubelet2.go",
+			`
+package kubelet
+
+func (kl *Kubelet) _() {
+	kl.
+}
+`,
+			`kl\.()`,
+		},
 	}
 
-	n := 1
-	beforeCompletion := func(env *Env) {
-		old := fmt.Sprintf("__MAGIC_STRING_%d", n)
-		new := fmt.Sprintf("__MAGIC_STRING_%d", n+1)
-		n++
-		env.RegexpReplace(file, old, new)
-	}
+	for _, test := range tests {
+		b.Run(test.repo, func(b *testing.B) {
+			repo := getRepo(b, test.repo)
+			_ = repo.sharedEnv(b) // ensure cache is warm
+			env := repo.newEnv(b, "completion."+test.repo, fake.EditorConfig{
+				Settings: map[string]interface{}{
+					"completeUnimported": false,
+				},
+			})
+			defer env.Close()
+
+			env.CreateBuffer(test.file, "// __REGTEST_PLACEHOLDER_0__\n"+test.content)
+			editPlaceholder := func() {
+				edits := atomic.AddInt64(&editID, 1)
+				env.EditBuffer(test.file, protocol.TextEdit{
+					Range: protocol.Range{
+						Start: protocol.Position{Line: 0, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 0},
+					},
+					// Increment the placeholder text, to ensure cache misses.
+					NewText: fmt.Sprintf("// __REGTEST_PLACEHOLDER_%d__\n", edits),
+				})
+			}
+			env.AfterChange()
 
-	benchmarkCompletion(completionBenchOptions{
-		file:             file,
-		locationRegexp:   `func \(c \*completer\) _\(\) {\n\tc\.inference\.kindMatches\((c)`,
-		setup:            setup,
-		beforeCompletion: beforeCompletion,
-	}, b)
+			// Run a completion to make sure the system is warm.
+			loc := env.RegexpSearch(test.file, test.locationRegexp)
+			completions := env.Completion(loc)
+
+			if testing.Verbose() {
+				fmt.Println("Results:")
+				for i := 0; i < len(completions.Items); i++ {
+					fmt.Printf("\t%d. %v\n", i, completions.Items[i])
+				}
+			}
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				editPlaceholder()
+				loc := env.RegexpSearch(test.file, test.locationRegexp)
+				env.Completion(loc)
+			}
+		})
+	}
 }