From 8ff6816f078f902a17a93fc302f24dee9993166a Mon Sep 17 00:00:00 2001 From: Shambhavi Srivastava Date: Mon, 7 Aug 2023 11:40:11 -0700 Subject: [PATCH] Implementing CopyInN PiperOrigin-RevId: 554542787 --- pkg/marshal/marshal.go | 8 +++++ pkg/marshal/primitive/primitive.go | 5 +++ .../generator_interfaces_array_newtype.go | 13 +++++-- .../gomarshal/generator_interfaces_dynamic.go | 14 ++++++-- .../generator_interfaces_primitive_newtype.go | 13 +++++-- .../gomarshal/generator_interfaces_struct.go | 18 +++++++--- tools/go_marshal/test/marshal_test.go | 34 +++++++++++++++++++ 7 files changed, 93 insertions(+), 12 deletions(-) diff --git a/pkg/marshal/marshal.go b/pkg/marshal/marshal.go index e60eacb3a..a8af4ede2 100644 --- a/pkg/marshal/marshal.go +++ b/pkg/marshal/marshal.go @@ -112,6 +112,14 @@ type Marshallable interface { // for UnmarshalBytes. CopyIn(cc CopyContext, addr hostarch.Addr) (int, error) + // CopyInN is like CopyIn, but explicitly requests a partial + // copy-in. Note that this may yield unexpected results for non-packed + // types and the caller may only want to allow this for packed types. See + // comment on UnmarshalBytes. + // + // The limit must be less than or equal to SizeBytes(). + CopyInN(cc CopyContext, addr hostarch.Addr, limit int) (int, error) + // CopyOut serializes a Marshallable type to a task's memory. This may only // be called from a task goroutine. This is more efficient than calling // MarshalUnsafe on Marshallable.Packed types, as the type being serialized diff --git a/pkg/marshal/primitive/primitive.go b/pkg/marshal/primitive/primitive.go index bea5e0b89..f5095768e 100644 --- a/pkg/marshal/primitive/primitive.go +++ b/pkg/marshal/primitive/primitive.go @@ -105,6 +105,11 @@ func (b *ByteSlice) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, err return cc.CopyInBytes(addr, *b) } +// CopyInN implements marshal.Marshallable.CopyInN. +func (b *ByteSlice) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) { + return cc.CopyInBytes(addr, (*b)[:limit]) +} + // CopyOut implements marshal.Marshallable.CopyOut. func (b *ByteSlice) CopyOut(cc marshal.CopyContext, addr hostarch.Addr) (int, error) { return cc.CopyOutBytes(addr, *b) diff --git a/tools/go_marshal/gomarshal/generator_interfaces_array_newtype.go b/tools/go_marshal/gomarshal/generator_interfaces_array_newtype.go index 5168387f6..1880606dd 100644 --- a/tools/go_marshal/gomarshal/generator_interfaces_array_newtype.go +++ b/tools/go_marshal/gomarshal/generator_interfaces_array_newtype.go @@ -124,17 +124,24 @@ func (g *interfaceGenerator) emitMarshallableForArrayNewtype(n *ast.Ident, a *as }) g.emit("}\n\n") - g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") - g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n") + g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName()) g.inIndent(func() { g.emitCastToByteSlice(g.r, "buf", fmt.Sprintf("%s.SizeBytes()", g.r)) - g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n") + g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n") g.emitKeepAlive(g.r) g.emit("return length, err\n") }) g.emit("}\n\n") + g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") + g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.inIndent(func() { + g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r) + }) + g.emit("}\n\n") + g.emit("// WriteTo implements io.WriterTo.WriteTo.\n") g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName()) g.inIndent(func() { diff --git a/tools/go_marshal/gomarshal/generator_interfaces_dynamic.go b/tools/go_marshal/gomarshal/generator_interfaces_dynamic.go index e29b05809..a1c8a8a15 100644 --- a/tools/go_marshal/gomarshal/generator_interfaces_dynamic.go +++ b/tools/go_marshal/gomarshal/generator_interfaces_dynamic.go @@ -65,10 +65,11 @@ func (g *interfaceGenerator) emitMarshallableForDynamicType() { }) g.emit("}\n\n") - g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") + g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n") + g.emit("//go:nosplit\n") g.recordUsedImport("marshal") g.recordUsedImport("hostarch") - g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName()) g.inIndent(func() { g.emit("// Type %s doesn't have a packed layout in memory, fall back to UnmarshalBytes.\n", g.typeName()) g.emit("buf := cc.CopyScratchBuffer(%s.SizeBytes()) // escapes: okay.\n", g.r) @@ -80,6 +81,15 @@ func (g *interfaceGenerator) emitMarshallableForDynamicType() { }) g.emit("}\n\n") + g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") + g.recordUsedImport("marshal") + g.recordUsedImport("hostarch") + g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.inIndent(func() { + g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r) + }) + g.emit("}\n\n") + g.emit("// WriteTo implements io.WriterTo.WriteTo.\n") g.recordUsedImport("io") g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName()) diff --git a/tools/go_marshal/gomarshal/generator_interfaces_primitive_newtype.go b/tools/go_marshal/gomarshal/generator_interfaces_primitive_newtype.go index e2ff7202c..d76f77e54 100644 --- a/tools/go_marshal/gomarshal/generator_interfaces_primitive_newtype.go +++ b/tools/go_marshal/gomarshal/generator_interfaces_primitive_newtype.go @@ -184,17 +184,24 @@ func (g *interfaceGenerator) emitMarshallableForPrimitiveNewtype(nt *ast.Ident) }) g.emit("}\n\n") - g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") - g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n") + g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName()) g.inIndent(func() { g.emitCastToByteSlice(g.r, "buf", fmt.Sprintf("%s.SizeBytes()", g.r)) - g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n") + g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n") g.emitKeepAlive(g.r) g.emit("return length, err\n") }) g.emit("}\n\n") + g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") + g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.inIndent(func() { + g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r) + }) + g.emit("}\n\n") + g.emit("// WriteTo implements io.WriterTo.WriteTo.\n") g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName()) g.inIndent(func() { diff --git a/tools/go_marshal/gomarshal/generator_interfaces_struct.go b/tools/go_marshal/gomarshal/generator_interfaces_struct.go index 1aac63d32..5f9a06d0f 100644 --- a/tools/go_marshal/gomarshal/generator_interfaces_struct.go +++ b/tools/go_marshal/gomarshal/generator_interfaces_struct.go @@ -322,6 +322,7 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) { fallback() } }) + g.emit("}\n\n") g.emit("// CopyOutN implements marshal.Marshallable.CopyOutN.\n") g.recordUsedImport("marshal") @@ -364,15 +365,15 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) { }) g.emit("}\n\n") - g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") + g.emit("// CopyInN implements marshal.Marshallable.CopyInN.\n") g.recordUsedImport("marshal") g.recordUsedImport("hostarch") - g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.emit("func (%s *%s) CopyInN(cc marshal.CopyContext, addr hostarch.Addr, limit int) (int, error) {\n", g.r, g.typeName()) g.inIndent(func() { fallback := func() { g.emit("// Type %s doesn't have a packed layout in memory, fall back to UnmarshalBytes.\n", g.typeName()) g.emit("buf := cc.CopyScratchBuffer(%s.SizeBytes()) // escapes: okay.\n", g.r) - g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n") + g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n") g.emit("// Unmarshal unconditionally. If we had a short copy-in, this results in a\n") g.emit("// partially unmarshalled struct.\n") g.emit("%s.UnmarshalBytes(buf) // escapes: fallback.\n", g.r) @@ -390,7 +391,7 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) { // Fast deserialization. g.emitCastToByteSlice(g.r, "buf", fmt.Sprintf("%s.SizeBytes()", g.r)) - g.emit("length, err := cc.CopyInBytes(addr, buf) // escapes: okay.\n") + g.emit("length, err := cc.CopyInBytes(addr, buf[:limit]) // escapes: okay.\n") g.emitKeepAlive(g.r) g.emit("return length, err\n") } else { @@ -399,6 +400,15 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) { }) g.emit("}\n\n") + g.emit("// CopyIn implements marshal.Marshallable.CopyIn.\n") + g.recordUsedImport("marshal") + g.recordUsedImport("hostarch") + g.emit("func (%s *%s) CopyIn(cc marshal.CopyContext, addr hostarch.Addr) (int, error) {\n", g.r, g.typeName()) + g.inIndent(func() { + g.emit("return %s.CopyInN(cc, addr, %s.SizeBytes())\n", g.r, g.r) + }) + g.emit("}\n\n") + g.emit("// WriteTo implements io.WriterTo.WriteTo.\n") g.recordUsedImport("io") g.emit("func (%s *%s) WriteTo(writer io.Writer) (int64, error) {\n", g.r, g.typeName()) diff --git a/tools/go_marshal/test/marshal_test.go b/tools/go_marshal/test/marshal_test.go index cb2aa16d1..571b3184a 100644 --- a/tools/go_marshal/test/marshal_test.go +++ b/tools/go_marshal/test/marshal_test.go @@ -222,6 +222,32 @@ func limitedCopyOut(t *testing.T, src marshal.Marshallable, limit int) { compareMemory(t, expectedMem, actualMem, n) } +// copyInN marshals src to task memory, requesting the marshalling to be +// limited to limit bytes. +func copyInN(t *testing.T, src, dst marshal.Marshallable, limit int) { + var cc mockCopyContext + cc.populate(src) + cc.setLimit(limit) + + n, err := dst.CopyInN(&cc, hostarch.Addr(0), limit) + if err != nil { + t.Errorf("CopyInN returned unexpected error: %v", err) + } + if n != limit { + t.Errorf("CopyInN copied unexpected number of bytes, expected %d, got %d", limit, n) + } + + expectedMem := unsafeMemory(src) + defer runtime.KeepAlive(src) + actualMem := unsafeMemory(dst) + defer runtime.KeepAlive(dst) + + t.Logf("Expected: %v + %v\n", expectedMem[:n], expectedMem[n:]) + t.Logf("Actual : %v + %v\n", actualMem[:n], actualMem[n:]) + + compareMemory(t, expectedMem, actualMem, n) +} + // copyOutN marshals src to task memory, requesting the marshalling to be // limited to limit bytes. func copyOutN(t *testing.T, src marshal.Marshallable, limit int) { @@ -296,6 +322,14 @@ func TestLimitedMarshalling(t *testing.T) { copyOutN(t, expected, expected.SizeBytes()/2) }) + + // Explicitly request partial copy-in. + t.Run(fmt.Sprintf("PartialCopyInN_%v", ty), func(t *testing.T) { + expected := reflect.New(ty).Interface().(marshal.Marshallable) + analysis.RandomizeValue(expected) + + copyInN(t, expected, expected, expected.SizeBytes()/2) + }) } }