• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1// Copyright 2011 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// Extract example functions from file ASTs.
6
7package doc
8
9import (
10	"cmp"
11	"go/ast"
12	"go/token"
13	"internal/lazyregexp"
14	"path"
15	"slices"
16	"strconv"
17	"strings"
18	"unicode"
19	"unicode/utf8"
20)
21
22// An Example represents an example function found in a test source file.
23type Example struct {
24	Name        string // name of the item being exemplified (including optional suffix)
25	Suffix      string // example suffix, without leading '_' (only populated by NewFromFiles)
26	Doc         string // example function doc string
27	Code        ast.Node
28	Play        *ast.File // a whole program version of the example
29	Comments    []*ast.CommentGroup
30	Output      string // expected output
31	Unordered   bool
32	EmptyOutput bool // expect empty output
33	Order       int  // original source code order
34}
35
36// Examples returns the examples found in testFiles, sorted by Name field.
37// The Order fields record the order in which the examples were encountered.
38// The Suffix field is not populated when Examples is called directly, it is
39// only populated by [NewFromFiles] for examples it finds in _test.go files.
40//
41// Playable Examples must be in a package whose name ends in "_test".
42// An Example is "playable" (the Play field is non-nil) in either of these
43// circumstances:
44//   - The example function is self-contained: the function references only
45//     identifiers from other packages (or predeclared identifiers, such as
46//     "int") and the test file does not include a dot import.
47//   - The entire test file is the example: the file contains exactly one
48//     example function, zero test, fuzz test, or benchmark function, and at
49//     least one top-level function, type, variable, or constant declaration
50//     other than the example function.
51func Examples(testFiles ...*ast.File) []*Example {
52	var list []*Example
53	for _, file := range testFiles {
54		hasTests := false // file contains tests, fuzz test, or benchmarks
55		numDecl := 0      // number of non-import declarations in the file
56		var flist []*Example
57		for _, decl := range file.Decls {
58			if g, ok := decl.(*ast.GenDecl); ok && g.Tok != token.IMPORT {
59				numDecl++
60				continue
61			}
62			f, ok := decl.(*ast.FuncDecl)
63			if !ok || f.Recv != nil {
64				continue
65			}
66			numDecl++
67			name := f.Name.Name
68			if isTest(name, "Test") || isTest(name, "Benchmark") || isTest(name, "Fuzz") {
69				hasTests = true
70				continue
71			}
72			if !isTest(name, "Example") {
73				continue
74			}
75			if params := f.Type.Params; len(params.List) != 0 {
76				continue // function has params; not a valid example
77			}
78			if f.Body == nil { // ast.File.Body nil dereference (see issue 28044)
79				continue
80			}
81			var doc string
82			if f.Doc != nil {
83				doc = f.Doc.Text()
84			}
85			output, unordered, hasOutput := exampleOutput(f.Body, file.Comments)
86			flist = append(flist, &Example{
87				Name:        name[len("Example"):],
88				Doc:         doc,
89				Code:        f.Body,
90				Play:        playExample(file, f),
91				Comments:    file.Comments,
92				Output:      output,
93				Unordered:   unordered,
94				EmptyOutput: output == "" && hasOutput,
95				Order:       len(flist),
96			})
97		}
98		if !hasTests && numDecl > 1 && len(flist) == 1 {
99			// If this file only has one example function, some
100			// other top-level declarations, and no tests or
101			// benchmarks, use the whole file as the example.
102			flist[0].Code = file
103			flist[0].Play = playExampleFile(file)
104		}
105		list = append(list, flist...)
106	}
107	// sort by name
108	slices.SortFunc(list, func(a, b *Example) int {
109		return cmp.Compare(a.Name, b.Name)
110	})
111	return list
112}
113
114var outputPrefix = lazyregexp.New(`(?i)^[[:space:]]*(unordered )?output:`)
115
116// Extracts the expected output and whether there was a valid output comment.
117func exampleOutput(b *ast.BlockStmt, comments []*ast.CommentGroup) (output string, unordered, ok bool) {
118	if _, last := lastComment(b, comments); last != nil {
119		// test that it begins with the correct prefix
120		text := last.Text()
121		if loc := outputPrefix.FindStringSubmatchIndex(text); loc != nil {
122			if loc[2] != -1 {
123				unordered = true
124			}
125			text = text[loc[1]:]
126			// Strip zero or more spaces followed by \n or a single space.
127			text = strings.TrimLeft(text, " ")
128			if len(text) > 0 && text[0] == '\n' {
129				text = text[1:]
130			}
131			return text, unordered, true
132		}
133	}
134	return "", false, false // no suitable comment found
135}
136
137// isTest tells whether name looks like a test, example, fuzz test, or
138// benchmark. It is a Test (say) if there is a character after Test that is not
139// a lower-case letter. (We don't want Testiness.)
140func isTest(name, prefix string) bool {
141	if !strings.HasPrefix(name, prefix) {
142		return false
143	}
144	if len(name) == len(prefix) { // "Test" is ok
145		return true
146	}
147	rune, _ := utf8.DecodeRuneInString(name[len(prefix):])
148	return !unicode.IsLower(rune)
149}
150
151// playExample synthesizes a new *ast.File based on the provided
152// file with the provided function body as the body of main.
153func playExample(file *ast.File, f *ast.FuncDecl) *ast.File {
154	body := f.Body
155
156	if !strings.HasSuffix(file.Name.Name, "_test") {
157		// We don't support examples that are part of the
158		// greater package (yet).
159		return nil
160	}
161
162	// Collect top-level declarations in the file.
163	topDecls := make(map[*ast.Object]ast.Decl)
164	typMethods := make(map[string][]ast.Decl)
165
166	for _, decl := range file.Decls {
167		switch d := decl.(type) {
168		case *ast.FuncDecl:
169			if d.Recv == nil {
170				topDecls[d.Name.Obj] = d
171			} else {
172				if len(d.Recv.List) == 1 {
173					t := d.Recv.List[0].Type
174					tname, _ := baseTypeName(t)
175					typMethods[tname] = append(typMethods[tname], d)
176				}
177			}
178		case *ast.GenDecl:
179			for _, spec := range d.Specs {
180				switch s := spec.(type) {
181				case *ast.TypeSpec:
182					topDecls[s.Name.Obj] = d
183				case *ast.ValueSpec:
184					for _, name := range s.Names {
185						topDecls[name.Obj] = d
186					}
187				}
188			}
189		}
190	}
191
192	// Find unresolved identifiers and uses of top-level declarations.
193	depDecls, unresolved := findDeclsAndUnresolved(body, topDecls, typMethods)
194
195	// Remove predeclared identifiers from unresolved list.
196	for n := range unresolved {
197		if predeclaredTypes[n] || predeclaredConstants[n] || predeclaredFuncs[n] {
198			delete(unresolved, n)
199		}
200	}
201
202	// Use unresolved identifiers to determine the imports used by this
203	// example. The heuristic assumes package names match base import
204	// paths for imports w/o renames (should be good enough most of the time).
205	var namedImports []ast.Spec
206	var blankImports []ast.Spec // _ imports
207
208	// To preserve the blank lines between groups of imports, find the
209	// start position of each group, and assign that position to all
210	// imports from that group.
211	groupStarts := findImportGroupStarts(file.Imports)
212	groupStart := func(s *ast.ImportSpec) token.Pos {
213		for i, start := range groupStarts {
214			if s.Path.ValuePos < start {
215				return groupStarts[i-1]
216			}
217		}
218		return groupStarts[len(groupStarts)-1]
219	}
220
221	for _, s := range file.Imports {
222		p, err := strconv.Unquote(s.Path.Value)
223		if err != nil {
224			continue
225		}
226		if p == "syscall/js" {
227			// We don't support examples that import syscall/js,
228			// because the package syscall/js is not available in the playground.
229			return nil
230		}
231		n := path.Base(p)
232		if s.Name != nil {
233			n = s.Name.Name
234			switch n {
235			case "_":
236				blankImports = append(blankImports, s)
237				continue
238			case ".":
239				// We can't resolve dot imports (yet).
240				return nil
241			}
242		}
243		if unresolved[n] {
244			// Copy the spec and its path to avoid modifying the original.
245			spec := *s
246			path := *s.Path
247			spec.Path = &path
248			spec.Path.ValuePos = groupStart(&spec)
249			namedImports = append(namedImports, &spec)
250			delete(unresolved, n)
251		}
252	}
253
254	// If there are other unresolved identifiers, give up because this
255	// synthesized file is not going to build.
256	if len(unresolved) > 0 {
257		return nil
258	}
259
260	// Include documentation belonging to blank imports.
261	var comments []*ast.CommentGroup
262	for _, s := range blankImports {
263		if c := s.(*ast.ImportSpec).Doc; c != nil {
264			comments = append(comments, c)
265		}
266	}
267
268	// Include comments that are inside the function body.
269	for _, c := range file.Comments {
270		if body.Pos() <= c.Pos() && c.End() <= body.End() {
271			comments = append(comments, c)
272		}
273	}
274
275	// Strip the "Output:" or "Unordered output:" comment and adjust body
276	// end position.
277	body, comments = stripOutputComment(body, comments)
278
279	// Include documentation belonging to dependent declarations.
280	for _, d := range depDecls {
281		switch d := d.(type) {
282		case *ast.GenDecl:
283			if d.Doc != nil {
284				comments = append(comments, d.Doc)
285			}
286		case *ast.FuncDecl:
287			if d.Doc != nil {
288				comments = append(comments, d.Doc)
289			}
290		}
291	}
292
293	// Synthesize import declaration.
294	importDecl := &ast.GenDecl{
295		Tok:    token.IMPORT,
296		Lparen: 1, // Need non-zero Lparen and Rparen so that printer
297		Rparen: 1, // treats this as a factored import.
298	}
299	importDecl.Specs = append(namedImports, blankImports...)
300
301	// Synthesize main function.
302	funcDecl := &ast.FuncDecl{
303		Name: ast.NewIdent("main"),
304		Type: f.Type,
305		Body: body,
306	}
307
308	decls := make([]ast.Decl, 0, 2+len(depDecls))
309	decls = append(decls, importDecl)
310	decls = append(decls, depDecls...)
311	decls = append(decls, funcDecl)
312
313	slices.SortFunc(decls, func(a, b ast.Decl) int {
314		return cmp.Compare(a.Pos(), b.Pos())
315	})
316	slices.SortFunc(comments, func(a, b *ast.CommentGroup) int {
317		return cmp.Compare(a.Pos(), b.Pos())
318	})
319
320	// Synthesize file.
321	return &ast.File{
322		Name:     ast.NewIdent("main"),
323		Decls:    decls,
324		Comments: comments,
325	}
326}
327
328// findDeclsAndUnresolved returns all the top-level declarations mentioned in
329// the body, and a set of unresolved symbols (those that appear in the body but
330// have no declaration in the program).
331//
332// topDecls maps objects to the top-level declaration declaring them (not
333// necessarily obj.Decl, as obj.Decl will be a Spec for GenDecls, but
334// topDecls[obj] will be the GenDecl itself).
335func findDeclsAndUnresolved(body ast.Node, topDecls map[*ast.Object]ast.Decl, typMethods map[string][]ast.Decl) ([]ast.Decl, map[string]bool) {
336	// This function recursively finds every top-level declaration used
337	// transitively by the body, populating usedDecls and usedObjs. Then it
338	// trims down the declarations to include only the symbols actually
339	// referenced by the body.
340
341	unresolved := make(map[string]bool)
342	var depDecls []ast.Decl
343	usedDecls := make(map[ast.Decl]bool)   // set of top-level decls reachable from the body
344	usedObjs := make(map[*ast.Object]bool) // set of objects reachable from the body (each declared by a usedDecl)
345
346	var inspectFunc func(ast.Node) bool
347	inspectFunc = func(n ast.Node) bool {
348		switch e := n.(type) {
349		case *ast.Ident:
350			if e.Obj == nil && e.Name != "_" {
351				unresolved[e.Name] = true
352			} else if d := topDecls[e.Obj]; d != nil {
353
354				usedObjs[e.Obj] = true
355				if !usedDecls[d] {
356					usedDecls[d] = true
357					depDecls = append(depDecls, d)
358				}
359			}
360			return true
361		case *ast.SelectorExpr:
362			// For selector expressions, only inspect the left hand side.
363			// (For an expression like fmt.Println, only add "fmt" to the
364			// set of unresolved names, not "Println".)
365			ast.Inspect(e.X, inspectFunc)
366			return false
367		case *ast.KeyValueExpr:
368			// For key value expressions, only inspect the value
369			// as the key should be resolved by the type of the
370			// composite literal.
371			ast.Inspect(e.Value, inspectFunc)
372			return false
373		}
374		return true
375	}
376
377	inspectFieldList := func(fl *ast.FieldList) {
378		if fl != nil {
379			for _, f := range fl.List {
380				ast.Inspect(f.Type, inspectFunc)
381			}
382		}
383	}
384
385	// Find the decls immediately referenced by body.
386	ast.Inspect(body, inspectFunc)
387	// Now loop over them, adding to the list when we find a new decl that the
388	// body depends on. Keep going until we don't find anything new.
389	for i := 0; i < len(depDecls); i++ {
390		switch d := depDecls[i].(type) {
391		case *ast.FuncDecl:
392			// Inspect type parameters.
393			inspectFieldList(d.Type.TypeParams)
394			// Inspect types of parameters and results. See #28492.
395			inspectFieldList(d.Type.Params)
396			inspectFieldList(d.Type.Results)
397
398			// Functions might not have a body. See #42706.
399			if d.Body != nil {
400				ast.Inspect(d.Body, inspectFunc)
401			}
402		case *ast.GenDecl:
403			for _, spec := range d.Specs {
404				switch s := spec.(type) {
405				case *ast.TypeSpec:
406					inspectFieldList(s.TypeParams)
407					ast.Inspect(s.Type, inspectFunc)
408					depDecls = append(depDecls, typMethods[s.Name.Name]...)
409				case *ast.ValueSpec:
410					if s.Type != nil {
411						ast.Inspect(s.Type, inspectFunc)
412					}
413					for _, val := range s.Values {
414						ast.Inspect(val, inspectFunc)
415					}
416				}
417			}
418		}
419	}
420
421	// Some decls include multiple specs, such as a variable declaration with
422	// multiple variables on the same line, or a parenthesized declaration. Trim
423	// the declarations to include only the specs that are actually mentioned.
424	// However, if there is a constant group with iota, leave it all: later
425	// constant declarations in the group may have no value and so cannot stand
426	// on their own, and removing any constant from the group could change the
427	// values of subsequent ones.
428	// See testdata/examples/iota.go for a minimal example.
429	var ds []ast.Decl
430	for _, d := range depDecls {
431		switch d := d.(type) {
432		case *ast.FuncDecl:
433			ds = append(ds, d)
434		case *ast.GenDecl:
435			containsIota := false // does any spec have iota?
436			// Collect all Specs that were mentioned in the example.
437			var specs []ast.Spec
438			for _, s := range d.Specs {
439				switch s := s.(type) {
440				case *ast.TypeSpec:
441					if usedObjs[s.Name.Obj] {
442						specs = append(specs, s)
443					}
444				case *ast.ValueSpec:
445					if !containsIota {
446						containsIota = hasIota(s)
447					}
448					// A ValueSpec may have multiple names (e.g. "var a, b int").
449					// Keep only the names that were mentioned in the example.
450					// Exception: the multiple names have a single initializer (which
451					// would be a function call with multiple return values). In that
452					// case, keep everything.
453					if len(s.Names) > 1 && len(s.Values) == 1 {
454						specs = append(specs, s)
455						continue
456					}
457					ns := *s
458					ns.Names = nil
459					ns.Values = nil
460					for i, n := range s.Names {
461						if usedObjs[n.Obj] {
462							ns.Names = append(ns.Names, n)
463							if s.Values != nil {
464								ns.Values = append(ns.Values, s.Values[i])
465							}
466						}
467					}
468					if len(ns.Names) > 0 {
469						specs = append(specs, &ns)
470					}
471				}
472			}
473			if len(specs) > 0 {
474				// Constant with iota? Keep it all.
475				if d.Tok == token.CONST && containsIota {
476					ds = append(ds, d)
477				} else {
478					// Synthesize a GenDecl with just the Specs we need.
479					nd := *d // copy the GenDecl
480					nd.Specs = specs
481					if len(specs) == 1 {
482						// Remove grouping parens if there is only one spec.
483						nd.Lparen = 0
484					}
485					ds = append(ds, &nd)
486				}
487			}
488		}
489	}
490	return ds, unresolved
491}
492
493func hasIota(s ast.Spec) bool {
494	has := false
495	ast.Inspect(s, func(n ast.Node) bool {
496		// Check that this is the special built-in "iota" identifier, not
497		// a user-defined shadow.
498		if id, ok := n.(*ast.Ident); ok && id.Name == "iota" && id.Obj == nil {
499			has = true
500			return false
501		}
502		return true
503	})
504	return has
505}
506
507// findImportGroupStarts finds the start positions of each sequence of import
508// specs that are not separated by a blank line.
509func findImportGroupStarts(imps []*ast.ImportSpec) []token.Pos {
510	startImps := findImportGroupStarts1(imps)
511	groupStarts := make([]token.Pos, len(startImps))
512	for i, imp := range startImps {
513		groupStarts[i] = imp.Pos()
514	}
515	return groupStarts
516}
517
518// Helper for findImportGroupStarts to ease testing.
519func findImportGroupStarts1(origImps []*ast.ImportSpec) []*ast.ImportSpec {
520	// Copy to avoid mutation.
521	imps := make([]*ast.ImportSpec, len(origImps))
522	copy(imps, origImps)
523	// Assume the imports are sorted by position.
524	slices.SortFunc(imps, func(a, b *ast.ImportSpec) int {
525		return cmp.Compare(a.Pos(), b.Pos())
526	})
527	// Assume gofmt has been applied, so there is a blank line between adjacent imps
528	// if and only if they are more than 2 positions apart (newline, tab).
529	var groupStarts []*ast.ImportSpec
530	prevEnd := token.Pos(-2)
531	for _, imp := range imps {
532		if imp.Pos()-prevEnd > 2 {
533			groupStarts = append(groupStarts, imp)
534		}
535		prevEnd = imp.End()
536		// Account for end-of-line comments.
537		if imp.Comment != nil {
538			prevEnd = imp.Comment.End()
539		}
540	}
541	return groupStarts
542}
543
544// playExampleFile takes a whole file example and synthesizes a new *ast.File
545// such that the example is function main in package main.
546func playExampleFile(file *ast.File) *ast.File {
547	// Strip copyright comment if present.
548	comments := file.Comments
549	if len(comments) > 0 && strings.HasPrefix(comments[0].Text(), "Copyright") {
550		comments = comments[1:]
551	}
552
553	// Copy declaration slice, rewriting the ExampleX function to main.
554	var decls []ast.Decl
555	for _, d := range file.Decls {
556		if f, ok := d.(*ast.FuncDecl); ok && isTest(f.Name.Name, "Example") {
557			// Copy the FuncDecl, as it may be used elsewhere.
558			newF := *f
559			newF.Name = ast.NewIdent("main")
560			newF.Body, comments = stripOutputComment(f.Body, comments)
561			d = &newF
562		}
563		decls = append(decls, d)
564	}
565
566	// Copy the File, as it may be used elsewhere.
567	f := *file
568	f.Name = ast.NewIdent("main")
569	f.Decls = decls
570	f.Comments = comments
571	return &f
572}
573
574// stripOutputComment finds and removes the "Output:" or "Unordered output:"
575// comment from body and comments, and adjusts the body block's end position.
576func stripOutputComment(body *ast.BlockStmt, comments []*ast.CommentGroup) (*ast.BlockStmt, []*ast.CommentGroup) {
577	// Do nothing if there is no "Output:" or "Unordered output:" comment.
578	i, last := lastComment(body, comments)
579	if last == nil || !outputPrefix.MatchString(last.Text()) {
580		return body, comments
581	}
582
583	// Copy body and comments, as the originals may be used elsewhere.
584	newBody := &ast.BlockStmt{
585		Lbrace: body.Lbrace,
586		List:   body.List,
587		Rbrace: last.Pos(),
588	}
589	newComments := make([]*ast.CommentGroup, len(comments)-1)
590	copy(newComments, comments[:i])
591	copy(newComments[i:], comments[i+1:])
592	return newBody, newComments
593}
594
595// lastComment returns the last comment inside the provided block.
596func lastComment(b *ast.BlockStmt, c []*ast.CommentGroup) (i int, last *ast.CommentGroup) {
597	if b == nil {
598		return
599	}
600	pos, end := b.Pos(), b.End()
601	for j, cg := range c {
602		if cg.Pos() < pos {
603			continue
604		}
605		if cg.End() > end {
606			break
607		}
608		i, last = j, cg
609	}
610	return
611}
612
613// classifyExamples classifies examples and assigns them to the Examples field
614// of the relevant Func, Type, or Package that the example is associated with.
615//
616// The classification process is ambiguous in some cases:
617//
618//   - ExampleFoo_Bar matches a type named Foo_Bar
619//     or a method named Foo.Bar.
620//   - ExampleFoo_bar matches a type named Foo_bar
621//     or Foo (with a "bar" suffix).
622//
623// Examples with malformed names are not associated with anything.
624func classifyExamples(p *Package, examples []*Example) {
625	if len(examples) == 0 {
626		return
627	}
628	// Mapping of names for funcs, types, and methods to the example listing.
629	ids := make(map[string]*[]*Example)
630	ids[""] = &p.Examples // package-level examples have an empty name
631	for _, f := range p.Funcs {
632		if !token.IsExported(f.Name) {
633			continue
634		}
635		ids[f.Name] = &f.Examples
636	}
637	for _, t := range p.Types {
638		if !token.IsExported(t.Name) {
639			continue
640		}
641		ids[t.Name] = &t.Examples
642		for _, f := range t.Funcs {
643			if !token.IsExported(f.Name) {
644				continue
645			}
646			ids[f.Name] = &f.Examples
647		}
648		for _, m := range t.Methods {
649			if !token.IsExported(m.Name) {
650				continue
651			}
652			ids[strings.TrimPrefix(nameWithoutInst(m.Recv), "*")+"_"+m.Name] = &m.Examples
653		}
654	}
655
656	// Group each example with the associated func, type, or method.
657	for _, ex := range examples {
658		// Consider all possible split points for the suffix
659		// by starting at the end of string (no suffix case),
660		// then trying all positions that contain a '_' character.
661		//
662		// An association is made on the first successful match.
663		// Examples with malformed names that match nothing are skipped.
664		for i := len(ex.Name); i >= 0; i = strings.LastIndexByte(ex.Name[:i], '_') {
665			prefix, suffix, ok := splitExampleName(ex.Name, i)
666			if !ok {
667				continue
668			}
669			exs, ok := ids[prefix]
670			if !ok {
671				continue
672			}
673			ex.Suffix = suffix
674			*exs = append(*exs, ex)
675			break
676		}
677	}
678
679	// Sort list of example according to the user-specified suffix name.
680	for _, exs := range ids {
681		slices.SortFunc(*exs, func(a, b *Example) int {
682			return cmp.Compare(a.Suffix, b.Suffix)
683		})
684	}
685}
686
687// nameWithoutInst returns name if name has no brackets. If name contains
688// brackets, then it returns name with all the contents between (and including)
689// the outermost left and right bracket removed.
690//
691// Adapted from debug/gosym/symtab.go:Sym.nameWithoutInst.
692func nameWithoutInst(name string) string {
693	start := strings.Index(name, "[")
694	if start < 0 {
695		return name
696	}
697	end := strings.LastIndex(name, "]")
698	if end < 0 {
699		// Malformed name, should contain closing bracket too.
700		return name
701	}
702	return name[0:start] + name[end+1:]
703}
704
705// splitExampleName attempts to split example name s at index i,
706// and reports if that produces a valid split. The suffix may be
707// absent. Otherwise, it must start with a lower-case letter and
708// be preceded by '_'.
709//
710// One of i == len(s) or s[i] == '_' must be true.
711func splitExampleName(s string, i int) (prefix, suffix string, ok bool) {
712	if i == len(s) {
713		return s, "", true
714	}
715	if i == len(s)-1 {
716		return "", "", false
717	}
718	prefix, suffix = s[:i], s[i+1:]
719	return prefix, suffix, isExampleSuffix(suffix)
720}
721
722func isExampleSuffix(s string) bool {
723	r, size := utf8.DecodeRuneInString(s)
724	return size > 0 && unicode.IsLower(r)
725}
726