internal/imports: sort import candidates by "relevance"

When proposing packages to import, we can propose more relevant packages
first. Introduce that concept to the pkg struct, and sort by it when
returning candidates.

In all cases we prefer stdlib packages first. Then, in module mode, we
prefer packages that are in the module's dependencies over those that
aren't. We could go further and prefer direct deps over indirect too,
but I didn't have the code for that handy.

I also changed the alphabetical sort from import path to package name,
because that's what the user sees first in the UI.

Updates golang/go#31906

Change-Id: Ia981ee9ffe3202e2a68eef3a36f65e81849a4ac2
Reviewed-on: https://go-review.googlesource.com/c/tools/+/204203
Run-TryBot: Heschi Kreinick <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/imports/fix.go b/internal/imports/fix.go
index 3e6e56d..72b43bd 100644
--- a/internal/imports/fix.go
+++ b/internal/imports/fix.go
@@ -589,15 +589,6 @@
 func getAllCandidates(filename string, env *ProcessEnv) ([]ImportFix, error) {
 	// TODO(heschi): filter out current package. (Don't forget x_test can import x.)
 
-	// Exclude goroot results -- getting them is relatively expensive, not cached,
-	// and generally redundant with the in-memory version.
-	exclude := []gopathwalk.RootType{gopathwalk.RootGOROOT}
-	// Only the go/packages resolver uses the first argument, and nobody uses that resolver.
-	pkgs, err := env.GetResolver().scan(nil, true, exclude)
-	if err != nil {
-		return nil, err
-	}
-
 	// Start off with the standard library.
 	var imports []ImportFix
 	for importPath := range stdlib {
@@ -609,6 +600,31 @@
 			FixType:   AddImport,
 		})
 	}
+	// Sort the stdlib bits solely by name.
+	sort.Slice(imports, func(i int, j int) bool {
+		return imports[i].StmtInfo.ImportPath < imports[j].StmtInfo.ImportPath
+	})
+
+	// Exclude goroot results -- getting them is relatively expensive, not cached,
+	// and generally redundant with the in-memory version.
+	exclude := []gopathwalk.RootType{gopathwalk.RootGOROOT}
+	// Only the go/packages resolver uses the first argument, and nobody uses that resolver.
+	pkgs, err := env.GetResolver().scan(nil, true, exclude)
+	if err != nil {
+		return nil, err
+	}
+	// Sort first by relevance, then by name, so that when we add them they're
+	// still in order.
+	sort.Slice(pkgs, func(i, j int) bool {
+		pi, pj := pkgs[i], pkgs[j]
+		if pi.relevance < pj.relevance {
+			return true
+		}
+		if pi.relevance > pj.relevance {
+			return false
+		}
+		return pi.packageName < pj.packageName
+	})
 
 	dupCheck := map[string]struct{}{}
 	for _, pkg := range pkgs {
@@ -627,9 +643,6 @@
 			FixType:   AddImport,
 		})
 	}
-	sort.Slice(imports, func(i int, j int) bool {
-		return imports[i].StmtInfo.ImportPath < imports[j].StmtInfo.ImportPath
-	})
 	return imports, nil
 }
 
@@ -1036,6 +1049,7 @@
 	dir             string // absolute file path to pkg directory ("/usr/lib/go/src/net/http")
 	importPathShort string // vendorless import path ("net/http", "a/b")
 	packageName     string // package name loaded from source if requested
+	relevance       int    // a weakly-defined score of how relevant a package is. 0 is most relevant.
 }
 
 type pkgDistance struct {
@@ -1097,6 +1111,10 @@
 		p := &pkg{
 			importPathShort: VendorlessPath(importpath),
 			dir:             dir,
+			relevance:       1,
+		}
+		if root.Type == gopathwalk.RootGOROOT {
+			p.relevance = 0
 		}
 		if loadNames {
 			var err error
diff --git a/internal/imports/fix_test.go b/internal/imports/fix_test.go
index 6dafd40..a29fc6e 100644
--- a/internal/imports/fix_test.go
+++ b/internal/imports/fix_test.go
@@ -2521,12 +2521,12 @@
 		name, path string
 	}
 	want := []res{
-		{"bar", "bar.com/bar"},
 		{"bytes", "bytes"},
 		{"rand", "crypto/rand"},
-		{"foo", "foo.com/foo"},
 		{"rand", "math/rand"},
 		{"http", "net/http"},
+		{"bar", "bar.com/bar"},
+		{"foo", "foo.com/foo"},
 	}
 
 	testConfig{
diff --git a/internal/imports/mod.go b/internal/imports/mod.go
index 2d394ab..333dac4 100644
--- a/internal/imports/mod.go
+++ b/internal/imports/mod.go
@@ -430,12 +430,15 @@
 			importPathShort: info.nonCanonicalImportPath,
 			dir:             info.dir,
 			packageName:     path.Base(info.nonCanonicalImportPath),
+			relevance:       0,
 		}, nil
 	}
 
 	importPath := info.nonCanonicalImportPath
+	relevance := 2
 	// Check if the directory is underneath a module that's in scope.
 	if mod := r.findModuleByDir(info.dir); mod != nil {
+		relevance = 1
 		// It is. If dir is the target of a replace directive,
 		// our guessed import path is wrong. Use the real one.
 		if mod.Dir == info.dir {
@@ -452,6 +455,7 @@
 		importPathShort: importPath,
 		dir:             info.dir,
 		packageName:     info.packageName, // may not be populated if the caller didn't ask for it
+		relevance:       relevance,
 	}
 	// We may have discovered a package that has a different version
 	// in scope already. Canonicalize to that one if possible.
diff --git a/internal/imports/mod_test.go b/internal/imports/mod_test.go
index dd7c39f..dcb8818 100644
--- a/internal/imports/mod_test.go
+++ b/internal/imports/mod_test.go
@@ -10,6 +10,7 @@
 	"log"
 	"os"
 	"path/filepath"
+	"reflect"
 	"regexp"
 	"strings"
 	"sync"
@@ -831,6 +832,54 @@
 	resolver.scan(nil, true, nil)
 }
 
+func TestGetCandidatesRanking(t *testing.T) {
+	mt := setup(t, `
+-- go.mod --
+module example.com
+
+require rsc.io/quote v1.5.1
+
+-- rpackage/x.go --
+package rpackage
+import _ "rsc.io/quote"
+`, "")
+	defer mt.cleanup()
+
+	if _, err := mt.env.invokeGo("mod", "download", "rsc.io/quote/v2@v2.0.1"); err != nil {
+		t.Fatal(err)
+	}
+
+	type res struct {
+		name, path string
+	}
+	want := []res{
+		// Stdlib
+		{"bytes", "bytes"},
+		{"http", "net/http"},
+		// In scope modules
+		{"language", "golang.org/x/text/language"},
+		{"quote", "rsc.io/quote"},
+		{"rpackage", "example.com/rpackage"},
+		// Out of scope modules
+		{"quote", "rsc.io/quote/v2"},
+	}
+	candidates, err := getAllCandidates("foo.go", mt.env)
+	if err != nil {
+		t.Fatalf("getAllCandidates() = %v", err)
+	}
+	var got []res
+	for _, c := range candidates {
+		for _, w := range want {
+			if c.StmtInfo.ImportPath == w.path {
+				got = append(got, res{c.IdentName, c.StmtInfo.ImportPath})
+			}
+		}
+	}
+	if !reflect.DeepEqual(want, got) {
+		t.Errorf("wanted candidates in order %v, got %v", want, got)
+	}
+}
+
 func BenchmarkScanModCache(b *testing.B) {
 	env := &ProcessEnv{
 		Debug:  true,
diff --git a/internal/lsp/testdata/unimported/unimported.go b/internal/lsp/testdata/unimported/unimported.go
index f720392..ed0670c 100644
--- a/internal/lsp/testdata/unimported/unimported.go
+++ b/internal/lsp/testdata/unimported/unimported.go
@@ -1,13 +1,13 @@
 package unimported
 
 func _() {
-	//@unimported("", bytes, context, cryptoslashrand, externalpackage, time, unsafe)
+	//@unimported("", bytes, context, cryptoslashrand, time, unsafe, externalpackage)
 }
 
 // Create markers for unimported std lib packages. Only for use by this test.
 /* bytes */ //@item(bytes, "bytes", "\"bytes\"", "package")
 /* context */ //@item(context, "context", "\"context\"", "package")
 /* rand */ //@item(cryptoslashrand, "rand", "\"crypto/rand\"", "package")
-/* pkg */ //@item(externalpackage, "pkg", "\"example.com/extramodule/pkg\"", "package" )
-/* unsafe */ //@item(unsafe, "unsafe", "\"unsafe\"", "package")
 /* time */ //@item(time, "time", "\"time\"", "package")
+/* unsafe */ //@item(unsafe, "unsafe", "\"unsafe\"", "package")
+/* pkg */ //@item(externalpackage, "pkg", "\"example.com/extramodule/pkg\"", "package" )