vendor/github.com/ugorji/go/codec/codecgen/gen.go

   1 // Copyright (c) 2012-2018 Ugorji Nwoke. All rights reserved.
   2 // Use of this source code is governed by a MIT license found in the LICENSE file.
   3
   4 // codecgen generates codec.Selfer implementations for a set of types.
   5 package main
   6
   7 import (
   8         "bufio"
   9         "bytes"
  10         "errors"
  11         "flag"
  12         "fmt"
  13         "go/ast"
  14         "go/parser"
  15         "go/token"
  16         "math/rand"
  17         "os"
  18         "os/exec"
  19         "path/filepath"
  20         "regexp"
  21         "strconv"
  22         "strings"
  23         "text/template"
  24         "time"
  25 )
  26
  27 const genCodecPkg = "codec1978" // keep this in sync with codec.genCodecPkg
  28
  29 const genFrunMainTmpl = `//+build ignore
  30
  31 // Code generated - temporary main package for codecgen - DO NOT EDIT.
  32
  33 package main
  34 {{ if .Types }}import "{{ .ImportPath }}"{{ end }}
  35 func main() {
  36         {{ $.PackageName }}.CodecGenTempWrite{{ .RandString }}()
  37 }
  38 `
  39
  40 // const genFrunPkgTmpl = `//+build codecgen
  41 const genFrunPkgTmpl = `
  42
  43 // Code generated - temporary package for codecgen - DO NOT EDIT.
  44
  45 package {{ $.PackageName }}
  46
  47 import (
  48         {{ if not .CodecPkgFiles }}{{ .CodecPkgName }} "{{ .CodecImportPath }}"{{ end }}
  49         "os"
  50         "reflect"
  51         "bytes"
  52         "strings"
  53         "go/format"
  54 )
  55
  56 func CodecGenTempWrite{{ .RandString }}() {
  57         os.Remove("{{ .OutFile }}")
  58         fout, err := os.Create("{{ .OutFile }}")
  59         if err != nil {
  60                 panic(err)
  61         }
  62         defer fout.Close()
  63
  64         var typs []reflect.Type
  65         var typ reflect.Type
  66         var numfields int
  67 {{ range $index, $element := .Types }}
  68         var t{{ $index }} {{ . }}
  69 typ = reflect.TypeOf(t{{ $index }})
  70         typs = append(typs, typ)
  71         if typ.Kind() == reflect.Struct { numfields += typ.NumField() } else { numfields += 1 }
  72 {{ end }}
  73
  74         // println("initializing {{ .OutFile }}, buf size: {{ .AllFilesSize }}*16",
  75         //      {{ .AllFilesSize }}*16, "num fields: ", numfields)
  76         var out = bytes.NewBuffer(make([]byte, 0, numfields*1024)) // {{ .AllFilesSize }}*16
  77         {{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}Gen(out,
  78                 "{{ .BuildTag }}", "{{ .PackageName }}", "{{ .RandString }}", {{ .NoExtensions }},
  79                 {{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}NewTypeInfos(strings.Split("{{ .StructTags }}", ",")),
  80                  typs...)
  81
  82         bout, err := format.Source(out.Bytes())
  83         // println("... lengths: before formatting: ", len(out.Bytes()), ", after formatting", len(bout))
  84         if err != nil {
  85                 fout.Write(out.Bytes())
  86                 panic(err)
  87         }
  88         fout.Write(bout)
  89 }
  90
  91 `
  92
  93 // Generate is given a list of *.go files to parse, and an output file (fout).
  94 //
  95 // It finds all types T in the files, and it creates 2 tmp files (frun).
  96 //   - main package file passed to 'go run'
  97 //   - package level file which calls *genRunner.Selfer to write Selfer impls for each T.
  98 // We use a package level file so that it can reference unexported types in the package being worked on.
  99 // Tool then executes: "go run __frun__" which creates fout.
 100 // fout contains Codec(En|De)codeSelf implementations for every type T.
 101 //
 102 func Generate(outfile, buildTag, codecPkgPath string,
 103         uid int64,
 104         goRunTag string, st string,
 105         regexName, notRegexName *regexp.Regexp,
 106         deleteTempFile, noExtensions bool,
 107         infiles ...string) (err error) {
 108         // For each file, grab AST, find each type, and write a call to it.
 109         if len(infiles) == 0 {
 110                 return
 111         }
 112         if outfile == "" || codecPkgPath == "" {
 113                 err = errors.New("outfile and codec package path cannot be blank")
 114                 return
 115         }
 116         if uid < 0 {
 117                 uid = -uid
 118         } else if uid == 0 {
 119                 rr := rand.New(rand.NewSource(time.Now().UnixNano()))
 120                 uid = 101 + rr.Int63n(9777)
 121         }
 122         // We have to parse dir for package, before opening the temp file for writing (else ImportDir fails).
 123         // Also, ImportDir(...) must take an absolute path.
 124         lastdir := filepath.Dir(outfile)
 125         absdir, err := filepath.Abs(lastdir)
 126         if err != nil {
 127                 return
 128         }
 129         importPath, err := pkgPath(absdir)
 130         if err != nil {
 131                 return
 132         }
 133         type tmplT struct {
 134                 CodecPkgName    string
 135                 CodecImportPath string
 136                 ImportPath      string
 137                 OutFile         string
 138                 PackageName     string
 139                 RandString      string
 140                 BuildTag        string
 141                 StructTags      string
 142                 Types           []string
 143                 AllFilesSize    int64
 144                 CodecPkgFiles   bool
 145                 NoExtensions    bool
 146         }
 147         tv := tmplT{
 148                 CodecPkgName:    genCodecPkg,
 149                 OutFile:         outfile,
 150                 CodecImportPath: codecPkgPath,
 151                 BuildTag:        buildTag,
 152                 RandString:      strconv.FormatInt(uid, 10),
 153                 StructTags:      st,
 154                 NoExtensions:    noExtensions,
 155         }
 156         tv.ImportPath = importPath
 157         if tv.ImportPath == tv.CodecImportPath {
 158                 tv.CodecPkgFiles = true
 159                 tv.CodecPkgName = "codec"
 160         } else {
 161                 // HACK: always handle vendoring. It should be typically on in go 1.6, 1.7
 162                 tv.ImportPath = stripVendor(tv.ImportPath)
 163         }
 164         astfiles := make([]*ast.File, len(infiles))
 165         var fi os.FileInfo
 166         for i, infile := range infiles {
 167                 if filepath.Dir(infile) != lastdir {
 168                         err = errors.New("in files must all be in same directory as outfile")
 169                         return
 170                 }
 171                 if fi, err = os.Stat(infile); err != nil {
 172                         return
 173                 }
 174                 tv.AllFilesSize += fi.Size()
 175
 176                 fset := token.NewFileSet()
 177                 astfiles[i], err = parser.ParseFile(fset, infile, nil, 0)
 178                 if err != nil {
 179                         return
 180                 }
 181                 if i == 0 {
 182                         tv.PackageName = astfiles[i].Name.Name
 183                         if tv.PackageName == "main" {
 184                                 // codecgen cannot be run on types in the 'main' package.
 185                                 // A temporary 'main' package must be created, and should reference the fully built
 186                                 // package containing the types.
 187                                 // Also, the temporary main package will conflict with the main package which already has a main method.
 188                                 err = errors.New("codecgen cannot be run on types in the 'main' package")
 189                                 return
 190                         }
 191                 }
 192         }
 193
 194         // keep track of types with selfer methods
 195         // selferMethods := []string{"CodecEncodeSelf", "CodecDecodeSelf"}
 196         selferEncTyps := make(map[string]bool)
 197         selferDecTyps := make(map[string]bool)
 198         for _, f := range astfiles {
 199                 for _, d := range f.Decls {
 200                         // if fd, ok := d.(*ast.FuncDecl); ok && fd.Recv != nil && fd.Recv.NumFields() == 1 {
 201                         if fd, ok := d.(*ast.FuncDecl); ok && fd.Recv != nil && len(fd.Recv.List) == 1 {
 202                                 recvType := fd.Recv.List[0].Type
 203                                 if ptr, ok := recvType.(*ast.StarExpr); ok {
 204                                         recvType = ptr.X
 205                                 }
 206                                 if id, ok := recvType.(*ast.Ident); ok {
 207                                         switch fd.Name.Name {
 208                                         case "CodecEncodeSelf":
 209                                                 selferEncTyps[id.Name] = true
 210                                         case "CodecDecodeSelf":
 211                                                 selferDecTyps[id.Name] = true
 212                                         }
 213                                 }
 214                         }
 215                 }
 216         }
 217
 218         // now find types
 219         for _, f := range astfiles {
 220                 for _, d := range f.Decls {
 221                         if gd, ok := d.(*ast.GenDecl); ok {
 222                                 for _, dd := range gd.Specs {
 223                                         if td, ok := dd.(*ast.TypeSpec); ok {
 224                                                 // if len(td.Name.Name) == 0 || td.Name.Name[0] > 'Z' || td.Name.Name[0] < 'A' {
 225                                                 if len(td.Name.Name) == 0 {
 226                                                         continue
 227                                                 }
 228
 229                                                 // only generate for:
 230                                                 //   struct: StructType
 231                                                 //   primitives (numbers, bool, string): Ident
 232                                                 //   map: MapType
 233                                                 //   slice, array: ArrayType
 234                                                 //   chan: ChanType
 235                                                 // do not generate:
 236                                                 //   FuncType, InterfaceType, StarExpr (ptr), etc
 237                                                 //
 238                                                 // We generate for all these types (not just structs), because they may be a field
 239                                                 // in another struct which doesn't have codecgen run on it, and it will be nice
 240                                                 // to take advantage of the fact that the type is a Selfer.
 241                                                 switch td.Type.(type) {
 242                                                 case *ast.StructType, *ast.Ident, *ast.MapType, *ast.ArrayType, *ast.ChanType:
 243                                                         // only add to tv.Types iff
 244                                                         //   - it matches per the -r parameter
 245                                                         //   - it doesn't match per the -nr parameter
 246                                                         //   - it doesn't have any of the Selfer methods in the file
 247                                                         if regexName.FindStringIndex(td.Name.Name) != nil &&
 248                                                                 notRegexName.FindStringIndex(td.Name.Name) == nil &&
 249                                                                 !selferEncTyps[td.Name.Name] &&
 250                                                                 !selferDecTyps[td.Name.Name] {
 251                                                                 tv.Types = append(tv.Types, td.Name.Name)
 252                                                         }
 253                                                 }
 254                                         }
 255                                 }
 256                         }
 257                 }
 258         }
 259
 260         if len(tv.Types) == 0 {
 261                 return
 262         }
 263
 264         // we cannot use ioutil.TempFile, because we cannot guarantee the file suffix (.go).
 265         // Also, we cannot create file in temp directory,
 266         // because go run will not work (as it needs to see the types here).
 267         // Consequently, create the temp file in the current directory, and remove when done.
 268
 269         // frun, err = ioutil.TempFile("", "codecgen-")
 270         // frunName := filepath.Join(os.TempDir(), "codecgen-"+strconv.FormatInt(time.Now().UnixNano(), 10)+".go")
 271
 272         frunMainName := "codecgen-main-" + tv.RandString + ".generated.go"
 273         frunPkgName := "codecgen-pkg-" + tv.RandString + ".generated.go"
 274         if deleteTempFile {
 275                 defer os.Remove(frunMainName)
 276                 defer os.Remove(frunPkgName)
 277         }
 278         // var frunMain, frunPkg *os.File
 279         if _, err = gen1(frunMainName, genFrunMainTmpl, &tv); err != nil {
 280                 return
 281         }
 282         if _, err = gen1(frunPkgName, genFrunPkgTmpl, &tv); err != nil {
 283                 return
 284         }
 285
 286         // remove outfile, so "go run ..." will not think that types in outfile already exist.
 287         os.Remove(outfile)
 288
 289         // execute go run frun
 290         cmd := exec.Command("go", "run", "-tags", "codecgen.exec safe "+goRunTag, frunMainName) //, frunPkg.Name())
 291         var buf bytes.Buffer
 292         cmd.Stdout = &buf
 293         cmd.Stderr = &buf
 294         if err = cmd.Run(); err != nil {
 295                 err = fmt.Errorf("error running 'go run %s': %v, console: %s",
 296                         frunMainName, err, buf.Bytes())
 297                 return
 298         }
 299         os.Stdout.Write(buf.Bytes())
 300         return
 301 }
 302
 303 func gen1(frunName, tmplStr string, tv interface{}) (frun *os.File, err error) {
 304         os.Remove(frunName)
 305         if frun, err = os.Create(frunName); err != nil {
 306                 return
 307         }
 308         defer frun.Close()
 309
 310         t := template.New("")
 311         if t, err = t.Parse(tmplStr); err != nil {
 312                 return
 313         }
 314         bw := bufio.NewWriter(frun)
 315         if err = t.Execute(bw, tv); err != nil {
 316                 bw.Flush()
 317                 return
 318         }
 319         if err = bw.Flush(); err != nil {
 320                 return
 321         }
 322         return
 323 }
 324
 325 // copied from ../gen.go (keep in sync).
 326 func stripVendor(s string) string {
 327         // HACK: Misbehaviour occurs in go 1.5. May have to re-visit this later.
 328         // if s contains /vendor/ OR startsWith vendor/, then return everything after it.
 329         const vendorStart = "vendor/"
 330         const vendorInline = "/vendor/"
 331         if i := strings.LastIndex(s, vendorInline); i >= 0 {
 332                 s = s[i+len(vendorInline):]
 333         } else if strings.HasPrefix(s, vendorStart) {
 334                 s = s[len(vendorStart):]
 335         }
 336         return s
 337 }
 338
 339 func main() {
 340         o := flag.String("o", "", "out file")
 341         c := flag.String("c", genCodecPath, "codec path")
 342         t := flag.String("t", "", "build tag to put in file")
 343         r := flag.String("r", ".*", "regex for type name to match")
 344         nr := flag.String("nr", "^$", "regex for type name to exclude")
 345         rt := flag.String("rt", "", "tags for go run")
 346         st := flag.String("st", "codec,json", "struct tag keys to introspect")
 347         x := flag.Bool("x", false, "keep temp file")
 348         _ = flag.Bool("u", false, "Allow unsafe use. ***IGNORED*** - kept for backwards compatibility: ")
 349         d := flag.Int64("d", 0, "random identifier for use in generated code")
 350         nx := flag.Bool("nx", false, "do not support extensions - support of extensions may cause extra allocation")
 351
 352         flag.Parse()
 353         err := Generate(*o, *t, *c, *d, *rt, *st,
 354                 regexp.MustCompile(*r), regexp.MustCompile(*nr), !*x, *nx, flag.Args()...)
 355         if err != nil {
 356                 fmt.Fprintf(os.Stderr, "codecgen error: %v\n", err)
 357                 os.Exit(1)
 358         }
 359 }