frostfs-api-go/util/protogen/internalgengo/json.go

package internalgengo

import (
	"fmt"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/reflect/protoreflect"
)

var (
	jwriterPackage = protogen.GoImportPath("github.com/mailru/easyjson/jwriter")
	jlexerPackage  = protogen.GoImportPath("github.com/mailru/easyjson/jlexer")
)

func emitJSONMethods(g *protogen.GeneratedFile, msg *protogen.Message) {
	emitJSONMarshal(g, msg)
	emitJSONUnmarshal(g, msg)
}

func emitJSONUnmarshal(g *protogen.GeneratedFile, msg *protogen.Message) {
	g.P("// UnmarshalJSON implements the json.Unmarshaler interface.")
	g.P("func (x *", msg.GoIdent.GoName, ") UnmarshalJSON(data []byte) error {")
	g.P("r := ", jlexerPackage.Ident("Lexer"), "{Data: data}")
	g.P("x.UnmarshalEasyJSON(&r)")
	g.P("return r.Error()")
	g.P("}")

	g.P("func (x *", msg.GoIdent.GoName, ") UnmarshalEasyJSON(in *", jlexerPackage.Ident("Lexer"), ") {")

	g.P("isTopLevel := in.IsStart()")
	g.P("if in.IsNull() {")
	g.P("if isTopLevel { in.Consumed() }")
	g.P("in.Skip()")
	g.P("return")
	g.P("}")

	g.P("in.Delim('{')")
	g.P("for !in.IsDelim('}') {")

	g.P("key := in.UnsafeFieldName(false)")
	g.P("in.WantColon()")
	g.P("if in.IsNull() { in.Skip(); in.WantComma(); continue }")
	g.P("switch key {")
	for _, f := range msg.Fields {
		g.P(`case "`, fieldJSONName(f), `":`)
		if f.Oneof != nil {
			g.P("xx := new(", f.GoIdent, ")")
			g.P("x." + f.Oneof.GoName + " = xx")
			emitJSONFieldRead(g, f, "xx")
			continue
		}
		emitJSONFieldRead(g, f, "x")
	}
	g.P("}")
	g.P("in.WantComma()")
	g.P("}")
	g.P("in.Delim('}')")
	g.P("if isTopLevel { in.Consumed() }")
	g.P("}")
}

func emitJSONFieldRead(g *protogen.GeneratedFile, f *protogen.Field, name string) {
	g.P("{")
	defer g.P("}")

	if f.Desc.IsList() {
		g.P("var f ", fieldType(g, f)[2:])
		g.P("var list ", fieldType(g, f))
		g.P("in.Delim('[')")
		defer g.P("in.Delim(']')")

		g.P("for !in.IsDelim(']') {")
	} else {
		g.P("var f ", fieldType(g, f))
	}

	var template string
	switch f.Desc.Kind() {
	case protoreflect.BoolKind:
		template = "%s = in.Bool()"
	case protoreflect.EnumKind:
		g.Import(strconvPackage)

		enumType := fieldType(g, f).String()
		g.P("var parsedValue " + enumType)
		g.P(`switch v := in.Interface().(type) {
		case string:
			if vv, ok := `+enumType+`_value[v]; ok {
				parsedValue = `+enumType+`(vv)
				break
			}
			vv, err := `, strconvPackage.Ident("ParseInt"), `(v, 10, 32)
			if err != nil {
				in.AddError(err)
				return
			}
			parsedValue = `+enumType+`(vv)
		case float64:
			parsedValue = `+enumType+`(v)
		}`)
		template = "%s = parsedValue"
	case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind:
		template = "%s = in.Int32()"
	case protoreflect.Uint32Kind, protoreflect.Fixed32Kind:
		template = "%s = in.Uint32()"
	case protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind:
		template = "%s = in.Int64()"
	case protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
		template = "%s = in.Uint64()"
	case protoreflect.FloatKind:
		template = "%s = in.Float32()"
	case protoreflect.DoubleKind:
		template = "%s = in.Float64()"
	case protoreflect.StringKind:
		template = "%s = in.String()"
	case protoreflect.BytesKind:
		template = "%s = in.Bytes()"
	case protoreflect.MessageKind:
		if f.Desc.IsList() {
			g.P("f = ", fieldType(g, f)[2:], "{}")
		} else {
			g.P("f = new(", fieldType(g, f)[1:], ")")
		}
		template = "%s.UnmarshalEasyJSON(in)"
	case protoreflect.GroupKind:
		panic("unimplemented")
	}
	g.P(fmt.Sprintf(template, "f"))
	if f.Desc.IsList() {
		g.P("list = append(list, f)")
		g.P("in.WantComma()")
		g.P("}")
		g.P(name, ".", f.GoName, " = list")
	} else {
		g.P(name, ".", f.GoName, " = f")
	}
}

func emitJSONMarshal(g *protogen.GeneratedFile, msg *protogen.Message) {
	g.P("// MarshalJSON implements the json.Marshaler interface.")
	g.P("func (x *", msg.GoIdent.GoName, ") MarshalJSON() ([]byte, error) {")
	g.P("w := ", jwriterPackage.Ident("Writer"), "{}")
	g.P("x.MarshalEasyJSON(&w)")
	g.P("return w.Buffer.BuildBytes(), w.Error")
	g.P("}")

	g.P("func (x *", msg.GoIdent.GoName, ") MarshalEasyJSON(out *", jwriterPackage.Ident("Writer"), ") {")
	g.P(`if x == nil { out.RawString("null"); return }`)

	g.P("out.RawByte('{')")
	for i, f := range msg.Fields {
		if f.Oneof != nil {
			if f.Oneof.Fields[0] != f {
				continue
			}

			g.P("switch xx := x.", f.Oneof.GoName, ".(type) {")
			for _, ff := range f.Oneof.Fields {
				g.P("case *", ff.GoIdent, ":")
				emitJSONFieldWrite(g, ff, "xx", i == 0)
			}
			g.P("}")
			continue
		}
		emitJSONFieldWrite(g, f, "x", i == 0)
	}
	g.P("out.RawByte('}')")
	g.P("}")
}

func emitJSONFieldWrite(g *protogen.GeneratedFile, f *protogen.Field, name string, first bool) {
	g.P("{")
	defer g.P("}")

	g.P("const prefix string = ", `",\"`, fieldJSONName(f), `\":"`)
	if first {
		g.P("out.RawString(prefix[1:])")
	} else {
		g.P("out.RawString(prefix)")
	}

	selector := name + "." + f.GoName
	if f.Desc.IsList() {
		selector += "[i]"
		g.P("out.RawByte('[')")
		defer g.P("out.RawByte(']')")

		g.P("for i := range ", name, ".", f.GoName, " {")
		g.P("if i != 0 { out.RawByte(',') }")
		defer g.P("}")
	}

	var template string
	switch f.Desc.Kind() {
	case protoreflect.BoolKind:
		template = "out.Bool(%s)"
	case protoreflect.EnumKind:
		enumType := fieldType(g, f).String()
		template = `v := int32(%s)
			if vv, ok := ` + enumType + `_name[v]; ok {
				out.String(vv)
			} else {
				out.Int32(v)
			}`
	case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind:
		template = "out.Int32(%s)"
	case protoreflect.Uint32Kind, protoreflect.Fixed32Kind:
		template = "out.Uint32(%s)"
	case protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind:
		template = "out.Int64(%s)"
	case protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
		template = "out.Uint64(%s)"
	case protoreflect.FloatKind:
		template = "out.Float32(%s)"
	case protoreflect.DoubleKind:
		template = "out.Float64(%s)"
	case protoreflect.StringKind:
		template = "out.String(%s)"
	case protoreflect.BytesKind:
		template = "out.Base64Bytes(%s)"
	case protoreflect.MessageKind:
		template = "%s.MarshalEasyJSON(out)"
	case protoreflect.GroupKind:
		panic("unimplemented")
	}
	g.P(fmt.Sprintf(template, selector))
}

func fieldJSONName(f *protogen.Field) string {
	if f.Desc.HasJSONName() {
		return f.Desc.JSONName()
	}
	return string(f.Desc.Name())
}