Merge pull request #834 from capnproto/attach-and-clone

Introduce kj::attachVal(), kj::attachRef(), and capnp::clone() utility functions

c++/src/capnp/arena.h

@@ -288,6 +288,10 @@ public:
     return &localCapTable;
   }
 
+  kj::Own<_::CapTableBuilder> releaseLocalCapTable() {
+    return kj::heap<LocalCapTable>(kj::mv(localCapTable));
+  }
+
   SegmentBuilder* getSegment(SegmentId id);
   // Get the segment with the given id.  Crashes or throws an exception if no such segment exists.
 
@@ -321,7 +325,7 @@ private:
   MessageBuilder* message;
   ReadLimiter dummyLimiter;
 
-  class LocalCapTable: public CapTableBuilder {
+  class LocalCapTable final: public CapTableBuilder {
 #if !CAPNP_LITE
   public:
     kj::Maybe<kj::Own<ClientHook>> extractCap(uint index) override;

c++/src/capnp/capability-test.c++

@@ -1080,6 +1080,27 @@ KJ_TEST("Promise<RemotePromise<T>> automatically reduces to RemotePromise<T> wit
   EXPECT_EQ(1, chainedCallCount);
 }
 
+KJ_TEST("clone() with caps") {
+  int dummy = 0;
+  MallocMessageBuilder builder(2048);
+  auto root = builder.getRoot<AnyPointer>().initAs<List<test::TestInterface>>(3);
+  root.set(0, kj::heap<TestInterfaceImpl>(dummy));
+  root.set(1, kj::heap<TestInterfaceImpl>(dummy));
+  root.set(2, kj::heap<TestInterfaceImpl>(dummy));
+
+  auto copyPtr = clone(root.asReader());
+  auto& copy = *copyPtr;
+
+  KJ_ASSERT(copy.size() == 3);
+  KJ_EXPECT(ClientHook::from(copy[0]).get() == ClientHook::from(root[0]).get());
+  KJ_EXPECT(ClientHook::from(copy[1]).get() == ClientHook::from(root[1]).get());
+  KJ_EXPECT(ClientHook::from(copy[2]).get() == ClientHook::from(root[2]).get());
+
+  KJ_EXPECT(ClientHook::from(copy[0]).get() != ClientHook::from(root[1]).get());
+  KJ_EXPECT(ClientHook::from(copy[1]).get() != ClientHook::from(root[2]).get());
+  KJ_EXPECT(ClientHook::from(copy[2]).get() != ClientHook::from(root[0]).get());
+}
+
 }  // namespace
 }  // namespace _
 }  // namespace capnp

c++/src/capnp/message-test.c++

@@ -168,6 +168,14 @@ TEST(Message, ReadWriteDataStruct) {
   checkTestMessageAllZero(defaultValue<TestAllTypes>());
 }
 
+KJ_TEST("clone()") {
+  MallocMessageBuilder builder(2048);
+  initTestMessage(builder.getRoot<TestAllTypes>());
+
+  auto copy = clone(builder.getRoot<TestAllTypes>().asReader());
+  checkTestMessage(*copy);
+}
+
 // TODO(test):  More tests.
 
 }  // namespace

c++/src/capnp/message.c++

@@ -178,6 +178,10 @@ bool MessageBuilder::isCanonical() {
                                   .isCanonical(&readHead);
 }
 
+kj::Own<_::CapTableBuilder> MessageBuilder::releaseBuiltinCapTable() {
+  return arena()->releaseLocalCapTable();
+}
+
 // =======================================================================================
 
 SegmentArrayMessageReader::SegmentArrayMessageReader(

c++/src/capnp/message.h

@@ -39,6 +39,7 @@ namespace capnp {
 namespace _ {  // private
   class ReaderArena;
   class BuilderArena;
+  struct CloneImpl;
 }
 
 class StructSchema;
@@ -254,6 +255,14 @@ private:
   _::BuilderArena* arena() { return reinterpret_cast<_::BuilderArena*>(arenaSpace); }
   _::SegmentBuilder* getRootSegment();
   AnyPointer::Builder getRootInternal();
+
+  kj::Own<_::CapTableBuilder> releaseBuiltinCapTable();
+  // Hack for clone() to extract the cap table.
+
+  friend struct _::CloneImpl;
+  // We can't declare clone() as a friend directly because old versions of GCC incorrectly demand
+  // that the first declaration (even if it is a friend declaration) specify the default type args,
+  // whereas correct compilers do not permit default type args to be specified on a friend decl.
 };
 
 template <typename RootType>
@@ -315,6 +324,10 @@ static typename Type::Reader defaultValue();
 //
 // TODO(cleanup):  Find a better home for this function?
 
+template <typename Reader, typename = FromReader<Reader>>
+kj::Own<kj::Decay<Reader>> clone(Reader&& reader);
+// Make a deep copy of the given Reader on the heap, producing an owned pointer.
+
 // =======================================================================================
 
 class SegmentArrayMessageReader: public MessageReader {
@@ -506,6 +519,33 @@ static typename Type::Reader defaultValue() {
   return typename Type::Reader(_::StructReader());
 }
 
+namespace _ {
+  struct CloneImpl {
+    static inline kj::Own<_::CapTableBuilder> releaseBuiltinCapTable(MessageBuilder& message) {
+      return message.releaseBuiltinCapTable();
+    }
+  };
+};
+
+template <typename Reader, typename>
+kj::Own<kj::Decay<Reader>> clone(Reader&& reader) {
+  auto size = reader.totalSize();
+  auto buffer = kj::heapArray<capnp::word>(size.wordCount + 1);
+  memset(buffer.asBytes().begin(), 0, buffer.asBytes().size());
+  if (size.capCount == 0) {
+    copyToUnchecked(reader, buffer);
+    auto result = readMessageUnchecked<FromReader<Reader>>(buffer.begin());
+    return kj::attachVal(result, kj::mv(buffer));
+  } else {
+    FlatMessageBuilder builder(buffer);
+    builder.setRoot(kj::fwd<Reader>(reader));
+    builder.requireFilled();
+    auto capTable = _::CloneImpl::releaseBuiltinCapTable(builder);
+    AnyPointer::Reader raw(_::PointerReader::getRootUnchecked(buffer.begin()).imbue(capTable));
+    return kj::attachVal(raw.getAs<FromReader<Reader>>(), kj::mv(buffer), kj::mv(capTable));
+  }
+}
+
 template <typename T>
 kj::Array<word> canonicalize(T&& reader) {
     return _::PointerHelpers<FromReader<T>>::getInternalReader(reader).canonicalize();

c++/src/kj/memory-test.c++

@@ -136,6 +136,68 @@ TEST(Memory, AttachNested) {
   KJ_EXPECT(destroyed3 == 3, destroyed3);
 }
 
+KJ_TEST("attachRef") {
+  uint counter = 0;
+  uint destroyed1 = 0;
+  uint destroyed2 = 0;
+  uint destroyed3 = 0;
+
+  auto obj1 = kj::heap<DestructionOrderRecorder>(counter, destroyed1);
+  auto obj2 = kj::heap<DestructionOrderRecorder>(counter, destroyed2);
+  auto obj3 = kj::heap<DestructionOrderRecorder>(counter, destroyed3);
+
+  int i = 123;
+
+  Own<int> combined = attachRef(i, kj::mv(obj1), kj::mv(obj2), kj::mv(obj3));
+
+  KJ_EXPECT(combined.get() == &i);
+
+  KJ_EXPECT(obj1.get() == nullptr);
+  KJ_EXPECT(obj2.get() == nullptr);
+  KJ_EXPECT(obj3.get() == nullptr);
+  KJ_EXPECT(destroyed1 == 0);
+  KJ_EXPECT(destroyed2 == 0);
+  KJ_EXPECT(destroyed3 == 0);
+
+  combined = nullptr;
+
+  KJ_EXPECT(destroyed1 == 1, destroyed1);
+  KJ_EXPECT(destroyed2 == 2, destroyed2);
+  KJ_EXPECT(destroyed3 == 3, destroyed3);
+}
+
+KJ_TEST("attachVal") {
+  uint counter = 0;
+  uint destroyed1 = 0;
+  uint destroyed2 = 0;
+  uint destroyed3 = 0;
+
+  auto obj1 = kj::heap<DestructionOrderRecorder>(counter, destroyed1);
+  auto obj2 = kj::heap<DestructionOrderRecorder>(counter, destroyed2);
+  auto obj3 = kj::heap<DestructionOrderRecorder>(counter, destroyed3);
+
+  int i = 123;
+
+  Own<int> combined = attachVal(i, kj::mv(obj1), kj::mv(obj2), kj::mv(obj3));
+
+  int* ptr = combined.get();
+  KJ_EXPECT(ptr != &i);
+  KJ_EXPECT(*ptr == i);
+
+  KJ_EXPECT(obj1.get() == nullptr);
+  KJ_EXPECT(obj2.get() == nullptr);
+  KJ_EXPECT(obj3.get() == nullptr);
+  KJ_EXPECT(destroyed1 == 0);
+  KJ_EXPECT(destroyed2 == 0);
+  KJ_EXPECT(destroyed3 == 0);
+
+  combined = nullptr;
+
+  KJ_EXPECT(destroyed1 == 1, destroyed1);
+  KJ_EXPECT(destroyed2 == 2, destroyed2);
+  KJ_EXPECT(destroyed3 == 3, destroyed3);
+}
+
 struct StaticType {
   int i;
 };

c++/src/kj/memory.h

@@ -426,6 +426,21 @@ Own<Decay<T>> heap(T&& orig) {
   return Own<T2>(new T2(kj::fwd<T>(orig)), _::HeapDisposer<T2>::instance);
 }
 
+template <typename T, typename... Attachments>
+Own<Decay<T>> attachVal(T&& value, Attachments&&... attachments);
+// Returns an Own<T> that takes ownership of `value` and `attachments`, and points to `value`.
+//
+// This is equivalent to heap(value).attach(attachments), but only does one allocation rather than
+// two.
+
+template <typename T, typename... Attachments>
+Own<T> attachRef(T& value, Attachments&&... attachments);
+// Like attach() but `value` is not moved; the resulting Own<T> points to its existing location.
+// This is preferred if `value` is already owned by one of `attachments`.
+//
+// This is equivalent to Own<T>(&value, kj::NullDisposer::instance).attach(attachments), but
+// is easier to write and allocates slightly less memory.
+
 // =======================================================================================
 // SpaceFor<T> -- assists in manual allocation
 
@@ -519,4 +534,17 @@ Own<T> Own<T>::attach(Attachments&&... attachments) {
   return Own<T>(ptrCopy, *bundle);
 }
 
+template <typename T, typename... Attachments>
+Own<T> attachRef(T& value, Attachments&&... attachments) {
+  auto bundle = new _::DisposableOwnedBundle<Attachments...>(kj::fwd<Attachments>(attachments)...);
+  return Own<T>(&value, *bundle);
+}
+
+template <typename T, typename... Attachments>
+Own<Decay<T>> attachVal(T&& value, Attachments&&... attachments) {
+  auto bundle = new _::DisposableOwnedBundle<T, Attachments...>(
+      kj::fwd<T>(value), kj::fwd<Attachments>(attachments)...);
+  return Own<Decay<T>>(&bundle->first, *bundle);
+}
+
 }  // namespace kj