gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/tools/go_generics/go_merge/main.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"bytes"
    19  	"flag"
    20  	"fmt"
    21  	"go/ast"
    22  	"go/format"
    23  	"go/parser"
    24  	"go/token"
    25  	"os"
    26  	"path/filepath"
    27  	"strconv"
    28  
    29  	"gvisor.dev/gvisor/tools/constraintutil"
    30  )
    31  
    32  var (
    33  	output = flag.String("o", "", "output `file`")
    34  )
    35  
    36  func fatalf(s string, args ...any) {
    37  	fmt.Fprintf(os.Stderr, s, args...)
    38  	os.Exit(1)
    39  }
    40  
    41  func main() {
    42  	flag.Usage = func() {
    43  		fmt.Fprintf(os.Stderr, "Usage: %s [options] <input1> [<input2> ...]\n", os.Args[0])
    44  		flag.PrintDefaults()
    45  	}
    46  
    47  	flag.Parse()
    48  	if *output == "" || len(flag.Args()) == 0 {
    49  		flag.Usage()
    50  		os.Exit(1)
    51  	}
    52  
    53  	// Load all files.
    54  	files := make(map[string]*ast.File)
    55  	fset := token.NewFileSet()
    56  	var name string
    57  	for _, fname := range flag.Args() {
    58  		f, err := parser.ParseFile(fset, fname, nil, parser.ParseComments|parser.DeclarationErrors|parser.SpuriousErrors)
    59  		if err != nil {
    60  			fatalf("%v\n", err)
    61  		}
    62  
    63  		files[fname] = f
    64  		if name == "" {
    65  			name = f.Name.Name
    66  		} else if name != f.Name.Name {
    67  			fatalf("Expected '%s' for package name instead of '%s'.\n", name, f.Name.Name)
    68  		}
    69  	}
    70  
    71  	// Merge all files into one.
    72  	pkg := &ast.Package{
    73  		Name:  name,
    74  		Files: files,
    75  	}
    76  	f := ast.MergePackageFiles(pkg, ast.FilterUnassociatedComments|ast.FilterFuncDuplicates|ast.FilterImportDuplicates)
    77  
    78  	// Create a new declaration slice with all imports at the top, merging any
    79  	// redundant imports.
    80  	imports := make(map[string]*ast.ImportSpec)
    81  	var importNames []string // Keep imports in the original order to get deterministic output.
    82  	var anonImports []*ast.ImportSpec
    83  	for _, d := range f.Decls {
    84  		if g, ok := d.(*ast.GenDecl); ok && g.Tok == token.IMPORT {
    85  			for _, s := range g.Specs {
    86  				i := s.(*ast.ImportSpec)
    87  				p, _ := strconv.Unquote(i.Path.Value)
    88  				var n string
    89  				if i.Name == nil {
    90  					n = filepath.Base(p)
    91  				} else {
    92  					n = i.Name.Name
    93  				}
    94  				if n == "_" {
    95  					anonImports = append(anonImports, i)
    96  				} else {
    97  					if i2, ok := imports[n]; ok {
    98  						if first, second := i.Path.Value, i2.Path.Value; first != second {
    99  							fatalf("Conflicting paths for import name '%s': '%s' vs. '%s'\n", n, first, second)
   100  						}
   101  					} else {
   102  						imports[n] = i
   103  						importNames = append(importNames, n)
   104  					}
   105  				}
   106  			}
   107  		}
   108  	}
   109  	newDecls := make([]ast.Decl, 0, len(f.Decls))
   110  	if l := len(imports) + len(anonImports); l > 0 {
   111  		// Non-NoPos Lparen is needed for Go to recognize more than one spec in
   112  		// ast.GenDecl.Specs.
   113  		d := &ast.GenDecl{
   114  			Tok:    token.IMPORT,
   115  			Lparen: token.NoPos + 1,
   116  			Specs:  make([]ast.Spec, 0, l),
   117  		}
   118  		for _, i := range importNames {
   119  			d.Specs = append(d.Specs, imports[i])
   120  		}
   121  		for _, i := range anonImports {
   122  			d.Specs = append(d.Specs, i)
   123  		}
   124  		newDecls = append(newDecls, d)
   125  	}
   126  	for _, d := range f.Decls {
   127  		if g, ok := d.(*ast.GenDecl); !ok || g.Tok != token.IMPORT {
   128  			newDecls = append(newDecls, d)
   129  		}
   130  	}
   131  	f.Decls = newDecls
   132  
   133  	// Infer build constraints for the output file.
   134  	bcexpr, err := constraintutil.CombineFromFiles(flag.Args())
   135  	if err != nil {
   136  		fatalf("Failed to read build constraints: %v\n", err)
   137  	}
   138  
   139  	// Write the output file.
   140  	var buf bytes.Buffer
   141  	if err := format.Node(&buf, fset, f); err != nil {
   142  		fatalf("fomatting: %v\n", err)
   143  	}
   144  	outf, err := os.OpenFile(*output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
   145  	if err != nil {
   146  		fatalf("opening output: %v\n", err)
   147  	}
   148  	defer outf.Close()
   149  	outf.WriteString(constraintutil.Lines(bcexpr))
   150  	if _, err := outf.Write(buf.Bytes()); err != nil {
   151  		fatalf("write: %v\n", err)
   152  	}
   153  }