compress_generate.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  1. //+build ignore
  2. // compression_generate.go is meant to run with go generate. It will use
  3. // go/{importer,types} to track down all the RR struct types. Then for each type
  4. // it will look to see if there are (compressible) names, if so it will add that
  5. // type to compressionLenHelperType and comressionLenSearchType which "fake" the
  6. // compression so that Len() is fast.
  7. package main
  8. import (
  9. "bytes"
  10. "fmt"
  11. "go/format"
  12. "go/importer"
  13. "go/types"
  14. "log"
  15. "os"
  16. )
  17. var packageHdr = `
  18. // *** DO NOT MODIFY ***
  19. // AUTOGENERATED BY go generate from compress_generate.go
  20. package dns
  21. `
  22. // getTypeStruct will take a type and the package scope, and return the
  23. // (innermost) struct if the type is considered a RR type (currently defined as
  24. // those structs beginning with a RR_Header, could be redefined as implementing
  25. // the RR interface). The bool return value indicates if embedded structs were
  26. // resolved.
  27. func getTypeStruct(t types.Type, scope *types.Scope) (*types.Struct, bool) {
  28. st, ok := t.Underlying().(*types.Struct)
  29. if !ok {
  30. return nil, false
  31. }
  32. if st.Field(0).Type() == scope.Lookup("RR_Header").Type() {
  33. return st, false
  34. }
  35. if st.Field(0).Anonymous() {
  36. st, _ := getTypeStruct(st.Field(0).Type(), scope)
  37. return st, true
  38. }
  39. return nil, false
  40. }
  41. func main() {
  42. // Import and type-check the package
  43. pkg, err := importer.Default().Import("github.com/miekg/dns")
  44. fatalIfErr(err)
  45. scope := pkg.Scope()
  46. domainTypes := map[string]bool{} // Types that have a domain name in them (either comressible or not).
  47. cdomainTypes := map[string]bool{} // Types that have a compressible domain name in them (subset of domainType)
  48. for _, name := range scope.Names() {
  49. o := scope.Lookup(name)
  50. if o == nil || !o.Exported() {
  51. continue
  52. }
  53. st, _ := getTypeStruct(o.Type(), scope)
  54. if st == nil {
  55. continue
  56. }
  57. if name == "PrivateRR" {
  58. continue
  59. }
  60. if scope.Lookup("Type"+o.Name()) == nil && o.Name() != "RFC3597" {
  61. log.Fatalf("Constant Type%s does not exist.", o.Name())
  62. }
  63. for i := 1; i < st.NumFields(); i++ {
  64. if _, ok := st.Field(i).Type().(*types.Slice); ok {
  65. if st.Tag(i) == `dns:"domain-name"` {
  66. domainTypes[o.Name()] = true
  67. }
  68. if st.Tag(i) == `dns:"cdomain-name"` {
  69. cdomainTypes[o.Name()] = true
  70. domainTypes[o.Name()] = true
  71. }
  72. continue
  73. }
  74. switch {
  75. case st.Tag(i) == `dns:"domain-name"`:
  76. domainTypes[o.Name()] = true
  77. case st.Tag(i) == `dns:"cdomain-name"`:
  78. cdomainTypes[o.Name()] = true
  79. domainTypes[o.Name()] = true
  80. }
  81. }
  82. }
  83. b := &bytes.Buffer{}
  84. b.WriteString(packageHdr)
  85. // compressionLenHelperType - all types that have domain-name/cdomain-name can be used for compressing names
  86. fmt.Fprint(b, "func compressionLenHelperType(c map[string]int, r RR) {\n")
  87. fmt.Fprint(b, "switch x := r.(type) {\n")
  88. for name, _ := range domainTypes {
  89. o := scope.Lookup(name)
  90. st, _ := getTypeStruct(o.Type(), scope)
  91. fmt.Fprintf(b, "case *%s:\n", name)
  92. for i := 1; i < st.NumFields(); i++ {
  93. out := func(s string) { fmt.Fprintf(b, "compressionLenHelper(c, x.%s)\n", st.Field(i).Name()) }
  94. if _, ok := st.Field(i).Type().(*types.Slice); ok {
  95. switch st.Tag(i) {
  96. case `dns:"domain-name"`:
  97. fallthrough
  98. case `dns:"cdomain-name"`:
  99. // For HIP we need to slice over the elements in this slice.
  100. fmt.Fprintf(b, `for i := range x.%s {
  101. compressionLenHelper(c, x.%s[i])
  102. }
  103. `, st.Field(i).Name(), st.Field(i).Name())
  104. }
  105. continue
  106. }
  107. switch {
  108. case st.Tag(i) == `dns:"cdomain-name"`:
  109. fallthrough
  110. case st.Tag(i) == `dns:"domain-name"`:
  111. out(st.Field(i).Name())
  112. }
  113. }
  114. }
  115. fmt.Fprintln(b, "}\n}\n\n")
  116. // compressionLenSearchType - search cdomain-tags types for compressible names.
  117. fmt.Fprint(b, "func compressionLenSearchType(c map[string]int, r RR) (int, bool) {\n")
  118. fmt.Fprint(b, "switch x := r.(type) {\n")
  119. for name, _ := range cdomainTypes {
  120. o := scope.Lookup(name)
  121. st, _ := getTypeStruct(o.Type(), scope)
  122. fmt.Fprintf(b, "case *%s:\n", name)
  123. j := 1
  124. for i := 1; i < st.NumFields(); i++ {
  125. out := func(s string, j int) {
  126. fmt.Fprintf(b, "k%d, ok%d := compressionLenSearch(c, x.%s)\n", j, j, st.Field(i).Name())
  127. }
  128. // There are no slice types with names that can be compressed.
  129. switch {
  130. case st.Tag(i) == `dns:"cdomain-name"`:
  131. out(st.Field(i).Name(), j)
  132. j++
  133. }
  134. }
  135. k := "k1"
  136. ok := "ok1"
  137. for i := 2; i < j; i++ {
  138. k += fmt.Sprintf(" + k%d", i)
  139. ok += fmt.Sprintf(" && ok%d", i)
  140. }
  141. fmt.Fprintf(b, "return %s, %s\n", k, ok)
  142. }
  143. fmt.Fprintln(b, "}\nreturn 0, false\n}\n\n")
  144. // gofmt
  145. res, err := format.Source(b.Bytes())
  146. if err != nil {
  147. b.WriteTo(os.Stderr)
  148. log.Fatal(err)
  149. }
  150. f, err := os.Create("zcompress.go")
  151. fatalIfErr(err)
  152. defer f.Close()
  153. f.Write(res)
  154. }
  155. func fatalIfErr(err error) {
  156. if err != nil {
  157. log.Fatal(err)
  158. }
  159. }