internal/refactor/inline: make soleUse asymptotically faster

Fixes golang/go#75773

Change-Id: Ie1e562ef5a9ba01b5d54cd7e69fffcda36f9a2e9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/717540
Auto-Submit: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>

Author: adonovan

internal/refactor/inline/inline.go

@@ -426,12 +426,35 @@ func newImportState(logf func(string, ...any), caller *Caller, callee *gobCallee
 	// For simplicity we ignore existing dot imports, so that a qualified
 	// identifier (QI) in the callee is always represented by a QI in the caller,
 	// allowing us to treat a QI like a selection on a package name.
-	is := &importState{
+	ist := &importState{
 		logf:      logf,
 		caller:    caller,
 		importMap: make(map[string][]string),
 	}
 
+	// Build an index of used-once PkgNames.
+	type pkgNameUse struct {
+		count int
+		id    *ast.Ident // an arbitrary use
+	}
+	pkgNameUses := make(map[*types.PkgName]pkgNameUse)
+	for id, obj := range caller.Info.Uses {
+		if pkgname, ok := obj.(*types.PkgName); ok {
+			u := pkgNameUses[pkgname]
+			u.id = id
+			u.count++
+			pkgNameUses[pkgname] = u
+		}
+	}
+	// soleUse returns the ident that refers to pkgname, if there is exactly one.
+	soleUse := func(pkgname *types.PkgName) *ast.Ident {
+		u := pkgNameUses[pkgname]
+		if u.count == 1 {
+			return u.id
+		}
+		return nil
+	}
+
 	for _, imp := range caller.File.Imports {
 		if pkgName, ok := importedPkgName(caller.Info, imp); ok &&
 			pkgName.Name() != "." &&
@@ -450,7 +473,7 @@ func newImportState(logf func(string, ...any), caller *Caller, callee *gobCallee
 			// need this import. Doing so eagerly simplifies the resulting logic.
 			needed := true
 			sel, ok := ast.Unparen(caller.Call.Fun).(*ast.SelectorExpr)
-			if ok && soleUse(caller.Info, pkgName) == sel.X {
+			if ok && soleUse(pkgName) == sel.X {
 				needed = false // no longer needed by caller
 				// Check to see if any of the inlined free objects need this package.
 				for _, obj := range callee.FreeObjs {
@@ -465,13 +488,13 @@ func newImportState(logf func(string, ...any), caller *Caller, callee *gobCallee
 			// return value holds these.
 			if needed {
 				path := pkgName.Imported().Path()
-				is.importMap[path] = append(is.importMap[path], pkgName.Name())
+				ist.importMap[path] = append(ist.importMap[path], pkgName.Name())
 			} else {
-				is.oldImports = append(is.oldImports, oldImport{pkgName: pkgName, spec: imp})
+				ist.oldImports = append(ist.oldImports, oldImport{pkgName: pkgName, spec: imp})
 			}
 		}
 	}
-	return is
+	return ist
 }
 
 // importName finds an existing import name to use in a particular shadowing
@@ -3736,18 +3759,4 @@ func hasNonTrivialReturn(returnInfo [][]returnOperandFlags) bool {
 	return false
 }
 
-// soleUse returns the ident that refers to obj, if there is exactly one.
-func soleUse(info *types.Info, obj types.Object) (sole *ast.Ident) {
-	// This is not efficient, but it is called infrequently.
-	for id, obj2 := range info.Uses {
-		if obj2 == obj {
-			if sole != nil {
-				return nil // not unique
-			}
-			sole = id
-		}
-	}
-	return sole
-}
-
 type unit struct{} // for representing sets as maps