gentools.go

     1  // Copyright 2026 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  // Package gentools provides shared helper utilities for Go code generator tools
     6  // in archsimd.
     7  //
     8  // Basic usage:
     9  //
    10  //	func main() {
    11  //	    gentools.RegisterFlags(nil)
    12  //	    flag.Parse()
    13  //
    14  //	    var files gentools.Files
    15  //	    defer files.FlushOrExit()
    16  //
    17  //	    buf := files.NewGoFile("src/simd/archsimd/ops_amd64.go")
    18  //	    fmt.Fprintln(buf, "package archsimd")
    19  //	    // ... write generated code to buf ...
    20  //	}
    21  //
    22  // By default (when -w is not specified), gentools outputs all generated files
    23  // as a txtar archive to standard output. Pass -w to write files directly into
    24  // the Go source tree.
    25  package gentools
    26  
    27  import (
    28  	"bytes"
    29  	"flag"
    30  	"fmt"
    31  	"go/format"
    32  	"go/scanner"
    33  	"go/token"
    34  	"internal/diff"
    35  	"io"
    36  	"io/fs"
    37  	"os"
    38  	"path/filepath"
    39  	"strings"
    40  	"sync"
    41  )
    42  
    43  // Options contains standard options and CLI flags for code generators.
    44  type Options struct {
    45  	GOROOT string // -goroot: root of the input Go source tree
    46  	outDir string // -outdir: root of the output tree (defaults to GOROOT)
    47  	Write  bool   // -w: write generated files to disk under GOROOT
    48  	Diff   bool   // -diff: check if generated files match disk, print diffs if not
    49  	Txtar  bool   // -txtar: write generated files to output as a txtar archive (default output mode)
    50  
    51  	Output    io.Writer // output writer for txtar and diff mode; defaults to os.Stdout if nil
    52  	ErrOutput io.Writer // error writer for formatting errors; defaults to os.Stderr if nil
    53  }
    54  
    55  var globalOptions *Options
    56  
    57  // RegisterFlags registers standard generator flags with the provided FlagSet
    58  // (or [flag.CommandLine] if fs is nil) and returns a pointer to the Options
    59  // struct.
    60  //
    61  // If fs is nil, the returned options are remembered globally as defaults for
    62  // zero-value Files instances.
    63  func RegisterFlags(fs *flag.FlagSet) *Options {
    64  	o := new(Options)
    65  	if fs == nil {
    66  		fs = flag.CommandLine
    67  		globalOptions = o
    68  	}
    69  	defaultGOROOT := findGOROOT()
    70  	fs.StringVar(&o.GOROOT, "goroot", defaultGOROOT, "source Go dev tree")
    71  	fs.StringVar(&o.outDir, "outdir", "", "output directory (default: set to -goroot)")
    72  	fs.BoolVar(&o.Write, "w", false, "write generated files directly to disk under -outdir")
    73  	fs.BoolVar(&o.Diff, "diff", false, "compare generated files against disk and print unified diffs")
    74  	fs.BoolVar(&o.Txtar, "txtar", false, "output generated files as a txtar archive to stdout (default mode)")
    75  	return o
    76  }
    77  
    78  // InputPath resolves relPath relative to either o.OutDir/src, if that file
    79  // exists, or o.GOROOT/src. In effect, o.OutDir is treated as an overlay on
    80  // o.GOROOT.
    81  func (o *Options) InputPath(relPath string) string {
    82  	if o.outDir != o.GOROOT {
    83  		path := o.OutputPath(relPath)
    84  		if _, err := os.Stat(path); err == nil {
    85  			return path
    86  		}
    87  	}
    88  	return filepath.Join(o.GOROOT, "src", relPath)
    89  }
    90  
    91  // ReadFile reads relPath from either o.OutDir/src or o.GOROOT/src.
    92  func (o *Options) ReadFile(relPath string) ([]byte, error) {
    93  	return os.ReadFile(o.InputPath(relPath))
    94  }
    95  
    96  // OutputPath returns relPath relative to o.OutDir/src.
    97  func (o *Options) OutputPath(relPath string) string {
    98  	outDir := o.outDir
    99  	if outDir == "" {
   100  		outDir = o.GOROOT
   101  	}
   102  	return filepath.Join(outDir, "src", relPath)
   103  }
   104  
   105  // WritingToInput returns true if Flush will write to the input tree.
   106  func (o *Options) WritingToInput() bool {
   107  	return o.Write && (o.outDir == "" || o.outDir == o.GOROOT)
   108  }
   109  
   110  type fileInfo struct {
   111  	relPath string
   112  	isGo    bool
   113  	buf     bytes.Buffer
   114  }
   115  
   116  // Files manages a collection of generated files for a single generator run.
   117  // The zero value of Files is ready for immediate use and automatically honors
   118  // the command-line flags registered via RegisterFlags.
   119  type Files struct {
   120  	// Options optionally overrides the generator options for this Files instance.
   121  	// If nil, the globally registered options from RegisterFlags are used automatically.
   122  	Options *Options
   123  
   124  	files []*fileInfo
   125  
   126  	// tmpDir is a temporary directory used for communicating with subprocess
   127  	// gentools.
   128  	tmpDirOnce sync.Once
   129  	tmpDir     string
   130  }
   131  
   132  func (f *Files) getOptions() Options {
   133  	var opts Options
   134  	if f != nil && f.Options != nil {
   135  		opts = *f.Options
   136  	} else if globalOptions != nil {
   137  		opts = *globalOptions
   138  	}
   139  
   140  	if opts.GOROOT == "" {
   141  		opts.GOROOT = findGOROOT()
   142  	}
   143  	if opts.Output == nil {
   144  		opts.Output = os.Stdout
   145  	}
   146  	if opts.ErrOutput == nil {
   147  		opts.ErrOutput = os.Stderr
   148  	}
   149  	if !(opts.Write || opts.Diff || opts.Txtar) {
   150  		opts.Txtar = true
   151  	}
   152  
   153  	return opts
   154  }
   155  
   156  // NewGoFile registers a Go source file at relPath (relative to GOROOT/src). It
   157  // returns a *bytes.Buffer for the generator to populate. During Flush(), Go
   158  // files are formatted with go/format.
   159  func (f *Files) NewGoFile(relPath string) *bytes.Buffer {
   160  	info := &fileInfo{
   161  		relPath: relPath,
   162  		isGo:    true,
   163  	}
   164  	f.files = append(f.files, info)
   165  	return &info.buf
   166  }
   167  
   168  // NewRawFile registers a non-Go file (e.g. .rules, YAML, txtar) at relPath
   169  // (relative to GOROOT/src). It returns a *bytes.Buffer for the generator to
   170  // populate. During Flush(), content is written directly without go/format.
   171  func (f *Files) NewRawFile(relPath string) *bytes.Buffer {
   172  	info := &fileInfo{
   173  		relPath: relPath,
   174  		isGo:    false,
   175  	}
   176  	f.files = append(f.files, info)
   177  	return &info.buf
   178  }
   179  
   180  // ExecFlags returns a sequence of flags that can be passed to a gentools
   181  // subprocess. This allows several gentools to be tied together by a larger
   182  // gentool, including if later gentools read the outputs of earlier gentools.
   183  //
   184  // Regardless of the output mode of f, this directs subprocesses to write to a
   185  // temporary directory. Flush then reads the contents of this temporary
   186  // directory back as if this process had written all of those files using f and
   187  // applies the configured output mode.
   188  func (f *Files) ExecFlags() []string {
   189  	f.tmpDirOnce.Do(func() {
   190  		tmpDir, err := os.MkdirTemp("", "")
   191  		if err != nil {
   192  			panic("failed to create tmpdir: " + err.Error())
   193  		}
   194  		f.tmpDir = tmpDir
   195  	})
   196  	return []string{"-goroot", f.getOptions().GOROOT, "-w", "-outdir", f.tmpDir}
   197  }
   198  
   199  // Flush outputs all registered files according to the mode in options.
   200  //
   201  // In default / -txtar mode, it outputs files as a txtar archive to Output. In
   202  // write mode (-w), it writes all files to disk under GOROOT. In diff mode
   203  // (-diff), it compares generated content against disk, prints diffs to Output,
   204  // and returns an error if out of date.
   205  func (f *Files) Flush() error {
   206  	opts := f.getOptions()
   207  
   208  	if (opts.Write || opts.Diff) && opts.GOROOT == "" {
   209  		return fmt.Errorf("GOROOT not found; pass -goroot flag")
   210  	}
   211  
   212  	type preparedFile struct {
   213  		relPath string
   214  		content []byte
   215  	}
   216  
   217  	prepared := make([]preparedFile, len(f.files))
   218  
   219  	// If we invoked subprocesses, read their output files.
   220  	if f.tmpDir != "" {
   221  		root := filepath.Join(f.tmpDir, "src")
   222  		err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
   223  			if d.IsDir() {
   224  				return nil
   225  			}
   226  			relPath, ok := strings.CutPrefix(path, root)
   227  			if !ok {
   228  				return fmt.Errorf("expected path %q to start with root %q", path, root)
   229  			}
   230  			content, err := os.ReadFile(path)
   231  			if err != nil {
   232  				return err
   233  			}
   234  			prepared = append(prepared, preparedFile{relPath, content})
   235  			return nil
   236  		})
   237  		if err != nil {
   238  			return err
   239  		}
   240  		os.RemoveAll(f.tmpDir)
   241  	}
   242  
   243  	for i, fi := range f.files {
   244  		raw := fi.buf.Bytes()
   245  		var content []byte
   246  		if fi.isGo {
   247  			formatted, err := format.Source(raw)
   248  			if err != nil {
   249  				printFormattingError(opts.ErrOutput, fi.relPath, raw, err)
   250  				return fmt.Errorf("error formatting %s: %w", fi.relPath, err)
   251  			}
   252  			content = formatted
   253  		} else {
   254  			content = raw
   255  		}
   256  
   257  		prepared[i] = preparedFile{
   258  			relPath: fi.relPath,
   259  			content: content,
   260  		}
   261  	}
   262  	f.files = nil
   263  
   264  	if opts.Diff {
   265  		hasDiffs := false
   266  		for _, pf := range prepared {
   267  			onDisk, err := opts.ReadFile(pf.relPath)
   268  			if err != nil && !os.IsNotExist(err) {
   269  				return fmt.Errorf("reading %s for diff: %w", pf.relPath, err)
   270  			}
   271  			srcPath := filepath.Join("src", pf.relPath)
   272  			d := diff.Diff(srcPath, onDisk, srcPath, pf.content)
   273  			if len(d) > 0 {
   274  				hasDiffs = true
   275  				opts.Output.Write(d)
   276  			}
   277  		}
   278  		if hasDiffs {
   279  			return fmt.Errorf("generated files differ from disk")
   280  		}
   281  	}
   282  
   283  	if opts.Txtar {
   284  		for i, pf := range prepared {
   285  			if i > 0 {
   286  				fmt.Fprintln(opts.Output)
   287  			}
   288  			srcPath := filepath.Join("src", pf.relPath)
   289  			fmt.Fprintf(opts.Output, "-- %s --\n", srcPath)
   290  			opts.Output.Write(pf.content)
   291  			// Ensure trailing \n
   292  			if len(pf.content) > 0 && !bytes.HasSuffix(pf.content, []byte("\n")) {
   293  				fmt.Fprintln(opts.Output)
   294  			}
   295  		}
   296  	}
   297  
   298  	if opts.Write {
   299  		for _, pf := range prepared {
   300  			path := opts.OutputPath(pf.relPath)
   301  			dir := filepath.Dir(path)
   302  			if err := os.MkdirAll(dir, 0755); err != nil {
   303  				return fmt.Errorf("creating directory %s: %w", dir, err)
   304  			}
   305  			if err := os.WriteFile(path, pf.content, 0644); err != nil {
   306  				return fmt.Errorf("writing %s: %w", path, err)
   307  			}
   308  		}
   309  	}
   310  
   311  	return nil
   312  }
   313  
   314  // FlushOrExit calls Flush(), prints any error to stderr, and exits with code 1 if Flush fails.
   315  //
   316  // It is intended to be deferred at the beginning of main (e.g., `defer files.FlushOrExit()`).
   317  // Hence, if invoked as part of a panic, it skips flushing and instead allows the panic to propagate.
   318  func (f *Files) FlushOrExit() {
   319  	if r := recover(); r != nil {
   320  		panic(r)
   321  	}
   322  	if err := f.Flush(); err != nil {
   323  		fmt.Fprintf(os.Stderr, "%v\n", err)
   324  		os.Exit(1)
   325  	}
   326  }
   327  
   328  // printFormattingError prints err, with 10 lines of context around the error
   329  // line and a caret mark ("^") to indicate the column offset of the error.
   330  func printFormattingError(out io.Writer, relPath string, raw []byte, err error) {
   331  	var pos token.Position
   332  	if el, ok := err.(scanner.ErrorList); ok && len(el) > 0 {
   333  		el.Sort()
   334  		pos = el[0].Pos
   335  	} else if e, ok := err.(*scanner.Error); ok {
   336  		pos = e.Pos
   337  	} else if e, ok := err.(scanner.Error); ok {
   338  		pos = e.Pos
   339  	}
   340  
   341  	lines := strings.Split(string(raw), "\n")
   342  	if len(lines) > 0 && lines[len(lines)-1] == "" {
   343  		lines = lines[:len(lines)-1]
   344  	}
   345  	if pos.Line <= 0 || pos.Line > len(lines) {
   346  		fmt.Fprintf(out, "error formatting %s: %v\n", relPath, err)
   347  		fmt.Fprintf(out, "%s\n", raw)
   348  		return
   349  	}
   350  
   351  	startLine := max(pos.Line-5, 1)
   352  	endLine := min(pos.Line+5, len(lines))
   353  
   354  	for i := startLine; i <= endLine; i++ {
   355  		line := lines[i-1]
   356  		fmt.Fprintf(out, "%s\n", line)
   357  		if i == pos.Line {
   358  			var indent strings.Builder
   359  			for _, ch := range line {
   360  				pos.Column--
   361  				if pos.Column == 0 {
   362  					break
   363  				}
   364  				if ch == '\t' {
   365  					indent.WriteByte('\t')
   366  				} else {
   367  					indent.WriteByte(' ')
   368  				}
   369  			}
   370  			fmt.Fprintf(out, "%s^\n", indent.String())
   371  			fmt.Fprintf(out, "%s\n", strings.TrimRight(err.Error(), "\n"))
   372  		}
   373  	}
   374  }
   375  
   376  func findGOROOT() string {
   377  	cwd, err := os.Getwd()
   378  	if err != nil {
   379  		return ""
   380  	}
   381  	dir := cwd
   382  	for {
   383  		parent := filepath.Dir(dir)
   384  		if parent == dir {
   385  			return ""
   386  		}
   387  		if filepath.Base(dir) == "src" {
   388  			if b, err := os.ReadFile(filepath.Join(dir, "go.mod")); err == nil {
   389  				for line := range strings.SplitSeq(string(b), "\n") {
   390  					fields := strings.Fields(line)
   391  					if len(fields) >= 2 && fields[0] == "module" && fields[1] == "std" {
   392  						return parent
   393  					}
   394  				}
   395  			}
   396  		}
   397  		dir = parent
   398  	}
   399  }
   400  
   401  func resolvePath(goroot, relPath string) string {
   402  	clean := cleanRelPath(relPath)
   403  	if goroot == "" {
   404  		return clean
   405  	}
   406  	return filepath.Join(goroot, clean)
   407  }
   408  
   409  func cleanRelPath(p string) string {
   410  	p = strings.ReplaceAll(p, "\\", "/")
   411  	return filepath.Join("src", p)
   412  }
   413
View as plain text