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