imports: prefer math/rand/v2 over math/rand

When a program refers to 'rand.Float64()' without an import statement,
import math/rand/v2 is the better answer.

Fixes: golang/go#66407
Change-Id: I92709c335b61aa10dd49ce2e5336ea233d8f9af8
Reviewed-on: https://go-review.googlesource.com/c/tools/+/573075
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: pjweinbgo

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

@@ -72,7 +72,7 @@ func a() {
 	Run(t, stuff, func(t *testing.T, env *Env) {
 		env.OpenFile("a.go")
 		was := env.BufferText("a.go")
-		env.Await(NoDiagnostics())
+		env.AfterChange(NoDiagnostics())
 		env.OrganizeImports("a.go")
 		is := env.BufferText("a.go")
 		if diff := compare.Text(was, is); diff != "" {
@@ -81,6 +81,51 @@ func a() {
 	})
 }
 
+func TestIssue66407(t *testing.T) {
+	const files = `
+-- go.mod --
+module foo
+go 1.21
+-- a.go --
+package foo
+
+func f(x float64) float64 {
+	return x +  rand.Float64()
+}
+-- b.go --
+package foo
+
+func g() {
+	_ = rand.Int63()
+}
+`
+	WithOptions(Modes(Default)).
+		Run(t, files, func(t *testing.T, env *Env) {
+			env.OpenFile("a.go")
+			was := env.BufferText("a.go")
+			env.OrganizeImports("a.go")
+			is := env.BufferText("a.go")
+			// expect complaint that module is before 1.22
+			env.AfterChange(Diagnostics(ForFile("a.go")))
+			diff := compare.Text(was, is)
+			// check that it found the 'right' rand
+			if !strings.Contains(diff, `import "math/rand/v2"`) {
+				t.Errorf("expected rand/v2, got %q", diff)
+			}
+			env.OpenFile("b.go")
+			was = env.BufferText("b.go")
+			env.OrganizeImports("b.go")
+			// a.go still has its module problem but b.go is fine
+			env.AfterChange(Diagnostics(ForFile("a.go")),
+				NoDiagnostics(ForFile("b.go")))
+			is = env.BufferText("b.go")
+			diff = compare.Text(was, is)
+			if !strings.Contains(diff, `import "math/rand"`) {
+				t.Errorf("expected math/rand, got %q", diff)
+			}
+		})
+}
+
 func TestVim1(t *testing.T) {
 	const vim1 = `package main
 

internal/imports/fix.go

@@ -301,6 +301,20 @@ func (p *pass) loadPackageNames(imports []*ImportInfo) error {
 	return nil
 }
 
+// if there is a trailing major version, remove it
+func withoutVersion(nm string) string {
+	if v := path.Base(nm); len(v) > 0 && v[0] == 'v' {
+		if _, err := strconv.Atoi(v[1:]); err == nil {
+			// this is, for instance, called with rand/v2 and returns rand
+			if len(v) < len(nm) {
+				xnm := nm[:len(nm)-len(v)-1]
+				return path.Base(xnm)
+			}
+		}
+	}
+	return nm
+}
+
 // importIdentifier returns the identifier that imp will introduce. It will
 // guess if the package name has not been loaded, e.g. because the source
 // is not available.
@@ -310,7 +324,7 @@ func (p *pass) importIdentifier(imp *ImportInfo) string {
 	}
 	known := p.knownPackages[imp.ImportPath]
 	if known != nil && known.name != "" {
-		return known.name
+		return withoutVersion(known.name)
 	}
 	return ImportPathToAssumedName(imp.ImportPath)
 }
@@ -1059,6 +1073,18 @@ func addStdlibCandidates(pass *pass, refs references) error {
 	if err != nil {
 		return err
 	}
+	localbase := func(nm string) string {
+		ans := path.Base(nm)
+		if ans[0] == 'v' {
+			// this is called, for instance, with math/rand/v2 and returns rand/v2
+			if _, err := strconv.Atoi(ans[1:]); err == nil {
+				ix := strings.LastIndex(nm, ans)
+				more := path.Base(nm[:ix])
+				ans = path.Join(more, ans)
+			}
+		}
+		return ans
+	}
 	add := func(pkg string) {
 		// Prevent self-imports.
 		if path.Base(pkg) == pass.f.Name.Name && filepath.Join(goenv["GOROOT"], "src", pkg) == pass.srcDir {
@@ -1067,13 +1093,17 @@ func addStdlibCandidates(pass *pass, refs references) error {
 		exports := symbolNameSet(stdlib.PackageSymbols[pkg])
 		pass.addCandidate(
 			&ImportInfo{ImportPath: pkg},
-			&packageInfo{name: path.Base(pkg), exports: exports})
+			&packageInfo{name: localbase(pkg), exports: exports})
 	}
 	for left := range refs {
 		if left == "rand" {
-			// Make sure we try crypto/rand before math/rand.
+			// Make sure we try crypto/rand before any version of math/rand as both have Int()
+			// and our policy is to recommend crypto
 			add("crypto/rand")
-			add("math/rand")
+			// if the user's no later than go1.21, this should be "math/rand"
+			// but we have no way of figuring out what the user is using
+			// TODO: investigate using the toolchain version to disambiguate in the stdlib
+			add("math/rand/v2")
 			continue
 		}
 		for importPath := range stdlib.PackageSymbols {

internal/imports/fix_test.go

@@ -1151,7 +1151,7 @@ var _ = rand.NewZipf
 `,
 		out: `package p
 
-import "math/rand"
+import "math/rand/v2"
 
 var _ = rand.NewZipf
 `,