go/doc: synthesize "package main" for examples

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

diff --git a/src/pkg/go/doc/example.go b/src/pkg/go/doc/example.go
index a7e0e250a2e9d74a6e82d1f397e6f831794f8d35..8fcee33af08eb6da191e55e3623b8e308a4c1b0a 100644 (file)
--- a/src/pkg/go/doc/example.go
+++ b/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
+}