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 emitJSONParseInteger(g *protogen.GeneratedFile, ident string, method string, bitSize int, typ string) {
	g.P("r := in.JsonNumber()")
	g.P("n := r.String()")
	g.P("v, err := ", strconvPackage.Ident(method), "(n, 10, ", bitSize, ")")
	g.P("if err != nil {")
	g.P("	in.AddError(err)")
	g.P("	return")
	g.P("}")
	g.P(ident, " := ", typ, "(v)")
}

func emitJSONReadEnum(g *protogen.GeneratedFile, name string, enumType string) {
	g.P(`switch v := in.Interface().(type) {
	case string:
		if vv, ok := `+enumType+`_value[v]; ok {
			`+name+` = `+enumType+`(vv)
			break
		}
		vv, err := `, strconvPackage.Ident("ParseInt"), `(v, 10, 32)
		if err != nil {
			in.AddError(err)
			return
		}
		`+name+` = `+enumType+`(vv)
	case float64:
		`+name+` = `+enumType+`(v)
	}`)
}

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)
		emitJSONReadEnum(g, "parsedValue", enumType)
		template = "%s = parsedValue"
	case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind:
		emitJSONParseInteger(g, "pv", "ParseInt", 32, "int32")
		template = "%s = pv"
	case protoreflect.Uint32Kind, protoreflect.Fixed32Kind:
		emitJSONParseInteger(g, "pv", "ParseUint", 32, "uint32")
		template = "%s = pv"
	case protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind:
		emitJSONParseInteger(g, "pv", "ParseInt", 64, "int64")
		template = "%s = pv"
	case protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
		emitJSONParseInteger(g, "pv", "ParseUint", 64, "uint64")
		template = "%s = pv"
	case protoreflect.FloatKind:
		template = "%s = in.Float32()"
	case protoreflect.DoubleKind:
		template = "%s = in.Float64()"
	case protoreflect.StringKind:
		template = "%s = in.String()"
	case protoreflect.BytesKind:
		// Since some time ago proto3 support optional keyword, thus the presence is not tracked by default:
		// https://github.com/protocolbuffers/protobuf-go/blob/fb995f184a1719ec42b247a3771d1036d92adf67/internal/impl/message_reflect_field.go#L327
		// We do not explicitly support `optional` keyword, protoc will fail on such fileds.
		// Thus, treat empty string as `[]byte(nil)`.
		template = `{
			tmp := in.Bytes()
			if len(tmp) == 0 {
				tmp = nil
			}
			%s = tmp
		}`
	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 }`)

	if len(msg.Fields) != 0 {
		g.P("first := true")
	}
	g.P("out.RawByte('{')")
	for _, 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")
			}
			g.P("}")
			continue
		}
		emitJSONFieldWrite(g, f, "x")
	}
	g.P("out.RawByte('}')")
	g.P("}")
}

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

	selector := name + "." + f.GoName

	// This code is responsible for ignoring default values.
	// We will restore it after having parametrized JSON marshaling.
	//
	// isNotDefault := notNil
	// if f.Desc.IsList() {
	// 	isNotDefault = notEmpty
	// } else if f.Desc.Kind() != protoreflect.MessageKind {
	// 	_, isNotDefault = easyprotoKindInfo(f.Desc.Kind())
	// }
	// g.P("if ", isNotDefault(selector), "{")
	// defer g.P("}")

	g.P("if !first { out.RawByte(','); } else { first = false; }")
	g.P("const prefix string = ", `"\"`, fieldJSONName(f), `\":"`)
	g.P("out.RawString(prefix)")
	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:
		g.P("out.RawByte('\"')")
		g.P("out.Buffer.Buf = ", strconvPackage.Ident("AppendInt"), "(out.Buffer.Buf, ", selector, ", 10)")
		g.P("out.RawByte('\"')")
		return
	case protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
		g.P("out.RawByte('\"')")
		g.P("out.Buffer.Buf = ", strconvPackage.Ident("AppendUint"), "(out.Buffer.Buf, ", selector, ", 10)")
		g.P("out.RawByte('\"')")
		return
	case protoreflect.FloatKind:
		template = "out.Float32(%s)"
	case protoreflect.DoubleKind:
		template = "out.Float64(%s)"
	case protoreflect.StringKind:
		template = "out.String(%s)"
	case protoreflect.BytesKind:
		g.P("if ", selector, "!= nil {")
		g.P("out.Base64Bytes(", selector, ")")
		g.P("} else { out.String(\"\") }")
		return
	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())
}