// lastComment returns the last comment inside the provided block. func lastComment(b *ast.BlockStmt, c []*ast.CommentGroup) (i int, last *ast.CommentGroup) { pos, end := b.Pos(), b.End() for j, cg := range c { if cg.Pos() < pos { continue } if cg.End() > end { break } i, last = j, cg } return }
// 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 } // Find top-level declarations in the file. topDecls := make(map[*ast.Object]bool) for _, decl := range file.Decls { switch d := decl.(type) { case *ast.FuncDecl: topDecls[d.Name.Obj] = true case *ast.GenDecl: for _, spec := range d.Specs { switch s := spec.(type) { case *ast.TypeSpec: topDecls[s.Name.Obj] = true case *ast.ValueSpec: for _, id := range s.Names { topDecls[id.Obj] = true } } } } } // Find unresolved identifiers and uses of top-level declarations. unresolved := make(map[string]bool) usesTopDecl := false 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 { 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 { 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 } // Remove predeclared identifiers from unresolved list. for n := range unresolved { if predeclaredTypes[n] || predeclaredConstants[n] || predeclaredFuncs[n] { delete(unresolved, n) } } // Use unresolved identifiers to determine the imports used by this // example. The heuristic assumes package names match base import // paths for imports w/o renames (should be good enough most of the time). namedImports := make(map[string]string) // [name]path var blankImports []ast.Spec // _ imports for _, s := range file.Imports { p, err := strconv.Unquote(s.Path.Value) if err != nil { continue } n := path.Base(p) if s.Name != nil { n = s.Name.Name switch n { case "_": blankImports = append(blankImports, s) continue case ".": // We can't resolve dot imports (yet). return nil } } if unresolved[n] { namedImports[n] = p delete(unresolved, n) } } // If there are other unresolved identifiers, give up because this // synthesized file is not going to build. if len(unresolved) > 0 { return nil } // Include documentation belonging to blank imports. var comments []*ast.CommentGroup for _, s := range blankImports { if c := s.(*ast.ImportSpec).Doc; c != nil { comments = append(comments, c) } } // Include comments that are inside the function body. for _, c := range file.Comments { if body.Pos() <= c.Pos() && c.End() <= body.End() { comments = append(comments, c) } } // Strip "Output:" commment and adjust body end position. body, comments = stripOutputComment(body, comments) // Synthesize import declaration. 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 namedImports { 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) } importDecl.Specs = append(importDecl.Specs, blankImports...) // Synthesize main function. funcDecl := &ast.FuncDecl{ Name: ast.NewIdent("main"), Type: &ast.FuncType{}, Body: body, } // Synthesize file. return &ast.File{ Name: ast.NewIdent("main"), Decls: []ast.Decl{importDecl, funcDecl}, Comments: comments, } }