hanayo/vendor/github.com/ugorji/go/codec/codecgen/gen.go
2019-02-23 13:29:15 +00:00

361 lines
11 KiB
Go

// Copyright (c) 2012-2018 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.
// codecgen generates codec.Selfer implementations for a set of types.
package main
import (
"bufio"
"bytes"
"errors"
"flag"
"fmt"
"go/ast"
"go/build"
"go/parser"
"go/token"
"math/rand"
"os"
"os/exec"
"path/filepath"
"regexp"
"strconv"
"strings"
"text/template"
"time"
)
const genCodecPkg = "codec1978" // keep this in sync with codec.genCodecPkg
const genFrunMainTmpl = `//+build ignore
// Code generated - temporary main package for codecgen - DO NOT EDIT.
package main
{{ if .Types }}import "{{ .ImportPath }}"{{ end }}
func main() {
{{ $.PackageName }}.CodecGenTempWrite{{ .RandString }}()
}
`
// const genFrunPkgTmpl = `//+build codecgen
const genFrunPkgTmpl = `
// Code generated - temporary package for codecgen - DO NOT EDIT.
package {{ $.PackageName }}
import (
{{ if not .CodecPkgFiles }}{{ .CodecPkgName }} "{{ .CodecImportPath }}"{{ end }}
"os"
"reflect"
"bytes"
"strings"
"go/format"
)
func CodecGenTempWrite{{ .RandString }}() {
os.Remove("{{ .OutFile }}")
fout, err := os.Create("{{ .OutFile }}")
if err != nil {
panic(err)
}
defer fout.Close()
var typs []reflect.Type
var typ reflect.Type
var numfields int
{{ range $index, $element := .Types }}
var t{{ $index }} {{ . }}
typ = reflect.TypeOf(t{{ $index }})
typs = append(typs, typ)
if typ.Kind() == reflect.Struct { numfields += typ.NumField() } else { numfields += 1 }
{{ end }}
// println("initializing {{ .OutFile }}, buf size: {{ .AllFilesSize }}*16",
// {{ .AllFilesSize }}*16, "num fields: ", numfields)
var out = bytes.NewBuffer(make([]byte, 0, numfields*1024)) // {{ .AllFilesSize }}*16
{{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}Gen(out,
"{{ .BuildTag }}", "{{ .PackageName }}", "{{ .RandString }}", {{ .NoExtensions }},
{{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}NewTypeInfos(strings.Split("{{ .StructTags }}", ",")),
typs...)
bout, err := format.Source(out.Bytes())
// println("... lengths: before formatting: ", len(out.Bytes()), ", after formatting", len(bout))
if err != nil {
fout.Write(out.Bytes())
panic(err)
}
fout.Write(bout)
}
`
// Generate is given a list of *.go files to parse, and an output file (fout).
//
// It finds all types T in the files, and it creates 2 tmp files (frun).
// - main package file passed to 'go run'
// - package level file which calls *genRunner.Selfer to write Selfer impls for each T.
// We use a package level file so that it can reference unexported types in the package being worked on.
// Tool then executes: "go run __frun__" which creates fout.
// fout contains Codec(En|De)codeSelf implementations for every type T.
//
func Generate(outfile, buildTag, codecPkgPath string,
uid int64,
goRunTag string, st string,
regexName, notRegexName *regexp.Regexp,
deleteTempFile, noExtensions bool,
infiles ...string) (err error) {
// For each file, grab AST, find each type, and write a call to it.
if len(infiles) == 0 {
return
}
if outfile == "" || codecPkgPath == "" {
err = errors.New("outfile and codec package path cannot be blank")
return
}
if uid < 0 {
uid = -uid
} else if uid == 0 {
rr := rand.New(rand.NewSource(time.Now().UnixNano()))
uid = 101 + rr.Int63n(9777)
}
// We have to parse dir for package, before opening the temp file for writing (else ImportDir fails).
// Also, ImportDir(...) must take an absolute path.
lastdir := filepath.Dir(outfile)
absdir, err := filepath.Abs(lastdir)
if err != nil {
return
}
pkg, err := build.Default.ImportDir(absdir, build.AllowBinary)
if err != nil {
return
}
type tmplT struct {
CodecPkgName string
CodecImportPath string
ImportPath string
OutFile string
PackageName string
RandString string
BuildTag string
StructTags string
Types []string
AllFilesSize int64
CodecPkgFiles bool
NoExtensions bool
}
tv := tmplT{
CodecPkgName: genCodecPkg,
OutFile: outfile,
CodecImportPath: codecPkgPath,
BuildTag: buildTag,
RandString: strconv.FormatInt(uid, 10),
StructTags: st,
NoExtensions: noExtensions,
}
tv.ImportPath = pkg.ImportPath
if tv.ImportPath == tv.CodecImportPath {
tv.CodecPkgFiles = true
tv.CodecPkgName = "codec"
} else {
// HACK: always handle vendoring. It should be typically on in go 1.6, 1.7
tv.ImportPath = stripVendor(tv.ImportPath)
}
astfiles := make([]*ast.File, len(infiles))
var fi os.FileInfo
for i, infile := range infiles {
if filepath.Dir(infile) != lastdir {
err = errors.New("in files must all be in same directory as outfile")
return
}
if fi, err = os.Stat(infile); err != nil {
return
}
tv.AllFilesSize += fi.Size()
fset := token.NewFileSet()
astfiles[i], err = parser.ParseFile(fset, infile, nil, 0)
if err != nil {
return
}
if i == 0 {
tv.PackageName = astfiles[i].Name.Name
if tv.PackageName == "main" {
// codecgen cannot be run on types in the 'main' package.
// A temporary 'main' package must be created, and should reference the fully built
// package containing the types.
// Also, the temporary main package will conflict with the main package which already has a main method.
err = errors.New("codecgen cannot be run on types in the 'main' package")
return
}
}
}
// keep track of types with selfer methods
// selferMethods := []string{"CodecEncodeSelf", "CodecDecodeSelf"}
selferEncTyps := make(map[string]bool)
selferDecTyps := make(map[string]bool)
for _, f := range astfiles {
for _, d := range f.Decls {
// if fd, ok := d.(*ast.FuncDecl); ok && fd.Recv != nil && fd.Recv.NumFields() == 1 {
if fd, ok := d.(*ast.FuncDecl); ok && fd.Recv != nil && len(fd.Recv.List) == 1 {
recvType := fd.Recv.List[0].Type
if ptr, ok := recvType.(*ast.StarExpr); ok {
recvType = ptr.X
}
if id, ok := recvType.(*ast.Ident); ok {
switch fd.Name.Name {
case "CodecEncodeSelf":
selferEncTyps[id.Name] = true
case "CodecDecodeSelf":
selferDecTyps[id.Name] = true
}
}
}
}
}
// now find types
for _, f := range astfiles {
for _, d := range f.Decls {
if gd, ok := d.(*ast.GenDecl); ok {
for _, dd := range gd.Specs {
if td, ok := dd.(*ast.TypeSpec); ok {
// if len(td.Name.Name) == 0 || td.Name.Name[0] > 'Z' || td.Name.Name[0] < 'A' {
if len(td.Name.Name) == 0 {
continue
}
// only generate for:
// struct: StructType
// primitives (numbers, bool, string): Ident
// map: MapType
// slice, array: ArrayType
// chan: ChanType
// do not generate:
// FuncType, InterfaceType, StarExpr (ptr), etc
//
// We generate for all these types (not just structs), because they may be a field
// in another struct which doesn't have codecgen run on it, and it will be nice
// to take advantage of the fact that the type is a Selfer.
switch td.Type.(type) {
case *ast.StructType, *ast.Ident, *ast.MapType, *ast.ArrayType, *ast.ChanType:
// only add to tv.Types iff
// - it matches per the -r parameter
// - it doesn't match per the -nr parameter
// - it doesn't have any of the Selfer methods in the file
if regexName.FindStringIndex(td.Name.Name) != nil &&
notRegexName.FindStringIndex(td.Name.Name) == nil &&
!selferEncTyps[td.Name.Name] &&
!selferDecTyps[td.Name.Name] {
tv.Types = append(tv.Types, td.Name.Name)
}
}
}
}
}
}
}
if len(tv.Types) == 0 {
return
}
// we cannot use ioutil.TempFile, because we cannot guarantee the file suffix (.go).
// Also, we cannot create file in temp directory,
// because go run will not work (as it needs to see the types here).
// Consequently, create the temp file in the current directory, and remove when done.
// frun, err = ioutil.TempFile("", "codecgen-")
// frunName := filepath.Join(os.TempDir(), "codecgen-"+strconv.FormatInt(time.Now().UnixNano(), 10)+".go")
frunMainName := "codecgen-main-" + tv.RandString + ".generated.go"
frunPkgName := "codecgen-pkg-" + tv.RandString + ".generated.go"
if deleteTempFile {
defer os.Remove(frunMainName)
defer os.Remove(frunPkgName)
}
// var frunMain, frunPkg *os.File
if _, err = gen1(frunMainName, genFrunMainTmpl, &tv); err != nil {
return
}
if _, err = gen1(frunPkgName, genFrunPkgTmpl, &tv); err != nil {
return
}
// remove outfile, so "go run ..." will not think that types in outfile already exist.
os.Remove(outfile)
// execute go run frun
cmd := exec.Command("go", "run", "-tags", "codecgen.exec safe "+goRunTag, frunMainName) //, frunPkg.Name())
var buf bytes.Buffer
cmd.Stdout = &buf
cmd.Stderr = &buf
if err = cmd.Run(); err != nil {
err = fmt.Errorf("error running 'go run %s': %v, console: %s",
frunMainName, err, buf.Bytes())
return
}
os.Stdout.Write(buf.Bytes())
return
}
func gen1(frunName, tmplStr string, tv interface{}) (frun *os.File, err error) {
os.Remove(frunName)
if frun, err = os.Create(frunName); err != nil {
return
}
defer frun.Close()
t := template.New("")
if t, err = t.Parse(tmplStr); err != nil {
return
}
bw := bufio.NewWriter(frun)
if err = t.Execute(bw, tv); err != nil {
bw.Flush()
return
}
if err = bw.Flush(); err != nil {
return
}
return
}
// copied from ../gen.go (keep in sync).
func stripVendor(s string) string {
// HACK: Misbehaviour occurs in go 1.5. May have to re-visit this later.
// if s contains /vendor/ OR startsWith vendor/, then return everything after it.
const vendorStart = "vendor/"
const vendorInline = "/vendor/"
if i := strings.LastIndex(s, vendorInline); i >= 0 {
s = s[i+len(vendorInline):]
} else if strings.HasPrefix(s, vendorStart) {
s = s[len(vendorStart):]
}
return s
}
func main() {
o := flag.String("o", "", "out file")
c := flag.String("c", genCodecPath, "codec path")
t := flag.String("t", "", "build tag to put in file")
r := flag.String("r", ".*", "regex for type name to match")
nr := flag.String("nr", "^$", "regex for type name to exclude")
rt := flag.String("rt", "", "tags for go run")
st := flag.String("st", "codec,json", "struct tag keys to introspect")
x := flag.Bool("x", false, "keep temp file")
_ = flag.Bool("u", false, "Allow unsafe use. ***IGNORED*** - kept for backwards compatibility: ")
d := flag.Int64("d", 0, "random identifier for use in generated code")
nx := flag.Bool("nx", false, "do not support extensions - support of extensions may cause extra allocation")
flag.Parse()
err := Generate(*o, *t, *c, *d, *rt, *st,
regexp.MustCompile(*r), regexp.MustCompile(*nr), !*x, *nx, flag.Args()...)
if err != nil {
fmt.Fprintf(os.Stderr, "codecgen error: %v\n", err)
os.Exit(1)
}
}