go/doc: make examples that depend on top-level decls playable

Currently, the following example cannot run in playground:

    func a() {
        fmt.Println("A")
    }

    func ExampleA() {
        a()
    }

This CL solves it.

Fixes #23095

Change-Id: I5a492ff886a743f20cb4ae646e8453bde9c5f0da
Reviewed-on: https://go-review.googlesource.com/83615
Run-TryBot: Robert Griesemer <gri@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Griesemer <gri@golang.org>

src/go/doc/example.go

@@ -77,7 +77,7 @@ func Examples(files ...*ast.File) []*Example {
 				Name:        name[len("Example"):],
 				Doc:         doc,
 				Code:        f.Body,
-				Play:        playExample(file, f.Body),
+				Play:        playExample(file, f),
 				Comments:    file.Comments,
 				Output:      output,
 				Unordered:   unordered,
@@ -140,27 +140,39 @@ func isTest(name, prefix string) bool {
 
 // playExample synthesizes a new *ast.File based on the provided
 // file with the provided function body as the body of main.
-func playExample(file *ast.File, body *ast.BlockStmt) *ast.File {
+func playExample(file *ast.File, f *ast.FuncDecl) *ast.File {
+	body := f.Body
+
 	if !strings.HasSuffix(file.Name.Name, "_test") {
 		// We don't support examples that are part of the
 		// greater package (yet).
 		return nil
 	}
 
-	// Find top-level declarations in the file.
-	topDecls := make(map[*ast.Object]bool)
+	// Collect top-level declarations in the file.
+	topDecls := make(map[*ast.Object]ast.Decl)
+	typMethods := make(map[string][]ast.Decl)
+
 	for _, decl := range file.Decls {
 		switch d := decl.(type) {
 		case *ast.FuncDecl:
-			topDecls[d.Name.Obj] = true
+			if d.Recv == nil {
+				topDecls[d.Name.Obj] = d
+			} else {
+				if len(d.Recv.List) == 1 {
+					t := d.Recv.List[0].Type
+					tname, _ := baseTypeName(t)
+					typMethods[tname] = append(typMethods[tname], d)
+				}
+			}
 		case *ast.GenDecl:
 			for _, spec := range d.Specs {
 				switch s := spec.(type) {
 				case *ast.TypeSpec:
-					topDecls[s.Name.Obj] = true
+					topDecls[s.Name.Obj] = d
 				case *ast.ValueSpec:
-					for _, id := range s.Names {
-						topDecls[id.Obj] = true
+					for _, name := range s.Names {
+						topDecls[name.Obj] = d
 					}
 				}
 			}
@@ -169,36 +181,59 @@ func playExample(file *ast.File, body *ast.BlockStmt) *ast.File {
 
 	// Find unresolved identifiers and uses of top-level declarations.
 	unresolved := make(map[string]bool)
-	usesTopDecl := false
+	var depDecls []ast.Decl
+	hasDepDecls := make(map[ast.Decl]bool)
+
 	var inspectFunc func(ast.Node) bool
 	inspectFunc = func(n ast.Node) bool {
-		// For selector expressions, only inspect the left hand side.
-		// (For an expression like fmt.Println, only add "fmt" to the
-		// set of unresolved names, not "Println".)
-		if e, ok := n.(*ast.SelectorExpr); ok {
+		switch e := n.(type) {
+		case *ast.Ident:
+			if e.Obj == nil {
+				unresolved[e.Name] = true
+			} else if d := topDecls[e.Obj]; d != nil {
+				if !hasDepDecls[d] {
+					hasDepDecls[d] = true
+					depDecls = append(depDecls, d)
+				}
+			}
+			return true
+		case *ast.SelectorExpr:
+			// For selector expressions, only inspect the left hand side.
+			// (For an expression like fmt.Println, only add "fmt" to the
+			// set of unresolved names, not "Println".)
 			ast.Inspect(e.X, inspectFunc)
 			return false
-		}
-		// For key value expressions, only inspect the value
-		// as the key should be resolved by the type of the
-		// composite literal.
-		if e, ok := n.(*ast.KeyValueExpr); ok {
+		case *ast.KeyValueExpr:
+			// For key value expressions, only inspect the value
+			// as the key should be resolved by the type of the
+			// composite literal.
 			ast.Inspect(e.Value, inspectFunc)
 			return false
 		}
-		if id, ok := n.(*ast.Ident); ok {
-			if id.Obj == nil {
-				unresolved[id.Name] = true
-			} else if topDecls[id.Obj] {
-				usesTopDecl = true
-			}
-		}
 		return true
 	}
 	ast.Inspect(body, inspectFunc)
-	if usesTopDecl {
-		// We don't support examples that are not self-contained (yet).
-		return nil
+	for i := 0; i < len(depDecls); i++ {
+		switch d := depDecls[i].(type) {
+		case *ast.FuncDecl:
+			ast.Inspect(d.Body, inspectFunc)
+		case *ast.GenDecl:
+			for _, spec := range d.Specs {
+				switch s := spec.(type) {
+				case *ast.TypeSpec:
+					ast.Inspect(s.Type, inspectFunc)
+
+					depDecls = append(depDecls, typMethods[s.Name.Name]...)
+				case *ast.ValueSpec:
+					if s.Type != nil {
+						ast.Inspect(s.Type, inspectFunc)
+					}
+					for _, val := range s.Values {
+						ast.Inspect(val, inspectFunc)
+					}
+				}
+			}
+		}
 	}
 
 	// Remove predeclared identifiers from unresolved list.
@@ -261,6 +296,20 @@ func playExample(file *ast.File, body *ast.BlockStmt) *ast.File {
 	// end position.
 	body, comments = stripOutputComment(body, comments)
 
+	// Include documentation belonging to dependent declarations.
+	for _, d := range depDecls {
+		switch d := d.(type) {
+		case *ast.GenDecl:
+			if d.Doc != nil {
+				comments = append(comments, d.Doc)
+			}
+		case *ast.FuncDecl:
+			if d.Doc != nil {
+				comments = append(comments, d.Doc)
+			}
+		}
+	}
+
 	// Synthesize import declaration.
 	importDecl := &ast.GenDecl{
 		Tok:    token.IMPORT,
@@ -279,14 +328,27 @@ func playExample(file *ast.File, body *ast.BlockStmt) *ast.File {
 	// Synthesize main function.
 	funcDecl := &ast.FuncDecl{
 		Name: ast.NewIdent("main"),
-		Type: &ast.FuncType{Params: &ast.FieldList{}}, // FuncType.Params must be non-nil
+		Type: f.Type,
 		Body: body,
 	}
 
+	decls := make([]ast.Decl, 0, 2+len(depDecls))
+	decls = append(decls, importDecl)
+	decls = append(decls, depDecls...)
+	decls = append(decls, funcDecl)
+
+	sort.Slice(decls, func(i, j int) bool {
+		return decls[i].Pos() < decls[j].Pos()
+	})
+
+	sort.Slice(comments, func(i, j int) bool {
+		return comments[i].Pos() < comments[j].Pos()
+	})
+
 	// Synthesize file.
 	return &ast.File{
 		Name:     ast.NewIdent("main"),
-		Decls:    []ast.Decl{importDecl, funcDecl},
+		Decls:    decls,
 		Comments: comments,
 	}
 }

src/go/doc/example_test.go

@@ -21,6 +21,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"sort"
 	"os/exec"
 )
 
@@ -67,6 +68,46 @@ var keyValueTopDecl = struct {
 
 func ExampleKeyValueTopDecl() {
 	fmt.Print(keyValueTopDecl)
+	// Output: a: "B", b: 2
+}
+
+// Person represents a person by name and age.
+type Person struct {
+    Name string
+    Age  int
+}
+
+// String returns a string representation of the Person.
+func (p Person) String() string {
+    return fmt.Sprintf("%s: %d", p.Name, p.Age)
+}
+
+// ByAge implements sort.Interface for []Person based on
+// the Age field.
+type ByAge []Person
+
+// Len returns the number of elements in ByAge.
+func (a (ByAge)) Len() int { return len(a) }
+
+// Swap swaps the elements in ByAge.
+func (a ByAge) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
+func (a ByAge) Less(i, j int) bool { return a[i].Age < a[j].Age }
+
+// people is the array of Person
+var people = []Person{
+	{"Bob", 31},
+	{"John", 42},
+	{"Michael", 17},
+	{"Jenny", 26},
+}
+
+func ExampleSort() {
+    fmt.Println(people)
+    sort.Sort(ByAge(people))
+    fmt.Println(people)
+    // Output:
+    // [Bob: 31 John: 42 Michael: 17 Jenny: 26]
+    // [Michael: 17 Jenny: 26 Bob: 31 John: 42]
 }
 `
 
@@ -93,8 +134,14 @@ var exampleTestCases = []struct {
 		Output: "Name: \"play\"\n",
 	},
 	{
-		Name: "KeyValueTopDecl",
-		Play: "<nil>",
+		Name:   "KeyValueTopDecl",
+		Play:   exampleKeyValueTopDeclPlay,
+		Output: "a: \"B\", b: 2\n",
+	},
+	{
+		Name:   "Sort",
+		Play:   exampleSortPlay,
+		Output: "[Bob: 31 John: 42 Michael: 17 Jenny: 26]\n[Michael: 17 Jenny: 26 Bob: 31 John: 42]\n",
 	},
 }
 
@@ -158,6 +205,69 @@ func main() {
 }
 `
 
+const exampleKeyValueTopDeclPlay = `package main
+
+import (
+	"fmt"
+)
+
+var keyValueTopDecl = struct {
+	a string
+	b int
+}{
+	a: "B",
+	b: 2,
+}
+
+func main() {
+	fmt.Print(keyValueTopDecl)
+}
+`
+
+const exampleSortPlay = `package main
+
+import (
+	"fmt"
+	"sort"
+)
+
+// Person represents a person by name and age.
+type Person struct {
+	Name string
+	Age  int
+}
+
+// String returns a string representation of the Person.
+func (p Person) String() string {
+	return fmt.Sprintf("%s: %d", p.Name, p.Age)
+}
+
+// ByAge implements sort.Interface for []Person based on
+// the Age field.
+type ByAge []Person
+
+// Len returns the number of elements in ByAge.
+func (a ByAge) Len() int { return len(a) }
+
+// Swap swaps the elements in ByAge.
+func (a ByAge) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
+func (a ByAge) Less(i, j int) bool { return a[i].Age < a[j].Age }
+
+// people is the array of Person
+var people = []Person{
+	{"Bob", 31},
+	{"John", 42},
+	{"Michael", 17},
+	{"Jenny", 26},
+}
+
+func main() {
+	fmt.Println(people)
+	sort.Sort(ByAge(people))
+	fmt.Println(people)
+}
+`
+
 func TestExamples(t *testing.T) {
 	fset := token.NewFileSet()
 	file, err := parser.ParseFile(fset, "test.go", strings.NewReader(exampleTestFile), parser.ParseComments)

src/go/doc/reader.go

@@ -104,6 +104,8 @@ func baseTypeName(x ast.Expr) (name string, imported bool) {
 			// assume type is imported
 			return t.Sel.Name, true
 		}
+	case *ast.ParenExpr:
+		return baseTypeName(t.X)
 	case *ast.StarExpr:
 		return baseTypeName(t.X)
 	}