go/doc: synthesize "package main" for examples

R=gri
CC=golang-dev
https://golang.org/cl/6525046

src/pkg/go/doc/example.go

@@ -9,8 +9,10 @@ package doc
 import (
 	"go/ast"
 	"go/token"
+	"path"
 	"regexp"
 	"sort"
+	"strconv"
 	"strings"
 	"unicode"
 	"unicode/utf8"
@@ -20,6 +22,7 @@ type Example struct {
 	Name     string // name of the item being exemplified
 	Doc      string // example function doc string
 	Code     ast.Node
+	Play     *ast.File // a whole program version of the example
 	Comments []*ast.CommentGroup
 	Output   string // expected output
 }
@@ -56,6 +59,7 @@ func Examples(files ...*ast.File) []*Example {
 				Name:     name[len("Example"):],
 				Doc:      doc,
 				Code:     f.Body,
+				Play:     playExample(file, f.Body),
 				Comments: file.Comments,
 				Output:   exampleOutput(f, file.Comments),
 			})
@@ -115,3 +119,91 @@ type exampleByName []*Example
 func (s exampleByName) Len() int           { return len(s) }
 func (s exampleByName) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
 func (s exampleByName) Less(i, j int) bool { return s[i].Name < s[j].Name }
+
+// 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 {
+	if !strings.HasSuffix(file.Name.Name, "_test") {
+		// We don't support examples that are part of the
+		// greater package (yet).
+		return nil
+	}
+
+	// Determine the imports we need based on unresolved identifiers.
+	// This is a heuristic that presumes package names match base import paths.
+	// (Should be good enough most of the time.)
+	var unresolved []*ast.Ident
+	ast.Inspect(body, func(n ast.Node) bool {
+		if e, ok := n.(*ast.SelectorExpr); ok {
+			if id, ok := e.X.(*ast.Ident); ok && id.Obj == nil {
+				unresolved = append(unresolved, id)
+			}
+		}
+		return true
+	})
+	imports := make(map[string]string) // [name]path
+	for _, s := range file.Imports {
+		p, err := strconv.Unquote(s.Path.Value)
+		if err != nil {
+			continue
+		}
+		n := path.Base(p)
+		if s.Name != nil {
+			if s.Name.Name == "." {
+				// We can't resolve dot imports (yet).
+				return nil
+			}
+			n = s.Name.Name
+		}
+		for _, id := range unresolved {
+			if n == id.Name {
+				imports[n] = p
+				break
+			}
+		}
+	}
+
+	// TODO(adg): look for other unresolved identifiers and, if found, give up.
+
+	// Synthesize new imports.
+	importDecl := &ast.GenDecl{
+		Tok:    token.IMPORT,
+		Lparen: 1, // Need non-zero Lparen and Rparen so that printer
+		Rparen: 1, // treats this as a factored import.
+	}
+	for n, p := range imports {
+		s := &ast.ImportSpec{Path: &ast.BasicLit{Value: strconv.Quote(p)}}
+		if path.Base(p) != n {
+			s.Name = ast.NewIdent(n)
+		}
+		importDecl.Specs = append(importDecl.Specs, s)
+	}
+
+	// Synthesize main function.
+	funcDecl := &ast.FuncDecl{
+		Name: ast.NewIdent("main"),
+		Type: &ast.FuncType{},
+		Body: body,
+	}
+
+	// Filter out comments that are outside the function body.
+	var comments []*ast.CommentGroup
+	for _, c := range file.Comments {
+		if c.Pos() < body.Pos() || c.Pos() >= body.End() {
+			continue
+		}
+		comments = append(comments, c)
+	}
+
+	// Synthesize file.
+	f := &ast.File{
+		Name:     ast.NewIdent("main"),
+		Decls:    []ast.Decl{importDecl, funcDecl},
+		Comments: comments,
+	}
+
+	// TODO(adg): look for resolved identifiers declared outside function scope
+	// and include their declarations in the new file.
+
+	return f
+}