WIPandrew/fastjson

Change-Id: I5295d47102d879f29f0a6818481e8b65eafd02dd

4 files changed, 610 insertions, 0 deletions

diff --git a/cmd/fastjson/fastjson.go b/cmd/fastjson/fastjson.go
new file mode 100644
index 000000000..1b017235c
--- /dev/null
+++ b/cmd/fastjson/fastjson.go
@@ -0,0 +1,329 @@
+package main
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"go/ast"
+	"go/token"
+	"go/types"
+	"log"
+	"os"
+	"strconv"
+	"strings"
+
+	"golang.org/x/tools/go/packages"
+	"golang.org/x/tools/imports"
+	"tailscale.com/util/codegen"
+)
+
+var (
+	flagTypes     = flag.String("type", "", "comma-separated list of types; required")
+	flagBuildTags = flag.String("tags", "", "compiler build tags to apply")
+)
+
+func main() {
+	log.SetFlags(0)
+	log.SetPrefix("cloner: ")
+	log.SetOutput(os.Stderr)
+	flag.Parse()
+	if len(*flagTypes) == 0 {
+		flag.Usage()
+		os.Exit(2)
+	}
+	typeNames := strings.Split(*flagTypes, ",")
+
+	pkg, namedTypes, err := loadTypes(".", *flagBuildTags)
+	if err != nil {
+		log.Fatal(err)
+	}
+	it := codegen.NewImportTracker(pkg.Types)
+	buf := new(bytes.Buffer)
+
+	for _, typeName := range typeNames {
+		typ, ok := namedTypes[typeName]
+		if !ok {
+			log.Fatalf("could not find type %s", typeName)
+		}
+		gen(buf, it, typ)
+	}
+
+	outBuf := new(bytes.Buffer)
+	outBuf.WriteString("// Code generated by TODO; DO NOT EDIT.\n")
+	outBuf.WriteString("\n")
+	fmt.Fprintf(outBuf, "package %s\n\n", pkg.Name)
+	it.Write(outBuf)
+	outBuf.Write(buf.Bytes())
+
+	// Best-effort gofmt the output
+	out := outBuf.Bytes()
+	out, err = imports.Process("/nonexistant/main.go", out, &imports.Options{
+		Comments:   true,
+		TabIndent:  true,
+		TabWidth:   8,
+		FormatOnly: true, // fancy gofmt only
+	})
+	if err != nil {
+		out = outBuf.Bytes()
+	}
+	fmt.Print(string(out))
+}
+
+func gen(buf *bytes.Buffer, it *codegen.ImportTracker, typ *types.Named) {
+	t, ok := typ.Underlying().(*types.Struct)
+	if !ok {
+		return
+	}
+
+	name := typ.Obj().Name()
+	fmt.Fprintf(buf, "// MarshalJSONInto marshals this %s into JSON in the provided buffer.\n", name)
+	fmt.Fprintf(buf, "func (self *%s) MarshalJSONInto(buf []byte) ([]byte, error) {\n", name)
+	fmt.Fprintf(buf, "\tvar err error\n")
+	fmt.Fprintf(buf, "\t_ = err\n")
+
+	g := &generator{
+		buf:         buf,
+		it:          it,
+		indentLevel: 1,
+	}
+
+	g.writef(`buf = append(buf, '{')`)
+	for i := 0; i < t.NumFields(); i++ {
+		fname := t.Field(i).Name()
+		ft := t.Field(i).Type()
+
+		g.writef("")
+		g.writef(`// Encode field %s of type %q`, fname, ft.String())
+
+		// Write the field name; we need to quote the field (for JSON)
+		// and then quote it again (for the generated Go code).
+		qfname := strconv.Quote(fname) + ":"
+		g.writef(`buf = append(buf, []byte(%q)...)`, qfname)
+
+		// Write the value
+		g.encode("self."+fname, ft)
+
+		if i < t.NumFields()-1 {
+			g.writef(`buf = append(buf, ',')`)
+		}
+	}
+	g.writef(`buf = append(buf, '}')`)
+
+	g.writef("return buf, nil")
+	fmt.Fprintf(buf, "}\n\n")
+}
+
+type generator struct {
+	buf         *bytes.Buffer
+	it          *codegen.ImportTracker
+	indentLevel int
+}
+
+func (g *generator) writef(format string, args ...any) {
+	fmt.Fprintf(g.buf, strings.Repeat("\t", g.indentLevel)+format+"\n", args...)
+}
+
+func (g *generator) indent() {
+	g.indentLevel++
+}
+
+func (g *generator) dedent() {
+	g.indentLevel--
+}
+
+func (g *generator) encode(accessor string, ft types.Type) {
+	switch ft := ft.Underlying().(type) {
+	case *types.Basic:
+		g.encodeBasicField(accessor, ft)
+	case *types.Slice:
+		g.encodeSlice(accessor, ft)
+	case *types.Map:
+		g.encodeMap(accessor, ft)
+	case *types.Struct:
+		g.encodeStruct(accessor)
+	case *types.Pointer:
+		g.encodePointer(accessor, ft)
+	default:
+		g.writef(`panic("TODO: %s (%T)")`, accessor, ft)
+	}
+}
+
+func (g *generator) encodePointer(accessor string, ft *types.Pointer) {
+	g.writef("if %s != nil {", accessor)
+	g.indent()
+	// Don't deref for a struct, since we're going to call a function
+	// anyway; otherwise, do.
+	if _, ok := ft.Elem().Underlying().(*types.Struct); ok {
+		g.encode(accessor, ft.Elem())
+	} else {
+		g.encode("(*"+accessor+")", ft.Elem())
+	}
+	g.dedent()
+	g.writef("} else {")
+	g.writef("\tbuf = append(buf, []byte(\"null\")...)")
+	g.writef("}")
+}
+
+func (g *generator) encodeMap(accessor string, ft *types.Map) {
+	kt := ft.Key().Underlying()
+	vt := ft.Elem().Underlying()
+
+	g.writef(`buf = append(buf, '{')`)
+
+	// Determine how we marshal our key type
+	marshalKey := func() {
+		g.encode("k", kt)
+	}
+
+	// Now check how we marshal our value
+	switch vt := vt.(type) {
+	case *types.Basic:
+		g.writef("for k, v := range %s {", accessor)
+		marshalKey()
+		g.writef("\tbuf = append(buf, ':')")
+		g.encodeBasicField("v", vt)
+		g.writef("}")
+	case *types.Struct:
+		g.writef("for k, v := range %s {", accessor)
+		marshalKey()
+		g.writef("\tbuf = append(buf, ':')")
+		g.encodeStruct("v")
+		g.writef("}")
+	default:
+		g.writef(`panic("TODO: %s (%T)")`, accessor, vt)
+	}
+
+	g.writef(`buf = append(buf, '}')`)
+}
+
+func (g *generator) encodeStruct(accessor string) {
+	// Assume that this struct also has a MarshalJSONInto method.
+	g.writef("buf, err = %s.MarshalJSONInto(buf)", accessor)
+	g.writef("if err != nil {")
+	g.writef("\treturn nil, err")
+	g.writef("}")
+}
+
+func (g *generator) encodeSlice(accessor string, sl *types.Slice) {
+	switch ft := sl.Elem().Underlying().(type) {
+	case *types.Basic:
+		// Slice of basic elements
+		switch ft.Kind() {
+		case types.Byte:
+			// base64-encode
+			g.it.Import("encoding/base64")
+
+			g.writef(`buf = append(buf, '"')`)
+			g.writef("{")
+
+			// buf = append(buf, make([]byte, N)...) is a fast way to grow the slice by N
+			g.writef("encodedLen := base64.StdEncoding.EncodedLen(len(%s))", accessor)
+			g.writef("offset := len(buf)")
+			g.writef("buf = append(buf, make([]byte, encodedLen)...)")
+			g.writef("base64.StdEncoding.Encode(buf[offset:], %s)", accessor)
+
+			g.writef("}")
+			g.writef(`buf = append(buf, '"')`)
+		default:
+			// All other basic elements are encoded
+			// one at a time via encodeBasicField
+			g.writef(`buf = append(buf, '[')`)
+			g.writef(`for i, elem := range %s {`, accessor)
+			g.writef("\tif i > 0 {")
+			g.writef("\t\tbuf = append(buf, ',')")
+			g.writef("\t}")
+			g.encodeBasicField("elem", ft)
+			g.writef(`}`)
+			g.writef(`buf = append(buf, ']')`)
+		}
+
+	case *types.Struct:
+		g.writef(`buf = append(buf, '[')`)
+		g.writef(`for i, elem := range %s {`, accessor)
+		g.writef("\tif i > 0 {")
+		g.writef("\t\tbuf = append(buf, ',')")
+		g.writef("\t}")
+		g.encodeStruct("elem")
+		g.writef(`}`)
+		g.writef(`buf = append(buf, ']')`)
+
+	default:
+		// TODO: if the type implements our interface,
+		// call that function for everything in the
+		// slice.
+		g.writef(`panic("TODO: %s (%T)")`, accessor, ft)
+	}
+}
+
+func (g *generator) encodeBasicField(accessor string, field *types.Basic) {
+	switch field.Kind() {
+	case types.Bool:
+		g.writef("if %s {", accessor)
+		g.writef(`buf = append(buf, []byte("true")...)`)
+		g.writef("} else {")
+		g.writef(`buf = append(buf, []byte("false")...)`)
+		g.writef("}")
+	case types.Int, types.Int8, types.Int16, types.Int32, types.Int64:
+		g.it.Import("strconv")
+		g.writef("buf = strconv.AppendInt(buf, int64(%s), 10)", accessor)
+	case types.Uint, types.Uint8, types.Uint16, types.Uint32, types.Uint64:
+		g.it.Import("strconv")
+		g.writef("buf = strconv.AppendUint(buf, uint64(%s), 10)", accessor)
+	case types.String:
+		g.it.Import("strconv")
+		g.writef("buf = strconv.AppendQuote(buf, %s)", accessor)
+	default:
+		g.writef(`panic("TODO: %s (%T)")`, accessor, field.Kind)
+	}
+}
+
+func loadTypes(pkgName, buildTags string) (*packages.Package, map[string]*types.Named, error) {
+	cfg := &packages.Config{
+		Mode: packages.NeedTypes |
+			packages.NeedTypesInfo |
+			packages.NeedSyntax |
+			packages.NeedName,
+		Tests: false,
+	}
+	if buildTags != "" {
+		cfg.BuildFlags = []string{"-tags=" + buildTags}
+	}
+
+	pkgs, err := packages.Load(cfg, pkgName)
+	if err != nil {
+		return nil, nil, err
+	}
+	if len(pkgs) != 1 {
+		return nil, nil, fmt.Errorf("wrong number of packages: %d", len(pkgs))
+	}
+	pkg := pkgs[0]
+	return pkg, namedTypes(pkg), nil
+}
+
+func namedTypes(pkg *packages.Package) map[string]*types.Named {
+	nt := make(map[string]*types.Named)
+	for _, file := range pkg.Syntax {
+		for _, d := range file.Decls {
+			decl, ok := d.(*ast.GenDecl)
+			if !ok || decl.Tok != token.TYPE {
+				continue
+			}
+			for _, s := range decl.Specs {
+				spec, ok := s.(*ast.TypeSpec)
+				if !ok {
+					continue
+				}
+				typeNameObj, ok := pkg.TypesInfo.Defs[spec.Name]
+				if !ok {
+					continue
+				}
+				typ, ok := typeNameObj.Type().(*types.Named)
+				if !ok {
+					continue
+				}
+				nt[spec.Name.Name] = typ
+			}
+		}
+	}
+	return nt
+}
diff --git a/cmd/fastjson/testcodegen/gen_test.go b/cmd/fastjson/testcodegen/gen_test.go
new file mode 100644
index 000000000..6357b1d22
--- /dev/null
+++ b/cmd/fastjson/testcodegen/gen_test.go
@@ -0,0 +1,71 @@
+package testcodegen
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func testObj() *PingRequest {
+	var ival int = 123
+	mp1 := &ival
+	mp2 := &mp1
+
+	obj := &PingRequest{
+		URL:        "https://example.com",
+		Log:        true,
+		Types:      "TODO",
+		IP:         "127.0.0.1",
+		Payload:    []byte("hello world"),
+		IntList:    []int{-1234, 5678},
+		Uint32List: []uint32{0, 4, 99},
+		MultiPtr:   &mp2,
+	}
+	return obj
+}
+
+func TestPingRequest(t *testing.T) {
+	obj := testObj()
+	out, err := obj.MarshalJSONInto(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	const expected = `{"URL":"https://example.com","URLIsNoise":true,"Log":true,"Types":"TODO","IP":"127.0.0.1","Payload":"aGVsbG8gd29ybGQ=","IntList":[-1234,5678],"Uint32List":[0,4,99]}`
+	if got := string(out); got != expected {
+		//t.Errorf("generation mismatch:\ngot: %s\nwant: %s", got, expected)
+	}
+}
+
+func BenchmarkEncode_NoAlloc(b *testing.B) {
+	obj := testObj()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		_, _ = obj.MarshalJSONInto(nil)
+	}
+}
+
+func BenchmarkEncode_Alloc(b *testing.B) {
+	obj := testObj()
+	buf := make([]byte, 0, 10)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		buf, _ = obj.MarshalJSONInto(buf[:0])
+	}
+}
+
+func BenchmarkStd(b *testing.B) {
+	obj := testObj()
+	_, err := json.Marshal(obj)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		_, _ = json.Marshal(obj)
+	}
+}
diff --git a/cmd/fastjson/testcodegen/json_gen.go b/cmd/fastjson/testcodegen/json_gen.go
new file mode 100644
index 000000000..02b9cb5bb
--- /dev/null
+++ b/cmd/fastjson/testcodegen/json_gen.go
@@ -0,0 +1,141 @@
+// Code generated by TODO; DO NOT EDIT.
+
+package testcodegen
+
+import (
+	"encoding/base64"
+	"strconv"
+)
+
+// MarshalJSONInto marshals this PingRequest into JSON in the provided buffer.
+func (self *PingRequest) MarshalJSONInto(buf []byte) ([]byte, error) {
+	var err error
+	_ = err
+	buf = append(buf, '{')
+
+	// Encode field URL of type "string"
+	buf = append(buf, []byte("\"URL\":")...)
+	buf = strconv.AppendQuote(buf, self.URL)
+	buf = append(buf, ',')
+
+	// Encode field URLIsNoise of type "bool"
+	buf = append(buf, []byte("\"URLIsNoise\":")...)
+	if self.URLIsNoise {
+		buf = append(buf, []byte("true")...)
+	} else {
+		buf = append(buf, []byte("false")...)
+	}
+	buf = append(buf, ',')
+
+	// Encode field Log of type "bool"
+	buf = append(buf, []byte("\"Log\":")...)
+	if self.Log {
+		buf = append(buf, []byte("true")...)
+	} else {
+		buf = append(buf, []byte("false")...)
+	}
+	buf = append(buf, ',')
+
+	// Encode field Types of type "string"
+	buf = append(buf, []byte("\"Types\":")...)
+	buf = strconv.AppendQuote(buf, self.Types)
+	buf = append(buf, ',')
+
+	// Encode field IP of type "string"
+	buf = append(buf, []byte("\"IP\":")...)
+	buf = strconv.AppendQuote(buf, self.IP)
+	buf = append(buf, ',')
+
+	// Encode field Payload of type "[]byte"
+	buf = append(buf, []byte("\"Payload\":")...)
+	buf = append(buf, '"')
+	{
+		encodedLen := base64.StdEncoding.EncodedLen(len(self.Payload))
+		offset := len(buf)
+		buf = append(buf, make([]byte, encodedLen)...)
+		base64.StdEncoding.Encode(buf[offset:], self.Payload)
+	}
+	buf = append(buf, '"')
+	buf = append(buf, ',')
+
+	// Encode field IntList of type "[]int"
+	buf = append(buf, []byte("\"IntList\":")...)
+	buf = append(buf, '[')
+	for i, elem := range self.IntList {
+		if i > 0 {
+			buf = append(buf, ',')
+		}
+		buf = strconv.AppendInt(buf, int64(elem), 10)
+	}
+	buf = append(buf, ']')
+	buf = append(buf, ',')
+
+	// Encode field Uint32List of type "[]uint32"
+	buf = append(buf, []byte("\"Uint32List\":")...)
+	buf = append(buf, '[')
+	for i, elem := range self.Uint32List {
+		if i > 0 {
+			buf = append(buf, ',')
+		}
+		buf = strconv.AppendUint(buf, uint64(elem), 10)
+	}
+	buf = append(buf, ']')
+	buf = append(buf, ',')
+
+	// Encode field StringPtr of type "*string"
+	buf = append(buf, []byte("\"StringPtr\":")...)
+	if self.StringPtr != nil {
+		buf = strconv.AppendQuote(buf, (*self.StringPtr))
+	} else {
+		buf = append(buf, []byte("null")...)
+	}
+	buf = append(buf, ',')
+
+	// Encode field StructPtr of type "*tailscale.com/cmd/fastjson/testcodegen.OtherStruct"
+	buf = append(buf, []byte("\"StructPtr\":")...)
+	if self.StructPtr != nil {
+		buf, err = self.StructPtr.MarshalJSONInto(buf)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		buf = append(buf, []byte("null")...)
+	}
+	buf = append(buf, ',')
+
+	// Encode field MultiPtr of type "***int"
+	buf = append(buf, []byte("\"MultiPtr\":")...)
+	if self.MultiPtr != nil {
+		if (*self.MultiPtr) != nil {
+			if (*(*self.MultiPtr)) != nil {
+				buf = strconv.AppendInt(buf, int64((*(*(*self.MultiPtr)))), 10)
+			} else {
+				buf = append(buf, []byte("null")...)
+			}
+		} else {
+			buf = append(buf, []byte("null")...)
+		}
+	} else {
+		buf = append(buf, []byte("null")...)
+	}
+	buf = append(buf, '}')
+	return buf, nil
+}
+
+// MarshalJSONInto marshals this OtherStruct into JSON in the provided buffer.
+func (self *OtherStruct) MarshalJSONInto(buf []byte) ([]byte, error) {
+	var err error
+	_ = err
+	buf = append(buf, '{')
+
+	// Encode field Name of type "string"
+	buf = append(buf, []byte("\"Name\":")...)
+	buf = strconv.AppendQuote(buf, self.Name)
+	buf = append(buf, ',')
+
+	// Encode field Age of type "int"
+	buf = append(buf, []byte("\"Age\":")...)
+	buf = strconv.AppendInt(buf, int64(self.Age), 10)
+	buf = append(buf, '}')
+	return buf, nil
+}
diff --git a/cmd/fastjson/testcodegen/ping_request.go b/cmd/fastjson/testcodegen/ping_request.go
new file mode 100644
index 000000000..a0095b740
--- /dev/null
+++ b/cmd/fastjson/testcodegen/ping_request.go
@@ -0,0 +1,69 @@
+package testcodegen
+
+// PingRequest with no IP and Types is a request to send an HTTP request to prove the
+// long-polling client is still connected.
+// PingRequest with Types and IP, will send a ping to the IP and send a POST
+// request containing a PingResponse to the URL containing results.
+type PingRequest struct {
+	// URL is the URL to reply to the PingRequest to.
+	// It will be a unique URL each time. No auth headers are necessary.
+	// If the client sees multiple PingRequests with the same URL,
+	// subsequent ones should be ignored.
+	//
+	// The HTTP method that the node should make back to URL depends on the other
+	// fields of the PingRequest. If Types is defined, then URL is the URL to
+	// send a POST request to. Otherwise, the node should just make a HEAD
+	// request to URL.
+	URL string
+
+	// URLIsNoise, if true, means that the client should hit URL over the Noise
+	// transport instead of TLS.
+	URLIsNoise bool `json:",omitempty"`
+
+	// Log is whether to log about this ping in the success case.
+	// For failure cases, the client will log regardless.
+	Log bool `json:",omitempty"`
+
+	// Types is the types of ping that are initiated. Can be any PingType, comma
+	// separated, e.g. "disco,TSMP"
+	//
+	// As a special case, if Types is "c2n", then this PingRequest is a
+	// client-to-node HTTP request. The HTTP request should be handled by this
+	// node's c2n handler and the HTTP response sent in a POST to URL. For c2n,
+	// the value of URLIsNoise is ignored and only the Noise transport (back to
+	// the control plane) will be used, as if URLIsNoise were true.
+	Types string `json:",omitempty"`
+
+	// IP is the ping target, when needed by the PingType(s) given in Types.
+	IP string
+
+	// Payload is the ping payload.
+	//
+	// It is only used for c2n requests, in which case it's an HTTP/1.0 or
+	// HTTP/1.1-formatted HTTP request as parsable with http.ReadRequest.
+	Payload []byte `json:",omitempty"`
+
+	IntList    []int
+	Uint32List []uint32
+
+	StringPtr *string
+	StructPtr *OtherStruct
+	MultiPtr  ***int
+
+	/*
+		Kv1 map[string]int
+		Kv2 map[int]bool
+	*/
+
+	/*
+		Other       OtherStruct
+		OtherSlice  []OtherStruct
+		OtherMap    map[string]OtherStruct
+		OtherKeyMap map[OtherStruct]bool
+	*/
+}
+
+type OtherStruct struct {
+	Name string
+	Age  int
+}