internal/imports: optimize package path filtering

As reported in golang/go#67889, a short unresolved identifier can cause
goimports fixes to be quite slow, because the substring match heuristic
used to filter import paths matches a large fraction of the module
cache.

Fix this by improving the precision of the matching heuristic used to
filter potential packages. We now match only full segments of the import
path (ignoring '.' and '-'), or subsegments delimited by '.' and '-'.

Add a gopls benchmark that reproduces this initial slowness, along with
a command to force a scan of the module cache. On my (overpowered) linux
development machine, with a 5GB module cache, this change reduces the
benchmark time ~90%, from 2s to 200ms. With a smaller machine, slower
operating system, or larger module cache, I imagine the starting
point could be significantly more than 2s.

Fixes golang/go#67889

Change-Id: Id8f7ea20040b059b42366333adeb4add684dee61
Reviewed-on: https://go-review.googlesource.com/c/tools/+/591715
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>

Author: findleyr

gopls/doc/commands.md

@@ -568,6 +568,10 @@ Args:
 }
 ```
 
+## `gopls.scan_imports`: **force a sychronous scan of the imports cache.**
+
+This command is intended for use by gopls tests only.
+
 ## `gopls.start_debugging`: **Start the gopls debug server**
 
 Start the gopls debug server if it isn't running, and return the debug

gopls/internal/cache/view.go

@@ -14,6 +14,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"log"
 	"os"
 	"os/exec"
 	"path"
@@ -481,6 +482,14 @@ func (v *View) shutdown() {
 	v.snapshotMu.Unlock()
 }
 
+// ScanImports scans the module cache synchronously.
+// For use in tests.
+func (v *View) ScanImports() {
+	gomodcache := v.folder.Env.GOMODCACHE
+	dirCache := v.importsState.modCache.dirCache(gomodcache)
+	imports.ScanModuleCache(gomodcache, dirCache, log.Printf)
+}
+
 // IgnoredFile reports if a file would be ignored by a `go list` of the whole
 // workspace.
 //

gopls/internal/doc/api.json

@@ -1107,6 +1107,13 @@
 			"ArgDoc": "{\n\t// The test file containing the tests to run.\n\t\"URI\": string,\n\t// Specific test names to run, e.g. TestFoo.\n\t\"Tests\": []string,\n\t// Specific benchmarks to run, e.g. BenchmarkFoo.\n\t\"Benchmarks\": []string,\n}",
 			"ResultDoc": ""
 		},
+		{
+			"Command": "gopls.scan_imports",
+			"Title": "force a sychronous scan of the imports cache.",
+			"Doc": "This command is intended for use by gopls tests only.",
+			"ArgDoc": "",
+			"ResultDoc": ""
+		},
 		{
 			"Command": "gopls.start_debugging",
 			"Title": "Start the gopls debug server",

gopls/internal/protocol/command/command_gen.go

@@ -259,6 +259,11 @@ type Interface interface {
 	// specified function symbol (plus any nested lambdas and defers).
 	// The machine architecture is determined by the view.
 	Assembly(_ context.Context, viewID, packageID, symbol string) error
+
+	// ScanImports: force a sychronous scan of the imports cache.
+	//
+	// This command is intended for use by gopls tests only.
+	ScanImports(context.Context) error
 }
 
 type RunTestsArgs struct {

gopls/internal/protocol/command/interface.go

@@ -1484,3 +1484,10 @@ func (c *commandHandler) Assembly(ctx context.Context, viewID, packageID, symbol
 	openClientBrowser(ctx, c.s.client, url)
 	return nil
 }
+
+func (c *commandHandler) ScanImports(ctx context.Context) error {
+	for _, v := range c.s.session.Views() {
+		v.ScanImports()
+	}
+	return nil
+}

gopls/internal/server/command.go

@@ -0,0 +1,89 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package bench
+
+import (
+	"context"
+	"flag"
+	"testing"
+	"time"
+
+	"golang.org/x/tools/gopls/internal/protocol"
+	"golang.org/x/tools/gopls/internal/protocol/command"
+	. "golang.org/x/tools/gopls/internal/test/integration"
+	"golang.org/x/tools/gopls/internal/test/integration/fake"
+)
+
+var gopath = flag.String("gopath", "", "if set, run goimports scan with this GOPATH value")
+
+func BenchmarkInitialGoimportsScan(b *testing.B) {
+	if *gopath == "" {
+		// This test doesn't make much sense with a tiny module cache.
+		// For now, don't bother trying to construct a huge cache, since it likely
+		// wouldn't work well on the perf builder. Instead, this benchmark only
+		// runs with a pre-existing GOPATH.
+		b.Skip("imports scan requires an explicit GOPATH to be set with -gopath")
+	}
+
+	repo := getRepo(b, "tools") // since this a test of module cache scanning, any repo will do
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		func() {
+			// Unfortunately we (intentionally) don't support resetting the module
+			// cache scan state, so in order to have an accurate benchmark we must
+			// effectively restart gopls on every iteration.
+			//
+			// Warning: this can cause this benchmark to run quite slowly if the
+			// observed time (when the timer is running) is a tiny fraction of the
+			// actual time.
+			b.StopTimer()
+			config := fake.EditorConfig{
+				Env: map[string]string{"GOPATH": *gopath},
+			}
+			env := repo.newEnv(b, config, "imports", false)
+			defer env.Close()
+			env.Await(InitialWorkspaceLoad)
+
+			// Create a buffer with a dangling selctor where the receiver is a single
+			// character ('a') that matches a large fraction of the module cache.
+			env.CreateBuffer("internal/lsp/cache/temp.go", `
+// This is a temp file to exercise goimports scan of the module cache.
+package cache
+
+func _() {
+	_ = a.B // a dangling selector causes goimports to scan many packages
+}
+`)
+			env.AfterChange()
+
+			// Force a scan of the imports cache, so that the goimports algorithm
+			// observes all directories.
+			env.ExecuteCommand(&protocol.ExecuteCommandParams{
+				Command: command.ScanImports.String(),
+			}, nil)
+
+			if stopAndRecord := startProfileIfSupported(b, env, "importsscan"); stopAndRecord != nil {
+				defer stopAndRecord()
+			}
+
+			b.StartTimer()
+			if false {
+				// golang/go#67923: testing resuming imports scanning after a
+				// cancellation.
+				//
+				// Cancelling and then resuming the scan should take around the same
+				// amount of time.
+				ctx, cancel := context.WithTimeout(env.Ctx, 50*time.Millisecond)
+				defer cancel()
+				if err := env.Editor.OrganizeImports(ctx, "internal/lsp/cache/temp.go"); err != nil {
+					b.Logf("organize imports failed: %v", err)
+				}
+			}
+			env.OrganizeImports("internal/lsp/cache/temp.go")
+		}()
+	}
+}

gopls/internal/test/integration/bench/imports_test.go

@@ -46,7 +46,7 @@ func main() {
 			t.Fatal(err)
 		}
 		env.ExecuteCommand(&protocol.ExecuteCommandParams{
-			Command:   "gopls.add_import",
+			Command:   command.AddImport.String(),
 			Arguments: cmd.Arguments,
 		}, nil)
 		got := env.BufferText("main.go")

gopls/internal/test/integration/misc/import_test.go

@@ -1776,58 +1776,24 @@ func pkgIsCandidate(filename string, refs references, pkg *pkg) bool {
 	}
 
 	// Speed optimization to minimize disk I/O:
-	// the last two components on disk must contain the
-	// package name somewhere.
 	//
-	// This permits mismatch naming like directory
-	// "go-foo" being package "foo", or "pkg.v3" being "pkg",
-	// or directory "google.golang.org/api/cloudbilling/v1"
-	// being package "cloudbilling", but doesn't
-	// permit a directory "foo" to be package
-	// "bar", which is strongly discouraged
-	// anyway. There's no reason goimports needs
-	// to be slow just to accommodate that.
+	// Use the matchesPath heuristic to filter to package paths that could
+	// reasonably match a dangling reference.
+	//
+	// This permits mismatch naming like directory "go-foo" being package "foo",
+	// or "pkg.v3" being "pkg", or directory
+	// "google.golang.org/api/cloudbilling/v1" being package "cloudbilling", but
+	// doesn't permit a directory "foo" to be package "bar", which is strongly
+	// discouraged anyway. There's no reason goimports needs to be slow just to
+	// accommodate that.
 	for pkgIdent := range refs {
-		lastTwo := lastTwoComponents(pkg.importPathShort)
-		if strings.Contains(lastTwo, pkgIdent) {
-			return true
-		}
-		if hasHyphenOrUpperASCII(lastTwo) && !hasHyphenOrUpperASCII(pkgIdent) {
-			lastTwo = lowerASCIIAndRemoveHyphen(lastTwo)
-			if strings.Contains(lastTwo, pkgIdent) {
-				return true
-			}
-		}
-	}
-	return false
-}
-
-func hasHyphenOrUpperASCII(s string) bool {
-	for i := 0; i < len(s); i++ {
-		b := s[i]
-		if b == '-' || ('A' <= b && b <= 'Z') {
+		if matchesPath(pkgIdent, pkg.importPathShort) {
 			return true
 		}
 	}
 	return false
 }
 
-func lowerASCIIAndRemoveHyphen(s string) (ret string) {
-	buf := make([]byte, 0, len(s))
-	for i := 0; i < len(s); i++ {
-		b := s[i]
-		switch {
-		case b == '-':
-			continue
-		case 'A' <= b && b <= 'Z':
-			buf = append(buf, b+('a'-'A'))
-		default:
-			buf = append(buf, b)
-		}
-	}
-	return string(buf)
-}
-
 // canUse reports whether the package in dir is usable from filename,
 // respecting the Go "internal" and "vendor" visibility rules.
 func canUse(filename, dir string) bool {
@@ -1868,19 +1834,84 @@ func canUse(filename, dir string) bool {
 	return !strings.Contains(relSlash, "/vendor/") && !strings.Contains(relSlash, "/internal/") && !strings.HasSuffix(relSlash, "/internal")
 }
 
-// lastTwoComponents returns at most the last two path components
-// of v, using either / or \ as the path separator.
-func lastTwoComponents(v string) string {
+// matchesPath reports whether ident may match a potential package name
+// referred to by path, using heuristics to filter out unidiomatic package
+// names.
+//
+// Specifically, it checks whether either of the last two '/'- or '\'-delimited
+// path segments matches the identifier. The segment-matching heuristic must
+// allow for various conventions around segment naming, including go-foo,
+// foo-go, and foo.v3. To handle all of these, matching considers both (1) the
+// entire segment, ignoring '-' and '.', as well as (2) the last subsegment
+// separated by '-' or '.'. So the segment foo-go matches all of the following
+// identifiers: foo, go, and foogo. All matches are case insensitive (for ASCII
+// identifiers).
+//
+// See the docstring for [pkgIsCandidate] for an explanation of how this
+// heuristic filters potential candidate packages.
+func matchesPath(ident, path string) bool {
+	// Ignore case, for ASCII.
+	lowerIfASCII := func(b byte) byte {
+		if 'A' <= b && b <= 'Z' {
+			return b + ('a' - 'A')
+		}
+		return b
+	}
+
+	// match reports whether path[start:end] matches ident, ignoring [.-].
+	match := func(start, end int) bool {
+		ii := len(ident) - 1 // current byte in ident
+		pi := end - 1        // current byte in path
+		for ; pi >= start && ii >= 0; pi-- {
+			pb := path[pi]
+			if pb == '-' || pb == '.' {
+				continue
+			}
+			pb = lowerIfASCII(pb)
+			ib := lowerIfASCII(ident[ii])
+			if pb != ib {
+				return false
+			}
+			ii--
+		}
+		return ii < 0 && pi < start // all bytes matched
+	}
+
+	// segmentEnd and subsegmentEnd hold the end points of the current segment
+	// and subsegment intervals.
+	segmentEnd := len(path)
+	subsegmentEnd := len(path)
+
+	// Count slashes; we only care about the last two segments.
 	nslash := 0
-	for i := len(v) - 1; i >= 0; i-- {
-		if v[i] == '/' || v[i] == '\\' {
+
+	for i := len(path) - 1; i >= 0; i-- {
+		switch b := path[i]; b {
+		// TODO(rfindley): we handle backlashes here only because the previous
+		// heuristic handled backslashes. This is perhaps overly defensive, but is
+		// the result of many lessons regarding Chesterton's fence and the
+		// goimports codebase.
+		//
+		// However, this function is only ever called with something called an
+		// 'importPath'. Is it possible that this is a real import path, and
+		// therefore we need only consider forward slashes?
+		case '/', '\\':
+			if match(i+1, segmentEnd) || match(i+1, subsegmentEnd) {
+				return true
+			}
 			nslash++
 			if nslash == 2 {
-				return v[i:]
+				return false // did not match above
+			}
+			segmentEnd, subsegmentEnd = i, i // reset
+		case '-', '.':
+			if match(i+1, subsegmentEnd) {
+				return true
 			}
+			subsegmentEnd = i
 		}
 	}
-	return v
+	return match(0, segmentEnd) || match(0, subsegmentEnd)
 }
 
 type visitFn func(node ast.Node) ast.Visitor