cmd/doc: adapt directory search for modules

Previously, cmd/doc treated GOROOT/src and GOPATH/src
as the roots of the directory trees holding packages, assuming
that the import path would be the path elements after the src directory.

With modules, each module serves as its own root of a file tree,
and the import path prefix starts with the module path before
adding the path elements after the module root.

There are ways we could make this more efficient,
but for now this is a fairly small adjustment to get 'go doc'
working OK for modules for Go 1.11.

Fixes #26635.

Change-Id: Ifdee4194601312846c7b1fc67f2fe7a4a44269cc
Reviewed-on: https://go-review.googlesource.com/126799
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rob Pike <r@golang.org>

5 files changed, 180 insertions, 31 deletions

diff --git a/src/cmd/doc/dirs.go b/src/cmd/doc/dirs.go
index f5fb795dc7..24bd797eb5 100644
--- a/src/cmd/doc/dirs.go
+++ b/src/cmd/doc/dirs.go
@@ -5,32 +5,41 @@
 package main
 
 import (
+	"bytes"
 	"log"
 	"os"
-	"path"
+	"os/exec"
 	"path/filepath"
 	"strings"
+	"sync"
 )
 
+// A Dir describes a directory holding code by specifying
+// the expected import path and the file system directory.
+type Dir struct {
+	importPath string // import path for that dir
+	dir        string // file system directory
+}
+
 // Dirs is a structure for scanning the directory tree.
 // Its Next method returns the next Go source directory it finds.
 // Although it can be used to scan the tree multiple times, it
 // only walks the tree once, caching the data it finds.
 type Dirs struct {
-	scan   chan string // directories generated by walk.
-	paths  []string    // Cache of known paths.
-	offset int         // Counter for Next.
+	scan   chan Dir // Directories generated by walk.
+	hist   []Dir    // History of reported Dirs.
+	offset int      // Counter for Next.
 }
 
 var dirs Dirs
 
 // dirsInit starts the scanning of package directories in GOROOT and GOPATH. Any
 // extra paths passed to it are included in the channel.
-func dirsInit(extra ...string) {
-	dirs.paths = make([]string, 0, 1000)
-	dirs.paths = append(dirs.paths, extra...)
-	dirs.scan = make(chan string)
-	go dirs.walk()
+func dirsInit(extra ...Dir) {
+	dirs.hist = make([]Dir, 0, 1000)
+	dirs.hist = append(dirs.hist, extra...)
+	dirs.scan = make(chan Dir)
+	go dirs.walk(codeRoots())
 }
 
 // Reset puts the scan back at the beginning.
@@ -40,25 +49,24 @@ func (d *Dirs) Reset() {
 
 // Next returns the next directory in the scan. The boolean
 // is false when the scan is done.
-func (d *Dirs) Next() (string, bool) {
-	if d.offset < len(d.paths) {
-		path := d.paths[d.offset]
+func (d *Dirs) Next() (Dir, bool) {
+	if d.offset < len(d.hist) {
+		dir := d.hist[d.offset]
 		d.offset++
-		return path, true
+		return dir, true
 	}
-	path, ok := <-d.scan
+	dir, ok := <-d.scan
 	if !ok {
-		return "", false
+		return Dir{}, false
 	}
-	d.paths = append(d.paths, path)
+	d.hist = append(d.hist, dir)
 	d.offset++
-	return path, ok
+	return dir, ok
 }
 
 // walk walks the trees in GOROOT and GOPATH.
-func (d *Dirs) walk() {
-	d.bfsWalkRoot(buildCtx.GOROOT)
-	for _, root := range splitGopath() {
+func (d *Dirs) walk(roots []Dir) {
+	for _, root := range roots {
 		d.bfsWalkRoot(root)
 	}
 	close(d.scan)
@@ -66,13 +74,13 @@ func (d *Dirs) walk() {
 
 // bfsWalkRoot walks a single directory hierarchy in breadth-first lexical order.
 // Each Go source directory it finds is delivered on d.scan.
-func (d *Dirs) bfsWalkRoot(root string) {
-	root = path.Join(root, "src")
+func (d *Dirs) bfsWalkRoot(root Dir) {
+	root.dir = filepath.Clean(root.dir) // because filepath.Join will do it anyway
 
 	// this is the queue of directories to examine in this pass.
 	this := []string{}
 	// next is the queue of directories to examine in the next pass.
-	next := []string{root}
+	next := []string{root.dir}
 
 	for len(next) > 0 {
 		this, next = next, this[0:0]
@@ -105,14 +113,83 @@ func (d *Dirs) bfsWalkRoot(root string) {
 				if name[0] == '.' || name[0] == '_' || name == "testdata" {
 					continue
 				}
+				// Ignore vendor when using modules.
+				if usingModules && name == "vendor" {
+					continue
+				}
 				// Remember this (fully qualified) directory for the next pass.
 				next = append(next, filepath.Join(dir, name))
 			}
 			if hasGoFiles {
 				// It's a candidate.
-				d.scan <- dir
+				importPath := root.importPath
+				if len(dir) > len(root.dir) {
+					if importPath != "" {
+						importPath += "/"
+					}
+					importPath += filepath.ToSlash(dir[len(root.dir)+1:])
+				}
+				d.scan <- Dir{importPath, dir}
 			}
 		}
 
 	}
 }
+
+var testGOPATH = false // force GOPATH use for testing
+
+// codeRoots returns the code roots to search for packages.
+// In GOPATH mode this is GOROOT/src and GOPATH/src, with empty import paths.
+// In module mode, this is each module root, with an import path set to its module path.
+func codeRoots() []Dir {
+	codeRootsCache.once.Do(func() {
+		codeRootsCache.roots = findCodeRoots()
+	})
+	return codeRootsCache.roots
+}
+
+var codeRootsCache struct {
+	once  sync.Once
+	roots []Dir
+}
+
+var usingModules bool
+
+func findCodeRoots() []Dir {
+	list := []Dir{{"", filepath.Join(buildCtx.GOROOT, "src")}}
+
+	if !testGOPATH {
+		// Check for use of modules by 'go env GOMOD',
+		// which reports a go.mod file path if modules are enabled.
+		stdout, _ := exec.Command("go", "env", "GOMOD").Output()
+		usingModules = bytes.Contains(stdout, []byte("go.mod"))
+	}
+
+	if !usingModules {
+		for _, root := range splitGopath() {
+			list = append(list, Dir{"", filepath.Join(root, "src")})
+		}
+		return list
+	}
+
+	// Find module root directories from go list.
+	// Eventually we want golang.org/x/tools/go/packages
+	// to handle the entire file system search and become go/packages,
+	// but for now enumerating the module roots lets us fit modules
+	// into the current code with as few changes as possible.
+	cmd := exec.Command("go", "list", "-m", "-f={{.Path}}\t{{.Dir}}", "all")
+	cmd.Stderr = os.Stderr
+	out, _ := cmd.Output()
+	for _, line := range strings.Split(string(out), "\n") {
+		i := strings.Index(line, "\t")
+		if i < 0 {
+			continue
+		}
+		path, dir := line[:i], line[i+1:]
+		if dir != "" {
+			list = append(list, Dir{path, dir})
+		}
+	}
+
+	return list
+}
diff --git a/src/cmd/doc/doc_test.go b/src/cmd/doc/doc_test.go
index e68e95f3fb..6010f04b56 100644
--- a/src/cmd/doc/doc_test.go
+++ b/src/cmd/doc/doc_test.go
@@ -18,6 +18,7 @@ import (
 func TestMain(m *testing.M) {
 	// Clear GOPATH so we don't access the user's own packages in the test.
 	buildCtx.GOPATH = ""
+	testGOPATH = true // force GOPATH mode; module test is in cmd/go/testdata/script/mod_doc.txt
 
 	// Add $GOROOT/src/cmd/doc/testdata explicitly so we can access its contents in the test.
 	// Normally testdata directories are ignored, but sending it to dirs.scan directly is
@@ -26,7 +27,7 @@ func TestMain(m *testing.M) {
 	if err != nil {
 		panic(err)
 	}
-	dirsInit(testdataDir, filepath.Join(testdataDir, "nested"), filepath.Join(testdataDir, "nested", "nested"))
+	dirsInit(Dir{"testdata", testdataDir}, Dir{"testdata/nested", filepath.Join(testdataDir, "nested")}, Dir{"testdata/nested/nested", filepath.Join(testdataDir, "nested", "nested")})
 
 	os.Exit(m.Run())
 }
@@ -537,7 +538,7 @@ func TestDoc(t *testing.T) {
 		var flagSet flag.FlagSet
 		err := do(&b, &flagSet, test.args)
 		if err != nil {
-			t.Fatalf("%s: %s\n", test.name, err)
+			t.Fatalf("%s %v: %s\n", test.name, test.args, err)
 		}
 		output := b.Bytes()
 		failed := false
diff --git a/src/cmd/doc/main.go b/src/cmd/doc/main.go
index bf0c7723f8..982c8e054a 100644
--- a/src/cmd/doc/main.go
+++ b/src/cmd/doc/main.go
@@ -39,6 +39,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
 	"unicode"
@@ -189,6 +190,10 @@ func parseArgs(args []string) (pkg *build.Package, path, symbol string, more boo
 		// Done below.
 	case 2:
 		// Package must be findable and importable.
+		pkg, err := build.Import(args[0], "", build.ImportComment)
+		if err == nil {
+			return pkg, args[0], args[1], false
+		}
 		for {
 			packagePath, ok := findNextPackage(arg)
 			if !ok {
@@ -355,14 +360,22 @@ func findNextPackage(pkg string) (string, bool) {
 	if pkg == "" || isUpper(pkg) { // Upper case symbol cannot be a package name.
 		return "", false
 	}
-	pkgString := filepath.Clean(string(filepath.Separator) + pkg)
+	if filepath.IsAbs(pkg) {
+		if dirs.offset == 0 {
+			dirs.offset = -1
+			return pkg, true
+		}
+		return "", false
+	}
+	pkg = path.Clean(pkg)
+	pkgSuffix := "/" + pkg
 	for {
-		path, ok := dirs.Next()
+		d, ok := dirs.Next()
 		if !ok {
 			return "", false
 		}
-		if strings.HasSuffix(path, pkgString) {
-			return path, true
+		if d.importPath == pkg || strings.HasSuffix(d.importPath, pkgSuffix) {
+			return d.dir, true
 		}
 	}
 }
diff --git a/src/cmd/doc/pkg.go b/src/cmd/doc/pkg.go
index 8ff9ff57ac..14e41b9106 100644
--- a/src/cmd/doc/pkg.go
+++ b/src/cmd/doc/pkg.go
@@ -425,12 +425,36 @@ func (pkg *Package) packageClause(checkUserPath bool) {
 			return
 		}
 	}
+
 	importPath := pkg.build.ImportComment
 	if importPath == "" {
 		importPath = pkg.build.ImportPath
 	}
+
+	// If we're using modules, the import path derived from module code locations wins.
+	// If we did a file system scan, we knew the import path when we found the directory.
+	// But if we started with a directory name, we never knew the import path.
+	// Either way, we don't know it now, and it's cheap to (re)compute it.
+	if usingModules {
+		for _, root := range codeRoots() {
+			if pkg.build.Dir == root.dir {
+				importPath = root.importPath
+				break
+			}
+			if strings.HasPrefix(pkg.build.Dir, root.dir+string(filepath.Separator)) {
+				suffix := filepath.ToSlash(pkg.build.Dir[len(root.dir)+1:])
+				if root.importPath == "" {
+					importPath = suffix
+				} else {
+					importPath = root.importPath + "/" + suffix
+				}
+				break
+			}
+		}
+	}
+
 	pkg.Printf("package %s // import %q\n\n", pkg.name, importPath)
-	if importPath != pkg.build.ImportPath {
+	if !usingModules && importPath != pkg.build.ImportPath {
 		pkg.Printf("WARNING: package source is installed in %q\n", pkg.build.ImportPath)
 	}
 }
diff --git a/src/cmd/go/testdata/script/mod_doc.txt b/src/cmd/go/testdata/script/mod_doc.txt
new file mode 100644
index 0000000000..223283d5f2
--- /dev/null
+++ b/src/cmd/go/testdata/script/mod_doc.txt
@@ -0,0 +1,34 @@
+# go doc should find module documentation
+
+env GO111MODULE=on
+
+go doc y
+stdout 'Package y is.*alphabet'
+stdout 'import "x/y"'
+go doc x/y
+stdout 'Package y is.*alphabet'
+! go doc quote.Hello
+stderr 'doc: symbol quote is not a type' # because quote is not in local cache
+go list rsc.io/quote # now it is
+go doc quote.Hello
+stdout 'Hello returns a greeting'
+go doc quote
+stdout 'Package quote collects pithy sayings.'
+
+# Double-check go doc y when y is not in GOPATH/src.
+env GOPATH=$WORK/altgopath
+go doc x/y
+stdout 'Package y is.*alphabet'
+go doc y
+stdout 'Package y is.*alphabet'
+
+-- go.mod --
+module x
+require rsc.io/quote v1.5.2
+
+-- y/y.go --
+// Package y is the next to last package of the alphabet.
+package y
+
+-- x.go --
+package x