go/doc: add NewFromFiles with support for classifying examples

This CL is based on work started by Joe Tsai in CL 94855.
It's rebased on top of the latest master branch, and
addresses various code review comments and findings
from attempting to use the original CL in practice.

The testing package documents a naming convention for examples
so that documentation tools can associate them with:

• a package (Example or Example_suffix)
• a function F (ExampleF or ExampleF_suffix)
• a type T (ExampleT or ExampleT_suffix)
• a method T.M (ExampleT_M or ExampleT_M_suffix)

This naming convention is in widespread use and enforced
via existing go vet checks.

This change adds first-class support for classifying examples
to go/doc, the package responsible for computing package
documentation from Go AST.

There isn't a way to supply test files to New that works well.
External test files may have a package name with "_test" suffix,
so ast.NewPackage may end up using the wrong package name if given
test files. A workaround is to add test files to *ast.Package.Files
after it is returned from ast.NewPackage:

	pkg, _ := ast.NewPackage(fset, goFiles, ...)
	for name, f := range testGoFiles {
		pkg.Files[name] = f
	}
	p := doc.New(pkg, ...)

But that is not a good API.

After nearly 8 years, a new entry-point is added to the go/doc
package, the function NewFromFiles. It accepts a Go package in
the form of a list of parsed Go files (including _test.go files)
and an import path. The caller is responsible with filtering out
files based on build constraints, as was the case before with New.
NewFromFiles computes package documentation from .go files,
extracts examples from _test.go files and classifies them.

Examples fields are added to Package, Type, and Func. They are
documented to only be populated with examples found in _test.go
files provided to NewFromFiles.

The new behavior is:

1. NewFromFiles computes package documentation from provided
   parsed .go files. It extracts examples from _test.go files.
2. It assigns each Example to corresponding Package, Type,
   or Func.
3. It sets the Suffix field in each example to the suffix.
4. Malformed examples are skipped.

This change implements behavior that matches the current behavior
of existing godoc-like tools, and will enable them to rely on the
logic in go/doc instead of reimplementing it themselves.

Fixes #23864

Change-Id: Iae834f2ff92fbd1c93a9bb7c2bf47d619bee05cf
Reviewed-on: https://go-review.googlesource.com/c/go/+/204830
Run-TryBot: Dmitri Shuralyov <dmitshur@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Griesemer <gri@golang.org>

Author: dmitshur

src/go/doc/doc.go

@@ -6,8 +6,10 @@
 package doc
 
 import (
+	"fmt"
 	"go/ast"
 	"go/token"
+	"strings"
 )
 
 // Package is the documentation for an entire package.
@@ -28,6 +30,11 @@ type Package struct {
 	Types  []*Type
 	Vars   []*Value
 	Funcs  []*Func
+
+	// Examples is a sorted list of examples associated with
+	// the package. Examples are extracted from _test.go files
+	// provided to NewFromFiles.
+	Examples []*Example
 }
 
 // Value is the documentation for a (possibly grouped) var or const declaration.
@@ -50,6 +57,11 @@ type Type struct {
 	Vars    []*Value // sorted list of variables of (mostly) this type
 	Funcs   []*Func  // sorted list of functions returning this type
 	Methods []*Func  // sorted list of methods (including embedded ones) of this type
+
+	// Examples is a sorted list of examples associated with
+	// this type. Examples are extracted from _test.go files
+	// provided to NewFromFiles.
+	Examples []*Example
 }
 
 // Func is the documentation for a func declaration.
@@ -63,6 +75,11 @@ type Func struct {
 	Recv  string // actual   receiver "T" or "*T"
 	Orig  string // original receiver "T" or "*T"
 	Level int    // embedding level; 0 means not embedded
+
+	// Examples is a sorted list of examples associated with this
+	// function or method. Examples are extracted from _test.go files
+	// provided to NewFromFiles.
+	Examples []*Example
 }
 
 // A Note represents a marked comment starting with "MARKER(uid): note body".
@@ -75,7 +92,7 @@ type Note struct {
 	Body     string    // note body text
 }
 
-// Mode values control the operation of New.
+// Mode values control the operation of New and NewFromFiles.
 type Mode int
 
 const (
@@ -95,6 +112,8 @@ const (
 
 // New computes the package documentation for the given package AST.
 // New takes ownership of the AST pkg and may edit or overwrite it.
+// To have the Examples fields populated, use NewFromFiles and include
+// the package's _test.go files.
 //
 func New(pkg *ast.Package, importPath string, mode Mode) *Package {
 	var r reader
@@ -115,3 +134,86 @@ func New(pkg *ast.Package, importPath string, mode Mode) *Package {
 		Funcs:      sortedFuncs(r.funcs, true),
 	}
 }
+
+// NewFromFiles computes documentation for a package.
+//
+// The package is specified by a list of *ast.Files and corresponding
+// file set, which must not be nil. NewFromFiles does not skip files
+// based on build constraints, so it is the caller's responsibility to
+// provide only the files that are matched by the build context.
+// The import path of the package is specified by importPath.
+//
+// Examples found in _test.go files are associated with the corresponding
+// type, function, method, or the package, based on their name.
+// If the example has a suffix in its name, it is set in the
+// Example.Suffix field. Examples with malformed names are skipped.
+//
+// Optionally, a single extra argument of type Mode can be provided to
+// control low-level aspects of the documentation extraction behavior.
+//
+// NewFromFiles takes ownership of the AST files and may edit them,
+// unless the PreserveAST Mode bit is on.
+//
+func NewFromFiles(fset *token.FileSet, files []*ast.File, importPath string, opts ...interface{}) (*Package, error) {
+	// Check for invalid API usage.
+	if fset == nil {
+		panic(fmt.Errorf("doc.NewFromFiles: no token.FileSet provided (fset == nil)"))
+	}
+	var mode Mode
+	switch len(opts) { // There can only be 0 or 1 options, so a simple switch works for now.
+	case 0:
+		// Nothing to do.
+	case 1:
+		m, ok := opts[0].(Mode)
+		if !ok {
+			panic(fmt.Errorf("doc.NewFromFiles: option argument type must be doc.Mode"))
+		}
+		mode = m
+	default:
+		panic(fmt.Errorf("doc.NewFromFiles: there must not be more than 1 option argument"))
+	}
+
+	// Collect .go and _test.go files.
+	var (
+		goFiles     = make(map[string]*ast.File)
+		testGoFiles []*ast.File
+	)
+	for i := range files {
+		f := fset.File(files[i].Pos())
+		if f == nil {
+			return nil, fmt.Errorf("file files[%d] is not found in the provided file set", i)
+		}
+		switch name := f.Name(); {
+		case strings.HasSuffix(name, ".go") && !strings.HasSuffix(name, "_test.go"):
+			goFiles[name] = files[i]
+		case strings.HasSuffix(name, "_test.go"):
+			testGoFiles = append(testGoFiles, files[i])
+		default:
+			return nil, fmt.Errorf("file files[%d] filename %q does not have a .go extension", i, name)
+		}
+	}
+
+	// TODO(dmitshur,gri): A relatively high level call to ast.NewPackage with a simpleImporter
+	// ast.Importer implementation is made below. It might be possible to short-circuit and simplify.
+
+	// Compute package documentation.
+	pkg, _ := ast.NewPackage(fset, goFiles, simpleImporter, nil) // Ignore errors that can happen due to unresolved identifiers.
+	p := New(pkg, importPath, mode)
+	classifyExamples(p, Examples(testGoFiles...))
+	return p, nil
+}
+
+// simpleImporter returns a (dummy) package object named by the last path
+// component of the provided package path (as is the convention for packages).
+// This is sufficient to resolve package identifiers without doing an actual
+// import. It never returns an error.
+func simpleImporter(imports map[string]*ast.Object, path string) (*ast.Object, error) {
+	pkg := imports[path]
+	if pkg == nil {
+		// note that strings.LastIndex returns -1 if there is no "/"
+		pkg = ast.NewObj(ast.Pkg, path[strings.LastIndex(path, "/")+1:])
+		pkg.Data = ast.NewScope(nil) // required by ast.NewPackage for dot-import
+		imports[path] = pkg
+	}
+	return pkg, nil
+}

src/go/doc/doc_test.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"flag"
 	"fmt"
+	"go/ast"
 	"go/parser"
 	"go/printer"
 	"go/token"
@@ -99,8 +100,16 @@ func test(t *testing.T, mode Mode) {
 
 	// test packages
 	for _, pkg := range pkgs {
-		importpath := dataDir + "/" + pkg.Name
-		doc := New(pkg, importpath, mode)
+		importPath := dataDir + "/" + pkg.Name
+		var files []*ast.File
+		for _, f := range pkg.Files {
+			files = append(files, f)
+		}
+		doc, err := NewFromFiles(fset, files, importPath, mode)
+		if err != nil {
+			t.Error(err)
+			continue
+		}
 
 		// golden files always use / in filenames - canonicalize them
 		for i, filename := range doc.Filenames {

src/go/doc/example.go

@@ -18,9 +18,10 @@ import (
 	"unicode/utf8"
 )
 
-// An Example represents an example function found in a source files.
+// An Example represents an example function found in a test source file.
 type Example struct {
-	Name        string // name of the item being exemplified
+	Name        string // name of the item being exemplified (including optional suffix)
+	Suffix      string // example suffix, without leading '_' (only populated by NewFromFiles)
 	Doc         string // example function doc string
 	Code        ast.Node
 	Play        *ast.File // a whole program version of the example
@@ -31,8 +32,10 @@ type Example struct {
 	Order       int  // original source code order
 }
 
-// Examples returns the examples found in the files, sorted by Name field.
+// Examples returns the examples found in testFiles, sorted by Name field.
 // The Order fields record the order in which the examples were encountered.
+// The Suffix field is not populated when Examples is called directly, it is
+// only populated by NewFromFiles for examples it finds in _test.go files.
 //
 // Playable Examples must be in a package whose name ends in "_test".
 // An Example is "playable" (the Play field is non-nil) in either of these
@@ -44,9 +47,9 @@ type Example struct {
 //     example function, zero test or benchmark functions, and at least one
 //     top-level function, type, variable, or constant declaration other
 //     than the example function.
-func Examples(files ...*ast.File) []*Example {
+func Examples(testFiles ...*ast.File) []*Example {
 	var list []*Example
-	for _, file := range files {
+	for _, file := range testFiles {
 		hasTests := false // file contains tests or benchmarks
 		numDecl := 0      // number of non-import declarations in the file
 		var flist []*Example
@@ -441,3 +444,101 @@ func lastComment(b *ast.BlockStmt, c []*ast.CommentGroup) (i int, last *ast.Comm
 	}
 	return
 }
+
+// classifyExamples classifies examples and assigns them to the Examples field
+// of the relevant Func, Type, or Package that the example is associated with.
+//
+// The classification process is ambiguous in some cases:
+//
+// 	- ExampleFoo_Bar matches a type named Foo_Bar
+// 	  or a method named Foo.Bar.
+// 	- ExampleFoo_bar matches a type named Foo_bar
+// 	  or Foo (with a "bar" suffix).
+//
+// Examples with malformed names are not associated with anything.
+//
+func classifyExamples(p *Package, examples []*Example) {
+	if len(examples) == 0 {
+		return
+	}
+
+	// Mapping of names for funcs, types, and methods to the example listing.
+	ids := make(map[string]*[]*Example)
+	ids[""] = &p.Examples // package-level examples have an empty name
+	for _, f := range p.Funcs {
+		if !token.IsExported(f.Name) {
+			continue
+		}
+		ids[f.Name] = &f.Examples
+	}
+	for _, t := range p.Types {
+		if !token.IsExported(t.Name) {
+			continue
+		}
+		ids[t.Name] = &t.Examples
+		for _, f := range t.Funcs {
+			if !token.IsExported(f.Name) {
+				continue
+			}
+			ids[f.Name] = &f.Examples
+		}
+		for _, m := range t.Methods {
+			if !token.IsExported(m.Name) || m.Level != 0 { // avoid forwarded methods from embedding
+				continue
+			}
+			ids[strings.TrimPrefix(m.Recv, "*")+"_"+m.Name] = &m.Examples
+		}
+	}
+
+	// Group each example with the associated func, type, or method.
+	for _, ex := range examples {
+		// Consider all possible split points for the suffix
+		// by starting at the end of string (no suffix case),
+		// then trying all positions that contain a '_' character.
+		//
+		// An association is made on the first successful match.
+		// Examples with malformed names that match nothing are skipped.
+		for i := len(ex.Name); i >= 0; i = strings.LastIndexByte(ex.Name[:i], '_') {
+			prefix, suffix, ok := splitExampleName(ex.Name, i)
+			if !ok {
+				continue
+			}
+			exs, ok := ids[prefix]
+			if !ok {
+				continue
+			}
+			ex.Suffix = suffix
+			*exs = append(*exs, ex)
+			break
+		}
+	}
+
+	// Sort list of example according to the user-specified suffix name.
+	for _, exs := range ids {
+		sort.Slice((*exs), func(i, j int) bool {
+			return (*exs)[i].Suffix < (*exs)[j].Suffix
+		})
+	}
+}
+
+// splitExampleName attempts to split example name s at index i,
+// and reports if that produces a valid split. The suffix may be
+// absent. Otherwise, it must start with a lower-case letter and
+// be preceded by '_'.
+//
+// One of i == len(s) or s[i] == '_' must be true.
+func splitExampleName(s string, i int) (prefix, suffix string, ok bool) {
+	if i == len(s) {
+		return s, "", true
+	}
+	if i == len(s)-1 {
+		return "", "", false
+	}
+	prefix, suffix = s[:i], s[i+1:]
+	return prefix, suffix, isExampleSuffix(suffix)
+}
+
+func isExampleSuffix(s string) bool {
+	r, size := utf8.DecodeRuneInString(s)
+	return size > 0 && unicode.IsLower(r)
+}