gen.go 11 KB

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