From 8119ed2342c7c38e5d6fcf018341f2f86d6af3df Mon Sep 17 00:00:00 2001 From: Elias Naur Date: Fri, 1 Dec 2023 10:25:44 -0600 Subject: [PATCH] wip one-alloc unmarshal --- features/zeropb/zeropb.go | 67 +++++++++++++++++++++++++++++++++------ testpb/1.pulsar.go | 67 ++++++++++++++++++++++++++++++++------- testpb/2.pulsar.go | 6 +++- 3 files changed, 117 insertions(+), 23 deletions(-) diff --git a/features/zeropb/zeropb.go b/features/zeropb/zeropb.go index 4ace5ba..b2b8798 100644 --- a/features/zeropb/zeropb.go +++ b/features/zeropb/zeropb.go @@ -19,6 +19,8 @@ const ( runtimePackage = protogen.GoImportPath("github.com/cosmos/cosmos-proto/runtime/zeropb") ) +const unsafeOptimizations = true + func init() { generator.RegisterFeature("zeropb", func(gen *generator.GeneratedFile, _ *protogen.Plugin) generator.FeatureGenerator { return zeropbFeature{ @@ -179,8 +181,6 @@ func (g zeropbFeature) generateMarshalPrimitive(f *protogen.Field, name, offset } } -const unsafeOptimizations = true - func (g zeropbFeature) generateUnmarshal(m *protogen.Message) { g.gen.P("func (x *", m.GoIdent, ") UnmarshalZeroPB(buf []byte) (err error) {") g.gen.P(" defer func() {") @@ -190,14 +190,23 @@ func (g zeropbFeature) generateUnmarshal(m *protogen.Message) { g.gen.P(" }()") g.gen.P(" var mem []byte") if unsafeOptimizations { - g.gen.P("mem = make([]byte, 251)") + g.gen.P("memSize := _", m.GoIdent, "UnmarshalZeroPBSize(buf, 0)") + g.gen.P("mem = make([]byte, memSize)") } g.gen.P(" x.unmarshalZeroPB(buf, 0, ", runtimePackage.Ident("NewBuffer"), "(mem))") g.gen.P(" return nil") g.gen.P("}") g.gen.P() - g.gen.P("func (x *", m.GoIdent, ") unmarshalZeroPB(buf []byte, n uint16, mem *", runtimePackage.Ident("Buffer"), ") {") + g.gen.P("func _", m.GoIdent, "UnmarshalZeroPBSize(buf []byte, n uint16) (size uint16) {") offset := 0 + for _, f := range m.Fields { + g.generateUnmarshalFieldSize(f, m.GoIdent, offset) + offset += fieldSize(f) + } + g.gen.P(" return") + g.gen.P("}") + g.gen.P("func (x *", m.GoIdent, ") unmarshalZeroPB(buf []byte, n uint16, mem *", runtimePackage.Ident("Buffer"), ") {") + offset = 0 for _, f := range m.Fields { g.generateUnmarshalField(f, offset) offset += fieldSize(f) @@ -205,6 +214,47 @@ func (g zeropbFeature) generateUnmarshal(m *protogen.Message) { g.gen.P("}") } +func (g zeropbFeature) generateUnmarshalFieldSize(f *protogen.Field, typ protogen.GoIdent, offset int) { + d := f.Desc + switch { + case d.IsList(): + g.gen.P("n_", d.Index(), ", len_", d.Index(), " := ", runtimePackage.Ident("ReadSlice"), "(buf, n+", offset, ")") + g.gen.P("_ = n_", d.Index()) + g.gen.P("size += len_", d.Index(), "*uint16(", unsafePackage.Ident("Sizeof"), "((", typ, "{}).", f.GoName, "[0]))") + g.gen.P("for i := uint16(0); i < len_", d.Index(), "; i++ {") + // Skip segment header. + g.generateUnmarshalPrimitiveSize(f, fmt.Sprintf("n_%d+%d+uint16(i)*%d", d.Index(), segmentHeaderSize, fieldElemSize(f))) + g.gen.P("}") + case d.IsMap(): + g.gen.P("n_", d.Index(), ", len_", d.Index(), " := ", runtimePackage.Ident("ReadSlice"), "(buf, n+", offset, ")") + g.gen.P("{") + g.gen.P(" n := n_", d.Index(), "; _ = n") + g.gen.P(" for i := uint16(0); i < len_", d.Index(), "; i++ {") + g.generateUnmarshalPrimitiveSize(f.Message.Fields[0], "n") + g.gen.P(" n += ", fieldSize(f.Message.Fields[0])) + g.generateUnmarshalPrimitiveSize(f.Message.Fields[1], "n") + g.gen.P(" n += ", fieldSize(f.Message.Fields[1])) + g.gen.P(" }") + g.gen.P("}") + case d.ContainingOneof() != nil: + g.gen.P("// TODO: field ", f.GoName) + default: + g.generateUnmarshalPrimitiveSize(f, fmt.Sprintf("n+%d", offset)) + } +} + +func (g zeropbFeature) generateUnmarshalPrimitiveSize(f *protogen.Field, offset string) { + switch d := f.Desc; d.Kind() { + case protoreflect.StringKind, protoreflect.BytesKind: + g.gen.P("_, len_", d.Index(), " := ", runtimePackage.Ident("ReadSlice"), "(buf, ", offset, ")") + g.gen.P("size += len_", d.Index(), "*uint16(", unsafePackage.Ident("Sizeof"), "(byte(0)))") + case protoreflect.MessageKind: + g.gen.P("size += uint16(", unsafePackage.Ident("Sizeof"), "(", f.Message.GoIdent, "{}))") + g.gen.P("size += _", f.Message.GoIdent, "UnmarshalZeroPBSize(buf, ", offset, ")") + default: + } +} + func (g zeropbFeature) generateUnmarshalField(f *protogen.Field, offset int) { d := f.Desc switch { @@ -217,11 +267,8 @@ func (g zeropbFeature) generateUnmarshalField(f *protogen.Field, offset int) { } else { g.gen.P("x.", f.GoName, " = make([]", typ, ", len_", d.Index(), ")") } - g.gen.P("{") - g.gen.P(" for i := range x.", f.GoName, "{") - // Skip segment header. - g.generateUnmarshalPrimitive(f, "x."+f.GoName+"[i]", fmt.Sprintf("n_%d+4+uint16(i)*%d", d.Index(), fieldElemSize(f))) - g.gen.P(" }") + g.gen.P("for i := range x.", f.GoName, "{") + g.generateUnmarshalPrimitive(f, "x."+f.GoName+"[i]", fmt.Sprintf("n_%d+%d+uint16(i)*%d", d.Index(), segmentHeaderSize, fieldElemSize(f))) g.gen.P("}") case d.IsMap(): g.gen.P("n_", d.Index(), ", len_", d.Index(), " := ", runtimePackage.Ident("ReadSlice"), "(buf, n+", offset, ")") @@ -280,7 +327,7 @@ func (g zeropbFeature) generateUnmarshalPrimitive(f *protogen.Field, name, offse if d.Kind() == protoreflect.BytesKind { g.gen.P(name, " = mem_", d.Index(), ".Buf") } else { - g.gen.P(name, " = ", unsafePackage.Ident("String"), "(unsafe.SliceData(mem_", d.Index(), ".Buf), len(mem_", d.Index(), ".Buf))") + g.gen.P(name, " = ", unsafePackage.Ident("String"), "(", unsafePackage.Ident("SliceData"), "(mem_", d.Index(), ".Buf), len(mem_", d.Index(), ".Buf))") } } else { if d.Kind() == protoreflect.BytesKind { diff --git a/testpb/1.pulsar.go b/testpb/1.pulsar.go index 56b109b..825b7bc 100644 --- a/testpb/1.pulsar.go +++ b/testpb/1.pulsar.go @@ -3207,11 +3207,52 @@ func (x *A) UnmarshalZeroPB(buf []byte) (err error) { } }() var mem []byte - mem = make([]byte, 251) + memSize := _AUnmarshalZeroPBSize(buf, 0) + mem = make([]byte, memSize) x.unmarshalZeroPB(buf, 0, zeropb.NewBuffer(mem)) return nil } +func _AUnmarshalZeroPBSize(buf []byte, n uint16) (size uint16) { + _, len_14 := zeropb.ReadSlice(buf, n+80) + size += len_14 * uint16(unsafe.Sizeof(byte(0))) + _, len_15 := zeropb.ReadSlice(buf, n+84) + size += len_15 * uint16(unsafe.Sizeof(byte(0))) + size += uint16(unsafe.Sizeof(B{})) + size += _BUnmarshalZeroPBSize(buf, n+88) + n_17, len_17 := zeropb.ReadSlice(buf, n+92) + { + n := n_17 + _ = n + for i := uint16(0); i < len_17; i++ { + _, len_0 := zeropb.ReadSlice(buf, n) + size += len_0 * uint16(unsafe.Sizeof(byte(0))) + n += 4 + size += uint16(unsafe.Sizeof(B{})) + size += _BUnmarshalZeroPBSize(buf, n) + n += 4 + } + } + n_18, len_18 := zeropb.ReadSlice(buf, n+96) + _ = n_18 + size += len_18 * uint16(unsafe.Sizeof((A{}).LIST[0])) + for i := uint16(0); i < len_18; i++ { + size += uint16(unsafe.Sizeof(B{})) + size += _BUnmarshalZeroPBSize(buf, n_18+4+uint16(i)*4) + } + // TODO: field ONEOF_B + // TODO: field ONEOF_STRING + n_21, len_21 := zeropb.ReadSlice(buf, n+108) + _ = n_21 + size += len_21 * uint16(unsafe.Sizeof((A{}).LIST_ENUM[0])) + for i := uint16(0); i < len_21; i++ { + } + size += uint16(unsafe.Sizeof(ImportedMessage{})) + size += _ImportedMessageUnmarshalZeroPBSize(buf, n+112) + _, len_23 := zeropb.ReadSlice(buf, n+112) + size += len_23 * uint16(unsafe.Sizeof(byte(0))) + return +} func (x *A) unmarshalZeroPB(buf []byte, n uint16, mem *zeropb.Buffer) { x.Enum = Enumeration(binary.LittleEndian.Uint32(buf[n+0:])) bool_1 := binary.LittleEndian.Uint32(buf[n+4:]) @@ -3267,22 +3308,18 @@ func (x *A) unmarshalZeroPB(buf []byte, n uint16, mem *zeropb.Buffer) { n_18, len_18 := zeropb.ReadSlice(buf, n+96) mem_18 := mem.Alloc(int(len_18) * int(unsafe.Sizeof(x.LIST[0]))) x.LIST = unsafe.Slice((**B)(unsafe.Pointer(unsafe.SliceData(mem_18.Buf))), len_18) - { - for i := range x.LIST { - mem_18 := mem.Alloc(int(unsafe.Sizeof(*x.LIST[i]))) - x.LIST[i] = (*B)(unsafe.Pointer(unsafe.SliceData(mem_18.Buf))) - x.LIST[i].unmarshalZeroPB(buf, n_18+4+uint16(i)*4, mem) - } + for i := range x.LIST { + mem_18 := mem.Alloc(int(unsafe.Sizeof(*x.LIST[i]))) + x.LIST[i] = (*B)(unsafe.Pointer(unsafe.SliceData(mem_18.Buf))) + x.LIST[i].unmarshalZeroPB(buf, n_18+4+uint16(i)*4, mem) } // TODO: field ONEOF_B // TODO: field ONEOF_STRING n_21, len_21 := zeropb.ReadSlice(buf, n+108) mem_21 := mem.Alloc(int(len_21) * int(unsafe.Sizeof(x.LIST_ENUM[0]))) x.LIST_ENUM = unsafe.Slice((*Enumeration)(unsafe.Pointer(unsafe.SliceData(mem_21.Buf))), len_21) - { - for i := range x.LIST_ENUM { - x.LIST_ENUM[i] = Enumeration(binary.LittleEndian.Uint32(buf[n_21+4+uint16(i)*4:])) - } + for i := range x.LIST_ENUM { + x.LIST_ENUM[i] = Enumeration(binary.LittleEndian.Uint32(buf[n_21+4+uint16(i)*4:])) } mem_22 := mem.Alloc(int(unsafe.Sizeof(*x.Imported))) x.Imported = (*ImportedMessage)(unsafe.Pointer(unsafe.SliceData(mem_22.Buf))) @@ -3315,11 +3352,17 @@ func (x *B) UnmarshalZeroPB(buf []byte) (err error) { } }() var mem []byte - mem = make([]byte, 251) + memSize := _BUnmarshalZeroPBSize(buf, 0) + mem = make([]byte, memSize) x.unmarshalZeroPB(buf, 0, zeropb.NewBuffer(mem)) return nil } +func _BUnmarshalZeroPBSize(buf []byte, n uint16) (size uint16) { + _, len_0 := zeropb.ReadSlice(buf, n+0) + size += len_0 * uint16(unsafe.Sizeof(byte(0))) + return +} func (x *B) unmarshalZeroPB(buf []byte, n uint16, mem *zeropb.Buffer) { n_0, len_0 := zeropb.ReadSlice(buf, n+0) buf_0 := buf[n_0 : n_0+len_0] diff --git a/testpb/2.pulsar.go b/testpb/2.pulsar.go index 5e12ed8..4513cf8 100644 --- a/testpb/2.pulsar.go +++ b/testpb/2.pulsar.go @@ -502,10 +502,14 @@ func (x *ImportedMessage) UnmarshalZeroPB(buf []byte) (err error) { } }() var mem []byte - mem = make([]byte, 251) + memSize := _ImportedMessageUnmarshalZeroPBSize(buf, 0) + mem = make([]byte, memSize) x.unmarshalZeroPB(buf, 0, zeropb.NewBuffer(mem)) return nil } +func _ImportedMessageUnmarshalZeroPBSize(buf []byte, n uint16) (size uint16) { + return +} func (x *ImportedMessage) unmarshalZeroPB(buf []byte, n uint16, mem *zeropb.Buffer) { }