compress_generate.go

//+build ignore

// compression_generate.go is meant to run with go generate. It will use
// go/{importer,types} to track down all the RR struct types. Then for each type
// it will look to see if there are (compressible) names, if so it will add that
// type to compressionLenHelperType and comressionLenSearchType which "fake" the
// compression so that Len() is fast.
package main

import (
	"bytes"
	"fmt"
	"go/format"
	"go/importer"
	"go/types"
	"log"
	"os"
)

var packageHdr = `
// Code generated by "go run compress_generate.go"; DO NOT EDIT.

package dns

`

// getTypeStruct will take a type and the package scope, and return the
// (innermost) struct if the type is considered a RR type (currently defined as
// those structs beginning with a RR_Header, could be redefined as implementing
// the RR interface). The bool return value indicates if embedded structs were
// resolved.
func getTypeStruct(t types.Type, scope *types.Scope) (*types.Struct, bool) {
	st, ok := t.Underlying().(*types.Struct)
	if !ok {
		return nil, false
	}
	if st.Field(0).Type() == scope.Lookup("RR_Header").Type() {
		return st, false
	}
	if st.Field(0).Anonymous() {
		st, _ := getTypeStruct(st.Field(0).Type(), scope)
		return st, true
	}
	return nil, false
}

func main() {
	// Import and type-check the package
	pkg, err := importer.Default().Import("github.com/miekg/dns")
	fatalIfErr(err)
	scope := pkg.Scope()

	var domainTypes []string  // Types that have a domain name in them (either compressible or not).
	var cdomainTypes []string // Types that have a compressible domain name in them (subset of domainType)
Names:
	for _, name := range scope.Names() {
		o := scope.Lookup(name)
		if o == nil || !o.Exported() {
			continue
		}
		st, _ := getTypeStruct(o.Type(), scope)
		if st == nil {
			continue
		}
		if name == "PrivateRR" {
			continue
		}

		if scope.Lookup("Type"+o.Name()) == nil && o.Name() != "RFC3597" {
			log.Fatalf("Constant Type%s does not exist.", o.Name())
		}

		for i := 1; i < st.NumFields(); i++ {
			if _, ok := st.Field(i).Type().(*types.Slice); ok {
				if st.Tag(i) == `dns:"domain-name"` {
					domainTypes = append(domainTypes, o.Name())
					continue Names
				}
				if st.Tag(i) == `dns:"cdomain-name"` {
					cdomainTypes = append(cdomainTypes, o.Name())
					domainTypes = append(domainTypes, o.Name())
					continue Names
				}
				continue
			}

			switch {
			case st.Tag(i) == `dns:"domain-name"`:
				domainTypes = append(domainTypes, o.Name())
				continue Names
			case st.Tag(i) == `dns:"cdomain-name"`:
				cdomainTypes = append(cdomainTypes, o.Name())
				domainTypes = append(domainTypes, o.Name())
				continue Names
			}
		}
	}

	b := &bytes.Buffer{}
	b.WriteString(packageHdr)

	// compressionLenHelperType - all types that have domain-name/cdomain-name can be used for compressing names

	fmt.Fprint(b, "func compressionLenHelperType(c map[string]int, r RR, initLen int) int {\n")
	fmt.Fprint(b, "currentLen := initLen\n")
	fmt.Fprint(b, "switch x := r.(type) {\n")
	for _, name := range domainTypes {
		o := scope.Lookup(name)
		st, _ := getTypeStruct(o.Type(), scope)

		fmt.Fprintf(b, "case *%s:\n", name)
		for i := 1; i < st.NumFields(); i++ {
			out := func(s string) {
				fmt.Fprintf(b, "currentLen -= len(x.%s) + 1\n", st.Field(i).Name())
				fmt.Fprintf(b, "currentLen += compressionLenHelper(c, x.%s, currentLen)\n", st.Field(i).Name())
			}

			if _, ok := st.Field(i).Type().(*types.Slice); ok {
				switch st.Tag(i) {
				case `dns:"domain-name"`:
					fallthrough
				case `dns:"cdomain-name"`:
					// For HIP we need to slice over the elements in this slice.
					fmt.Fprintf(b, `for i := range x.%s {
	currentLen -= len(x.%s[i]) + 1
}
`, st.Field(i).Name(), st.Field(i).Name())
					fmt.Fprintf(b, `for i := range x.%s {
	currentLen += compressionLenHelper(c, x.%s[i], currentLen)
}
`, st.Field(i).Name(), st.Field(i).Name())
				}
				continue
			}

			switch {
			case st.Tag(i) == `dns:"cdomain-name"`:
				fallthrough
			case st.Tag(i) == `dns:"domain-name"`:
				out(st.Field(i).Name())
			}
		}
	}
	fmt.Fprintln(b, "}\nreturn currentLen - initLen\n}\n\n")

	// compressionLenSearchType - search cdomain-tags types for compressible names.

	fmt.Fprint(b, "func compressionLenSearchType(c map[string]int, r RR) (int, bool, int) {\n")
	fmt.Fprint(b, "switch x := r.(type) {\n")
	for _, name := range cdomainTypes {
		o := scope.Lookup(name)
		st, _ := getTypeStruct(o.Type(), scope)

		fmt.Fprintf(b, "case *%s:\n", name)
		j := 1
		for i := 1; i < st.NumFields(); i++ {
			out := func(s string, j int) {
				fmt.Fprintf(b, "k%d, ok%d, sz%d := compressionLenSearch(c, x.%s)\n", j, j, j, st.Field(i).Name())
			}

			// There are no slice types with names that can be compressed.

			switch {
			case st.Tag(i) == `dns:"cdomain-name"`:
				out(st.Field(i).Name(), j)
				j++
			}
		}
		k := "k1"
		ok := "ok1"
		sz := "sz1"
		for i := 2; i < j; i++ {
			k += fmt.Sprintf(" + k%d", i)
			ok += fmt.Sprintf(" && ok%d", i)
			sz += fmt.Sprintf(" + sz%d", i)
		}
		fmt.Fprintf(b, "return %s, %s, %s\n", k, ok, sz)
	}
	fmt.Fprintln(b, "}\nreturn 0, false, 0\n}\n\n")

	// gofmt
	res, err := format.Source(b.Bytes())
	if err != nil {
		b.WriteTo(os.Stderr)
		log.Fatal(err)
	}

	f, err := os.Create("zcompress.go")
	fatalIfErr(err)
	defer f.Close()
	f.Write(res)
}

func fatalIfErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}