1// Copyright 2014 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// Stringer is a tool to automate the creation of methods that satisfy the fmt.Stringer
6// interface. Given the name of a (signed or unsigned) integer type T that has constants
7// defined, stringer will create a new self-contained Go source file implementing
8//	func (t T) String() string
9// The file is created in the same package and directory as the package that defines T.
10// It has helpful defaults designed for use with go generate.
11//
12// Stringer works best with constants that are consecutive values such as created using iota,
13// but creates good code regardless. In the future it might also provide custom support for
14// constant sets that are bit patterns.
15//
16// For example, given this snippet,
17//
18//	package painkiller
19//
20//	type Pill int
21//
22//	const (
23//		Placebo Pill = iota
24//		Aspirin
25//		Ibuprofen
26//		Paracetamol
27//		Acetaminophen = Paracetamol
28//	)
29//
30// running this command
31//
32//	stringer -type=Pill
33//
34// in the same directory will create the file pill_string.go, in package painkiller,
35// containing a definition of
36//
37//	func (Pill) String() string
38//
39// That method will translate the value of a Pill constant to the string representation
40// of the respective constant name, so that the call fmt.Print(painkiller.Aspirin) will
41// print the string "Aspirin".
42//
43// Typically this process would be run using go generate, like this:
44//
45//	//go:generate stringer -type=Pill
46//
47// If multiple constants have the same value, the lexically first matching name will
48// be used (in the example, Acetaminophen will print as "Paracetamol").
49//
50// With no arguments, it processes the package in the current directory.
51// Otherwise, the arguments must name a single directory holding a Go package
52// or a set of Go source files that represent a single Go package.
53//
54// The -type flag accepts a comma-separated list of types so a single run can
55// generate methods for multiple types. The default output file is t_string.go,
56// where t is the lower-cased name of the first type listed. It can be overridden
57// with the -output flag.
58//
59package main // import "golang.org/x/tools/cmd/stringer"
60
61import (
62	"bytes"
63	"flag"
64	"fmt"
65	"go/ast"
66	"go/build"
67	"go/constant"
68	"go/format"
69	"go/parser"
70	"go/token"
71	"go/types"
72	"io/ioutil"
73	"log"
74	"os"
75	"path/filepath"
76	"sort"
77	"strings"
78)
79
80var (
81	typeNames   = flag.String("type", "", "comma-separated list of type names; must be set")
82	output      = flag.String("output", "", "output file name; default srcdir/<type>_string.go")
83	trimprefix  = flag.String("trimprefix", "", "trim the `prefix` from the generated constant names")
84	linecomment = flag.Bool("linecomment", false, "use line comment text as printed text when present")
85	buildTags   = flag.String("tags", "", "comma-separated list of build tags to apply")
86)
87
88// Usage is a replacement usage function for the flags package.
89func Usage() {
90	fmt.Fprintf(os.Stderr, "Usage of stringer:\n")
91	fmt.Fprintf(os.Stderr, "\tstringer [flags] -type T [directory]\n")
92	fmt.Fprintf(os.Stderr, "\tstringer [flags] -type T files... # Must be a single package\n")
93	fmt.Fprintf(os.Stderr, "For more information, see:\n")
94	fmt.Fprintf(os.Stderr, "\thttp://godoc.org/golang.org/x/tools/cmd/stringer\n")
95	fmt.Fprintf(os.Stderr, "Flags:\n")
96	flag.PrintDefaults()
97}
98
99func main() {
100	log.SetFlags(0)
101	log.SetPrefix("stringer: ")
102	flag.Usage = Usage
103	flag.Parse()
104	if len(*typeNames) == 0 {
105		flag.Usage()
106		os.Exit(2)
107	}
108	types := strings.Split(*typeNames, ",")
109	var tags []string
110	if len(*buildTags) > 0 {
111		tags = strings.Split(*buildTags, ",")
112	}
113
114	// We accept either one directory or a list of files. Which do we have?
115	args := flag.Args()
116	if len(args) == 0 {
117		// Default: process whole package in current directory.
118		args = []string{"."}
119	}
120
121	// Parse the package once.
122	var dir string
123	g := Generator{
124		trimPrefix:  *trimprefix,
125		lineComment: *linecomment,
126	}
127	if len(args) == 1 && isDirectory(args[0]) {
128		dir = args[0]
129		g.parsePackageDir(args[0], tags)
130	} else {
131		if len(tags) != 0 {
132			log.Fatal("-tags option applies only to directories, not when files are specified")
133		}
134		dir = filepath.Dir(args[0])
135		g.parsePackageFiles(args)
136	}
137
138	// Print the header and package clause.
139	g.Printf("// Code generated by \"stringer %s\"; DO NOT EDIT.\n", strings.Join(os.Args[1:], " "))
140	g.Printf("\n")
141	g.Printf("package %s", g.pkg.name)
142	g.Printf("\n")
143	g.Printf("import \"strconv\"\n") // Used by all methods.
144
145	// Run generate for each type.
146	for _, typeName := range types {
147		g.generate(typeName)
148	}
149
150	// Format the output.
151	src := g.format()
152
153	// Write to file.
154	outputName := *output
155	if outputName == "" {
156		baseName := fmt.Sprintf("%s_string.go", types[0])
157		outputName = filepath.Join(dir, strings.ToLower(baseName))
158	}
159	err := ioutil.WriteFile(outputName, src, 0644)
160	if err != nil {
161		log.Fatalf("writing output: %s", err)
162	}
163}
164
165// isDirectory reports whether the named file is a directory.
166func isDirectory(name string) bool {
167	info, err := os.Stat(name)
168	if err != nil {
169		log.Fatal(err)
170	}
171	return info.IsDir()
172}
173
174// Generator holds the state of the analysis. Primarily used to buffer
175// the output for format.Source.
176type Generator struct {
177	buf bytes.Buffer // Accumulated output.
178	pkg *Package     // Package we are scanning.
179
180	trimPrefix  string
181	lineComment bool
182}
183
184func (g *Generator) Printf(format string, args ...interface{}) {
185	fmt.Fprintf(&g.buf, format, args...)
186}
187
188// File holds a single parsed file and associated data.
189type File struct {
190	pkg  *Package  // Package to which this file belongs.
191	file *ast.File // Parsed AST.
192	// These fields are reset for each type being generated.
193	typeName string  // Name of the constant type.
194	values   []Value // Accumulator for constant values of that type.
195
196	trimPrefix  string
197	lineComment bool
198}
199
200type Package struct {
201	dir      string
202	name     string
203	defs     map[*ast.Ident]types.Object
204	files    []*File
205	typesPkg *types.Package
206}
207
208func buildContext(tags []string) *build.Context {
209	ctx := build.Default
210	ctx.BuildTags = tags
211	return &ctx
212}
213
214// parsePackageDir parses the package residing in the directory.
215func (g *Generator) parsePackageDir(directory string, tags []string) {
216	pkg, err := buildContext(tags).ImportDir(directory, 0)
217	if err != nil {
218		log.Fatalf("cannot process directory %s: %s", directory, err)
219	}
220	var names []string
221	names = append(names, pkg.GoFiles...)
222	names = append(names, pkg.CgoFiles...)
223	// TODO: Need to think about constants in test files. Maybe write type_string_test.go
224	// in a separate pass? For later.
225	// names = append(names, pkg.TestGoFiles...) // These are also in the "foo" package.
226	names = append(names, pkg.SFiles...)
227	names = prefixDirectory(directory, names)
228	g.parsePackage(directory, names, nil)
229}
230
231// parsePackageFiles parses the package occupying the named files.
232func (g *Generator) parsePackageFiles(names []string) {
233	g.parsePackage(".", names, nil)
234}
235
236// prefixDirectory places the directory name on the beginning of each name in the list.
237func prefixDirectory(directory string, names []string) []string {
238	if directory == "." {
239		return names
240	}
241	ret := make([]string, len(names))
242	for i, name := range names {
243		ret[i] = filepath.Join(directory, name)
244	}
245	return ret
246}
247
248// parsePackage analyzes the single package constructed from the named files.
249// If text is non-nil, it is a string to be used instead of the content of the file,
250// to be used for testing. parsePackage exits if there is an error.
251func (g *Generator) parsePackage(directory string, names []string, text interface{}) {
252	var files []*File
253	var astFiles []*ast.File
254	g.pkg = new(Package)
255	fs := token.NewFileSet()
256	for _, name := range names {
257		if !strings.HasSuffix(name, ".go") {
258			continue
259		}
260		parsedFile, err := parser.ParseFile(fs, name, text, parser.ParseComments)
261		if err != nil {
262			log.Fatalf("parsing package: %s: %s", name, err)
263		}
264		astFiles = append(astFiles, parsedFile)
265		files = append(files, &File{
266			file:        parsedFile,
267			pkg:         g.pkg,
268			trimPrefix:  g.trimPrefix,
269			lineComment: g.lineComment,
270		})
271	}
272	if len(astFiles) == 0 {
273		log.Fatalf("%s: no buildable Go files", directory)
274	}
275	g.pkg.name = astFiles[0].Name.Name
276	g.pkg.files = files
277	g.pkg.dir = directory
278	g.pkg.typeCheck(fs, astFiles)
279}
280
281// check type-checks the package so we can evaluate contants whose values we are printing.
282func (pkg *Package) typeCheck(fs *token.FileSet, astFiles []*ast.File) {
283	pkg.defs = make(map[*ast.Ident]types.Object)
284	config := types.Config{
285		IgnoreFuncBodies: true, // We only need to evaluate constants.
286		Importer:         defaultImporter(),
287		FakeImportC:      true,
288	}
289	info := &types.Info{
290		Defs: pkg.defs,
291	}
292	typesPkg, err := config.Check(pkg.dir, fs, astFiles, info)
293	if err != nil {
294		log.Fatalf("checking package: %s", err)
295	}
296	pkg.typesPkg = typesPkg
297}
298
299// generate produces the String method for the named type.
300func (g *Generator) generate(typeName string) {
301	values := make([]Value, 0, 100)
302	for _, file := range g.pkg.files {
303		// Set the state for this run of the walker.
304		file.typeName = typeName
305		file.values = nil
306		if file.file != nil {
307			ast.Inspect(file.file, file.genDecl)
308			values = append(values, file.values...)
309		}
310	}
311
312	if len(values) == 0 {
313		log.Fatalf("no values defined for type %s", typeName)
314	}
315	runs := splitIntoRuns(values)
316	// The decision of which pattern to use depends on the number of
317	// runs in the numbers. If there's only one, it's easy. For more than
318	// one, there's a tradeoff between complexity and size of the data
319	// and code vs. the simplicity of a map. A map takes more space,
320	// but so does the code. The decision here (crossover at 10) is
321	// arbitrary, but considers that for large numbers of runs the cost
322	// of the linear scan in the switch might become important, and
323	// rather than use yet another algorithm such as binary search,
324	// we punt and use a map. In any case, the likelihood of a map
325	// being necessary for any realistic example other than bitmasks
326	// is very low. And bitmasks probably deserve their own analysis,
327	// to be done some other day.
328	switch {
329	case len(runs) == 1:
330		g.buildOneRun(runs, typeName)
331	case len(runs) <= 10:
332		g.buildMultipleRuns(runs, typeName)
333	default:
334		g.buildMap(runs, typeName)
335	}
336}
337
338// splitIntoRuns breaks the values into runs of contiguous sequences.
339// For example, given 1,2,3,5,6,7 it returns {1,2,3},{5,6,7}.
340// The input slice is known to be non-empty.
341func splitIntoRuns(values []Value) [][]Value {
342	// We use stable sort so the lexically first name is chosen for equal elements.
343	sort.Stable(byValue(values))
344	// Remove duplicates. Stable sort has put the one we want to print first,
345	// so use that one. The String method won't care about which named constant
346	// was the argument, so the first name for the given value is the only one to keep.
347	// We need to do this because identical values would cause the switch or map
348	// to fail to compile.
349	j := 1
350	for i := 1; i < len(values); i++ {
351		if values[i].value != values[i-1].value {
352			values[j] = values[i]
353			j++
354		}
355	}
356	values = values[:j]
357	runs := make([][]Value, 0, 10)
358	for len(values) > 0 {
359		// One contiguous sequence per outer loop.
360		i := 1
361		for i < len(values) && values[i].value == values[i-1].value+1 {
362			i++
363		}
364		runs = append(runs, values[:i])
365		values = values[i:]
366	}
367	return runs
368}
369
370// format returns the gofmt-ed contents of the Generator's buffer.
371func (g *Generator) format() []byte {
372	src, err := format.Source(g.buf.Bytes())
373	if err != nil {
374		// Should never happen, but can arise when developing this code.
375		// The user can compile the output to see the error.
376		log.Printf("warning: internal error: invalid Go generated: %s", err)
377		log.Printf("warning: compile the package to analyze the error")
378		return g.buf.Bytes()
379	}
380	return src
381}
382
383// Value represents a declared constant.
384type Value struct {
385	name string // The name of the constant.
386	// The value is stored as a bit pattern alone. The boolean tells us
387	// whether to interpret it as an int64 or a uint64; the only place
388	// this matters is when sorting.
389	// Much of the time the str field is all we need; it is printed
390	// by Value.String.
391	value  uint64 // Will be converted to int64 when needed.
392	signed bool   // Whether the constant is a signed type.
393	str    string // The string representation given by the "go/constant" package.
394}
395
396func (v *Value) String() string {
397	return v.str
398}
399
400// byValue lets us sort the constants into increasing order.
401// We take care in the Less method to sort in signed or unsigned order,
402// as appropriate.
403type byValue []Value
404
405func (b byValue) Len() int      { return len(b) }
406func (b byValue) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
407func (b byValue) Less(i, j int) bool {
408	if b[i].signed {
409		return int64(b[i].value) < int64(b[j].value)
410	}
411	return b[i].value < b[j].value
412}
413
414// genDecl processes one declaration clause.
415func (f *File) genDecl(node ast.Node) bool {
416	decl, ok := node.(*ast.GenDecl)
417	if !ok || decl.Tok != token.CONST {
418		// We only care about const declarations.
419		return true
420	}
421	// The name of the type of the constants we are declaring.
422	// Can change if this is a multi-element declaration.
423	typ := ""
424	// Loop over the elements of the declaration. Each element is a ValueSpec:
425	// a list of names possibly followed by a type, possibly followed by values.
426	// If the type and value are both missing, we carry down the type (and value,
427	// but the "go/types" package takes care of that).
428	for _, spec := range decl.Specs {
429		vspec := spec.(*ast.ValueSpec) // Guaranteed to succeed as this is CONST.
430		if vspec.Type == nil && len(vspec.Values) > 0 {
431			// "X = 1". With no type but a value. If the constant is untyped,
432			// skip this vspec and reset the remembered type.
433			typ = ""
434
435			// If this is a simple type conversion, remember the type.
436			// We don't mind if this is actually a call; a qualified call won't
437			// be matched (that will be SelectorExpr, not Ident), and only unusual
438			// situations will result in a function call that appears to be
439			// a type conversion.
440			ce, ok := vspec.Values[0].(*ast.CallExpr)
441			if !ok {
442				continue
443			}
444			id, ok := ce.Fun.(*ast.Ident)
445			if !ok {
446				continue
447			}
448			typ = id.Name
449		}
450		if vspec.Type != nil {
451			// "X T". We have a type. Remember it.
452			ident, ok := vspec.Type.(*ast.Ident)
453			if !ok {
454				continue
455			}
456			typ = ident.Name
457		}
458		if typ != f.typeName {
459			// This is not the type we're looking for.
460			continue
461		}
462		// We now have a list of names (from one line of source code) all being
463		// declared with the desired type.
464		// Grab their names and actual values and store them in f.values.
465		for _, name := range vspec.Names {
466			if name.Name == "_" {
467				continue
468			}
469			// This dance lets the type checker find the values for us. It's a
470			// bit tricky: look up the object declared by the name, find its
471			// types.Const, and extract its value.
472			obj, ok := f.pkg.defs[name]
473			if !ok {
474				log.Fatalf("no value for constant %s", name)
475			}
476			info := obj.Type().Underlying().(*types.Basic).Info()
477			if info&types.IsInteger == 0 {
478				log.Fatalf("can't handle non-integer constant type %s", typ)
479			}
480			value := obj.(*types.Const).Val() // Guaranteed to succeed as this is CONST.
481			if value.Kind() != constant.Int {
482				log.Fatalf("can't happen: constant is not an integer %s", name)
483			}
484			i64, isInt := constant.Int64Val(value)
485			u64, isUint := constant.Uint64Val(value)
486			if !isInt && !isUint {
487				log.Fatalf("internal error: value of %s is not an integer: %s", name, value.String())
488			}
489			if !isInt {
490				u64 = uint64(i64)
491			}
492			v := Value{
493				name:   name.Name,
494				value:  u64,
495				signed: info&types.IsUnsigned == 0,
496				str:    value.String(),
497			}
498			if c := vspec.Comment; f.lineComment && c != nil && len(c.List) == 1 {
499				v.name = strings.TrimSpace(c.Text())
500			}
501			v.name = strings.TrimPrefix(v.name, f.trimPrefix)
502			f.values = append(f.values, v)
503		}
504	}
505	return false
506}
507
508// Helpers
509
510// usize returns the number of bits of the smallest unsigned integer
511// type that will hold n. Used to create the smallest possible slice of
512// integers to use as indexes into the concatenated strings.
513func usize(n int) int {
514	switch {
515	case n < 1<<8:
516		return 8
517	case n < 1<<16:
518		return 16
519	default:
520		// 2^32 is enough constants for anyone.
521		return 32
522	}
523}
524
525// declareIndexAndNameVars declares the index slices and concatenated names
526// strings representing the runs of values.
527func (g *Generator) declareIndexAndNameVars(runs [][]Value, typeName string) {
528	var indexes, names []string
529	for i, run := range runs {
530		index, name := g.createIndexAndNameDecl(run, typeName, fmt.Sprintf("_%d", i))
531		if len(run) != 1 {
532			indexes = append(indexes, index)
533		}
534		names = append(names, name)
535	}
536	g.Printf("const (\n")
537	for _, name := range names {
538		g.Printf("\t%s\n", name)
539	}
540	g.Printf(")\n\n")
541
542	if len(indexes) > 0 {
543		g.Printf("var (")
544		for _, index := range indexes {
545			g.Printf("\t%s\n", index)
546		}
547		g.Printf(")\n\n")
548	}
549}
550
551// declareIndexAndNameVar is the single-run version of declareIndexAndNameVars
552func (g *Generator) declareIndexAndNameVar(run []Value, typeName string) {
553	index, name := g.createIndexAndNameDecl(run, typeName, "")
554	g.Printf("const %s\n", name)
555	g.Printf("var %s\n", index)
556}
557
558// createIndexAndNameDecl returns the pair of declarations for the run. The caller will add "const" and "var".
559func (g *Generator) createIndexAndNameDecl(run []Value, typeName string, suffix string) (string, string) {
560	b := new(bytes.Buffer)
561	indexes := make([]int, len(run))
562	for i := range run {
563		b.WriteString(run[i].name)
564		indexes[i] = b.Len()
565	}
566	nameConst := fmt.Sprintf("_%s_name%s = %q", typeName, suffix, b.String())
567	nameLen := b.Len()
568	b.Reset()
569	fmt.Fprintf(b, "_%s_index%s = [...]uint%d{0, ", typeName, suffix, usize(nameLen))
570	for i, v := range indexes {
571		if i > 0 {
572			fmt.Fprintf(b, ", ")
573		}
574		fmt.Fprintf(b, "%d", v)
575	}
576	fmt.Fprintf(b, "}")
577	return b.String(), nameConst
578}
579
580// declareNameVars declares the concatenated names string representing all the values in the runs.
581func (g *Generator) declareNameVars(runs [][]Value, typeName string, suffix string) {
582	g.Printf("const _%s_name%s = \"", typeName, suffix)
583	for _, run := range runs {
584		for i := range run {
585			g.Printf("%s", run[i].name)
586		}
587	}
588	g.Printf("\"\n")
589}
590
591// buildOneRun generates the variables and String method for a single run of contiguous values.
592func (g *Generator) buildOneRun(runs [][]Value, typeName string) {
593	values := runs[0]
594	g.Printf("\n")
595	g.declareIndexAndNameVar(values, typeName)
596	// The generated code is simple enough to write as a Printf format.
597	lessThanZero := ""
598	if values[0].signed {
599		lessThanZero = "i < 0 || "
600	}
601	if values[0].value == 0 { // Signed or unsigned, 0 is still 0.
602		g.Printf(stringOneRun, typeName, usize(len(values)), lessThanZero)
603	} else {
604		g.Printf(stringOneRunWithOffset, typeName, values[0].String(), usize(len(values)), lessThanZero)
605	}
606}
607
608// Arguments to format are:
609//	[1]: type name
610//	[2]: size of index element (8 for uint8 etc.)
611//	[3]: less than zero check (for signed types)
612const stringOneRun = `func (i %[1]s) String() string {
613	if %[3]si >= %[1]s(len(_%[1]s_index)-1) {
614		return "%[1]s(" + strconv.FormatInt(int64(i), 10) + ")"
615	}
616	return _%[1]s_name[_%[1]s_index[i]:_%[1]s_index[i+1]]
617}
618`
619
620// Arguments to format are:
621//	[1]: type name
622//	[2]: lowest defined value for type, as a string
623//	[3]: size of index element (8 for uint8 etc.)
624//	[4]: less than zero check (for signed types)
625/*
626 */
627const stringOneRunWithOffset = `func (i %[1]s) String() string {
628	i -= %[2]s
629	if %[4]si >= %[1]s(len(_%[1]s_index)-1) {
630		return "%[1]s(" + strconv.FormatInt(int64(i + %[2]s), 10) + ")"
631	}
632	return _%[1]s_name[_%[1]s_index[i] : _%[1]s_index[i+1]]
633}
634`
635
636// buildMultipleRuns generates the variables and String method for multiple runs of contiguous values.
637// For this pattern, a single Printf format won't do.
638func (g *Generator) buildMultipleRuns(runs [][]Value, typeName string) {
639	g.Printf("\n")
640	g.declareIndexAndNameVars(runs, typeName)
641	g.Printf("func (i %s) String() string {\n", typeName)
642	g.Printf("\tswitch {\n")
643	for i, values := range runs {
644		if len(values) == 1 {
645			g.Printf("\tcase i == %s:\n", &values[0])
646			g.Printf("\t\treturn _%s_name_%d\n", typeName, i)
647			continue
648		}
649		g.Printf("\tcase %s <= i && i <= %s:\n", &values[0], &values[len(values)-1])
650		if values[0].value != 0 {
651			g.Printf("\t\ti -= %s\n", &values[0])
652		}
653		g.Printf("\t\treturn _%s_name_%d[_%s_index_%d[i]:_%s_index_%d[i+1]]\n",
654			typeName, i, typeName, i, typeName, i)
655	}
656	g.Printf("\tdefault:\n")
657	g.Printf("\t\treturn \"%s(\" + strconv.FormatInt(int64(i), 10) + \")\"\n", typeName)
658	g.Printf("\t}\n")
659	g.Printf("}\n")
660}
661
662// buildMap handles the case where the space is so sparse a map is a reasonable fallback.
663// It's a rare situation but has simple code.
664func (g *Generator) buildMap(runs [][]Value, typeName string) {
665	g.Printf("\n")
666	g.declareNameVars(runs, typeName, "")
667	g.Printf("\nvar _%s_map = map[%s]string{\n", typeName, typeName)
668	n := 0
669	for _, values := range runs {
670		for _, value := range values {
671			g.Printf("\t%s: _%s_name[%d:%d],\n", &value, typeName, n, n+len(value.name))
672			n += len(value.name)
673		}
674	}
675	g.Printf("}\n\n")
676	g.Printf(stringMap, typeName)
677}
678
679// Argument to format is the type name.
680const stringMap = `func (i %[1]s) String() string {
681	if str, ok := _%[1]s_map[i]; ok {
682		return str
683	}
684	return "%[1]s(" + strconv.FormatInt(int64(i), 10) + ")"
685}
686`
687