When an under-sized struct is found while traversing a builder, upgrade it to…

When an under-sized struct is found while traversing a builder, upgrade it to the expected size.  This is needed to allow easy copying from readers to builders.  Next, this will need to be implemented for lists as well.

c++/src/capnproto/encoding-test.c++

@@ -616,6 +616,255 @@ TEST(Encoding, BitListUpgrade) {
   checkList(reader.getObjectField<List<bool>>(), {true, false, true, true});
 }
 
+TEST(Encoding, UpgradeStructInBuilder) {
+  MallocMessageBuilder builder;
+  auto root = builder.initRoot<test::TestObject>();
+
+  {
+    auto oldVersion = root.initObjectField<test::TestOldVersion>();
+    oldVersion.setOld1(123);
+    oldVersion.setOld2("foo");
+    auto sub = oldVersion.initOld3();
+    sub.setOld1(456);
+    sub.setOld2("bar");
+  }
+
+  size_t size = builder.getSegmentsForOutput()[0].size();
+  size_t size2;
+
+  {
+    auto newVersion = root.getObjectField<test::TestNewVersion>();
+
+    // Size should have increased due to re-allocating the struct.
+    size_t size1 = builder.getSegmentsForOutput()[0].size();
+    EXPECT_GT(size1, size);
+
+    auto sub = newVersion.getOld3();
+
+    // Size should have increased due to re-allocating the sub-struct.
+    size2 = builder.getSegmentsForOutput()[0].size();
+    EXPECT_GT(size2, size1);
+
+    // Check contents.
+    EXPECT_EQ(123, newVersion.getOld1());
+    EXPECT_EQ("foo", newVersion.getOld2());
+    EXPECT_EQ(987, newVersion.getNew1());
+    EXPECT_EQ("baz", newVersion.getNew2());
+
+    EXPECT_EQ(456, sub.getOld1());
+    EXPECT_EQ("bar", sub.getOld2());
+    EXPECT_EQ(987, sub.getNew1());
+    EXPECT_EQ("baz", sub.getNew2());
+
+    newVersion.setOld1(234);
+    newVersion.setOld2("qux");
+    newVersion.setNew1(321);
+    newVersion.setNew2("quux");
+
+    sub.setOld1(567);
+    sub.setOld2("corge");
+    sub.setNew1(654);
+    sub.setNew2("grault");
+  }
+
+  // We set four small text fields and implicitly initialized two to defaults, so the size should
+  // have raised by six words.
+  size_t size3 = builder.getSegmentsForOutput()[0].size();
+  EXPECT_EQ(size2 + 6, size3);
+
+  {
+    // Go back to old version.  It should have the values set on the new version.
+    auto oldVersion = root.getObjectField<test::TestOldVersion>();
+    EXPECT_EQ(234, oldVersion.getOld1());
+    EXPECT_EQ("qux", oldVersion.getOld2());
+
+    auto sub = oldVersion.getOld3();
+    EXPECT_EQ(567, sub.getOld1());
+    EXPECT_EQ("corge", sub.getOld2());
+
+    // Overwrite the old fields.  The new fields should remain intact.
+    oldVersion.setOld1(345);
+    oldVersion.setOld2("garply");
+    sub.setOld1(678);
+    sub.setOld2("waldo");
+  }
+
+  // We set two small text fields, so the size should have raised by two words.
+  size_t size4 = builder.getSegmentsForOutput()[0].size();
+  EXPECT_EQ(size3 + 2, size4);
+
+  {
+    // Back to the new version again.
+    auto newVersion = root.getObjectField<test::TestNewVersion>();
+    EXPECT_EQ(345, newVersion.getOld1());
+    EXPECT_EQ("garply", newVersion.getOld2());
+    EXPECT_EQ(321, newVersion.getNew1());
+    EXPECT_EQ("quux", newVersion.getNew2());
+
+    auto sub = newVersion.getOld3();
+    EXPECT_EQ(678, sub.getOld1());
+    EXPECT_EQ("waldo", sub.getOld2());
+    EXPECT_EQ(654, sub.getNew1());
+    EXPECT_EQ("grault", sub.getNew2());
+  }
+
+  // Size should not have changed because we didn't write anything and the structs were already
+  // the right size.
+  EXPECT_EQ(size4, builder.getSegmentsForOutput()[0].size());
+}
+
+TEST(Encoding, UpgradeStructInBuilderMultiSegment) {
+  // Exactly like the previous test, except that we force multiple segments.  Since we force a
+  // separate segment for every object, every pointer is a far pointer, and far pointers are easily
+  // transferred, so this is acutally not such a complicated case.
+
+  MallocMessageBuilder builder(0, AllocationStrategy::FIXED_SIZE);
+  auto root = builder.initRoot<test::TestObject>();
+
+  // Start with a 1-word first segment and the root object in the second segment.
+  size_t size = builder.getSegmentsForOutput().size();
+  EXPECT_EQ(2u, size);
+
+  {
+    auto oldVersion = root.initObjectField<test::TestOldVersion>();
+    oldVersion.setOld1(123);
+    oldVersion.setOld2("foo");
+    auto sub = oldVersion.initOld3();
+    sub.setOld1(456);
+    sub.setOld2("bar");
+  }
+
+  // Allocated two structs and two strings.
+  size_t size2 = builder.getSegmentsForOutput().size();
+  EXPECT_EQ(size + 4, size2);
+
+  size_t size4;
+
+  {
+    auto newVersion = root.getObjectField<test::TestNewVersion>();
+
+    // Allocated a new struct.
+    size_t size3 = builder.getSegmentsForOutput().size();
+    EXPECT_EQ(size2 + 1, size3);
+
+    auto sub = newVersion.getOld3();
+
+    // Allocated another new struct for its string field.
+    size4 = builder.getSegmentsForOutput().size();
+    EXPECT_EQ(size3 + 1, size4);
+
+    // Check contents.
+    EXPECT_EQ(123, newVersion.getOld1());
+    EXPECT_EQ("foo", newVersion.getOld2());
+    EXPECT_EQ(987, newVersion.getNew1());
+    EXPECT_EQ("baz", newVersion.getNew2());
+
+    EXPECT_EQ(456, sub.getOld1());
+    EXPECT_EQ("bar", sub.getOld2());
+    EXPECT_EQ(987, sub.getNew1());
+    EXPECT_EQ("baz", sub.getNew2());
+
+    newVersion.setOld1(234);
+    newVersion.setOld2("qux");
+    newVersion.setNew1(321);
+    newVersion.setNew2("quux");
+
+    sub.setOld1(567);
+    sub.setOld2("corge");
+    sub.setNew1(654);
+    sub.setNew2("grault");
+  }
+
+  // Set four strings and implicitly initialized two.
+  size_t size5 = builder.getSegmentsForOutput().size();
+  EXPECT_EQ(size4 + 6, size5);
+
+  {
+    // Go back to old version.  It should have the values set on the new version.
+    auto oldVersion = root.getObjectField<test::TestOldVersion>();
+    EXPECT_EQ(234, oldVersion.getOld1());
+    EXPECT_EQ("qux", oldVersion.getOld2());
+
+    auto sub = oldVersion.getOld3();
+    EXPECT_EQ(567, sub.getOld1());
+    EXPECT_EQ("corge", sub.getOld2());
+
+    // Overwrite the old fields.  The new fields should remain intact.
+    oldVersion.setOld1(345);
+    oldVersion.setOld2("garply");
+    sub.setOld1(678);
+    sub.setOld2("waldo");
+  }
+
+  // Set two new strings.
+  size_t size6 = builder.getSegmentsForOutput().size();
+  EXPECT_EQ(size5 + 2, size6);
+
+  {
+    // Back to the new version again.
+    auto newVersion = root.getObjectField<test::TestNewVersion>();
+    EXPECT_EQ(345, newVersion.getOld1());
+    EXPECT_EQ("garply", newVersion.getOld2());
+    EXPECT_EQ(321, newVersion.getNew1());
+    EXPECT_EQ("quux", newVersion.getNew2());
+
+    auto sub = newVersion.getOld3();
+    EXPECT_EQ(678, sub.getOld1());
+    EXPECT_EQ("waldo", sub.getOld2());
+    EXPECT_EQ(654, sub.getNew1());
+    EXPECT_EQ("grault", sub.getNew2());
+  }
+
+  // Size should not have changed because we didn't write anything and the structs were already
+  // the right size.
+  EXPECT_EQ(size6, builder.getSegmentsForOutput().size());
+}
+
+TEST(Encoding, UpgradeStructInBuilderFarPointers) {
+  // Force allocation of a Far pointer.
+
+  MallocMessageBuilder builder(7, AllocationStrategy::FIXED_SIZE);
+  auto root = builder.initRoot<test::TestObject>();
+
+  root.initObjectField<test::TestOldVersion>().setOld2("foo");
+
+  // We should have allocated all but one word of the first segment.
+  EXPECT_EQ(1u, builder.getSegmentsForOutput().size());
+  EXPECT_EQ(6u, builder.getSegmentsForOutput()[0].size());
+
+  // Now if we upgrade...
+  EXPECT_EQ("foo", root.getObjectField<test::TestNewVersion>().getOld2());
+
+  // We should have allocated the new struct in a new segment, but allocated the far pointer
+  // landing pad back in the first segment.
+  ASSERT_EQ(2u, builder.getSegmentsForOutput().size());
+  EXPECT_EQ(7u, builder.getSegmentsForOutput()[0].size());
+  EXPECT_EQ(6u, builder.getSegmentsForOutput()[1].size());
+}
+
+TEST(Encoding, UpgradeStructInBuilderDoubleFarPointers) {
+  // Force allocation of a double-Far pointer.
+
+  MallocMessageBuilder builder(6, AllocationStrategy::FIXED_SIZE);
+  auto root = builder.initRoot<test::TestObject>();
+
+  root.initObjectField<test::TestOldVersion>().setOld2("foo");
+
+  // We should have allocated all of the first segment.
+  EXPECT_EQ(1u, builder.getSegmentsForOutput().size());
+  EXPECT_EQ(6u, builder.getSegmentsForOutput()[0].size());
+
+  // Now if we upgrade...
+  EXPECT_EQ("foo", root.getObjectField<test::TestNewVersion>().getOld2());
+
+  // We should have allocated the new struct in a new segment, and also allocated the far pointer
+  // landing pad in yet another segment.
+  ASSERT_EQ(3u, builder.getSegmentsForOutput().size());
+  EXPECT_EQ(6u, builder.getSegmentsForOutput()[0].size());
+  EXPECT_EQ(6u, builder.getSegmentsForOutput()[1].size());
+  EXPECT_EQ(2u, builder.getSegmentsForOutput()[2].size());
+}
+
 // =======================================================================================
 // Tests of generated code, not really of the encoding.
 // TODO(cleanup):  Move to a different test?

c++/src/capnproto/layout.c++

@@ -89,6 +89,9 @@ struct WireReference {
   CAPNPROTO_ALWAYS_INLINE(void setKindAndTarget(Kind kind, word* target)) {
     offsetAndKind.set(((target - reinterpret_cast<word*>(this) - 1) << 2) | kind);
   }
+  CAPNPROTO_ALWAYS_INLINE(void setKindWithZeroOffset(Kind kind)) {
+    offsetAndKind.set(kind);
+  }
 
   CAPNPROTO_ALWAYS_INLINE(ElementCount inlineCompositeListElementCount() const) {
     return (offsetAndKind.get() >> 2) * ELEMENTS;
@@ -314,9 +317,10 @@ struct WireHelpers {
     }
   }
 
-  // Not always-inline because it's recursive.
   static word* copyMessage(
       SegmentBuilder*& segment, WireReference*& dst, const WireReference* src) {
+    // Not always-inline because it's recursive.
+
     switch (src->kind()) {
       case WireReference::STRUCT: {
         if (src->isNull()) {
@@ -407,6 +411,54 @@ struct WireHelpers {
     return nullptr;
   }
 
+  static void transferPointer(SegmentBuilder* dstSegment, WireReference* dst,
+                              SegmentBuilder* srcSegment, WireReference* src) {
+    // Make *dst point to the same object as *src.  Both must reside in the same message, but can
+    // be in different segments.  Not always-inline because this is rarely used.
+
+    if (src->isNull()) {
+      memset(dst, 0, sizeof(WireReference));
+    } else if (src->kind() == WireReference::FAR) {
+      // Far pointers are position-independent, so we can just copy.
+      memcpy(dst, src, sizeof(WireReference));
+    } else if (dstSegment == srcSegment) {
+      // Same segment, so create a direct reference.
+      dst->setKindAndTarget(src->kind(), src->target());
+
+      // We can just copy the upper 32 bits.  (Use memcpy() to comply with aliasing rules.)
+      memcpy(&dst->upper32Bits, &src->upper32Bits, sizeof(src->upper32Bits));
+    } else {
+      // Need to create a far pointer.  Try to allocate it in the same segment as the source, so
+      // that it doesn't need to be a double-far.
+
+      WireReference* landingPad =
+          reinterpret_cast<WireReference*>(srcSegment->allocate(1 * WORDS));
+      if (landingPad == nullptr) {
+        // Darn, need a double-far.
+        SegmentBuilder* farSegment = srcSegment->getArena()->getSegmentWithAvailable(2 * WORDS);
+        landingPad = reinterpret_cast<WireReference*>(farSegment->allocate(2 * WORDS));
+        DCHECK(landingPad != nullptr,
+            "getSegmentWithAvailable() returned segment without space available.");
+
+        landingPad[0].setFar(false, srcSegment->getOffsetTo(src->target()));
+        landingPad[0].farRef.segmentId.set(srcSegment->getSegmentId());
+
+        landingPad[1].setKindWithZeroOffset(src->kind());
+        memcpy(&landingPad[1].upper32Bits, &src->upper32Bits, sizeof(src->upper32Bits));
+
+        dst->setFar(true, farSegment->getOffsetTo(reinterpret_cast<word*>(landingPad)));
+        dst->farRef.set(farSegment->getSegmentId());
+      } else {
+        // Simple landing pad is just a pointer.
+        landingPad->setKindAndTarget(src->kind(), src->target());
+        memcpy(&landingPad->upper32Bits, &src->upper32Bits, sizeof(src->upper32Bits));
+
+        dst->setFar(false, srcSegment->getOffsetTo(reinterpret_cast<word*>(landingPad)));
+        dst->farRef.set(srcSegment->getSegmentId());
+      }
+    }
+  }
+
   // -----------------------------------------------------------------
 
   static CAPNPROTO_ALWAYS_INLINE(StructBuilder initStructReference(
@@ -427,6 +479,7 @@ struct WireHelpers {
     word* ptr;
 
     if (ref->isNull()) {
+    useDefault:
       if (defaultValue == nullptr ||
           reinterpret_cast<const WireReference*>(defaultValue)->isNull()) {
         ptr = allocate(ref, segment, size.total(), WireReference::STRUCT);
@@ -434,21 +487,52 @@ struct WireHelpers {
       } else {
         ptr = copyMessage(segment, ref, reinterpret_cast<const WireReference*>(defaultValue));
       }
+      return StructBuilder(segment, ptr, reinterpret_cast<WireReference*>(ptr + size.data),
+                           size.data * BITS_PER_WORD, size.pointers, 0 * BITS);
     } else {
-      ptr = followFars(ref, segment);
+      WireReference* oldRef = ref;
+      SegmentBuilder* oldSegment = segment;
+      word* oldPtr = followFars(oldRef, oldSegment);
 
-      DPRECOND(ref->kind() == WireReference::STRUCT,
-          "Called getStruct{Field,Element}() but existing reference is not a struct.");
-      DPRECOND(
-          ref->structRef.dataSize.get() == size.data,
-          "Trying to update struct with incorrect data size.");
-      DPRECOND(
-          ref->structRef.refCount.get() == size.pointers,
-          "Trying to update struct with incorrect reference count.");
-    }
+      VALIDATE_INPUT(oldRef->kind() == WireReference::STRUCT,
+          "Message contains non-struct reference where struct reference was expected.") {
+        goto useDefault;
+      }
 
-    return StructBuilder(segment, ptr, reinterpret_cast<WireReference*>(ptr + size.data),
-                         size.data * BITS_PER_WORD, size.pointers, 0 * BITS);
+      WordCount oldDataSize = oldRef->structRef.dataSize.get();
+      WireReferenceCount oldPointerCount = oldRef->structRef.refCount.get();
+      WireReference* oldPointerSection =
+          reinterpret_cast<WireReference*>(oldPtr + oldDataSize);
+
+      if (oldDataSize < size.data || oldPointerCount < size.pointers) {
+        // The space allocated for this struct is too small.  Unlike with readers, we can't just
+        // run with it and do bounds checks at access time, because how would we handle writes?
+        // Instead, we have to copy the struct to a new space now.
+
+        WordCount newDataSize = std::max<WordCount>(oldDataSize, size.data);
+        WireReferenceCount newPointerCount =
+            std::max<WireReferenceCount>(oldPointerCount, size.pointers);
+        WordCount totalSize = newDataSize + newPointerCount * WORDS_PER_REFERENCE;
+
+        ptr = allocate(ref, segment, totalSize, WireReference::STRUCT);
+        ref->structRef.set(newDataSize, newPointerCount);
+
+        // Copy data section.
+        memcpy(ptr, oldPtr, oldDataSize * BYTES_PER_WORD / BYTES);
+
+        // Copy pointer section.
+        WireReference* newPointerSection = reinterpret_cast<WireReference*>(ptr + newDataSize);
+        for (uint i = 0; i < oldPointerCount / REFERENCES; i++) {
+          transferPointer(segment, newPointerSection + i, oldSegment, oldPointerSection + i);
+        }
+
+        return StructBuilder(segment, ptr, newPointerSection, newDataSize * BITS_PER_WORD,
+                             newPointerCount, 0 * BITS);
+      } else {
+        return StructBuilder(oldSegment, oldPtr, oldPointerSection, oldDataSize * BITS_PER_WORD,
+                             oldPointerCount, 0 * BITS);
+      }
+    }
   }
 
   static CAPNPROTO_ALWAYS_INLINE(ListBuilder initListReference(

c++/src/capnproto/macros.h

@@ -52,8 +52,13 @@ namespace internal {
 // expect the condition to be true/false enough of the time that it's worth hard-coding branch
 // prediction.
 
+#ifdef NDEBUG
 #define CAPNPROTO_ALWAYS_INLINE(prototype) inline prototype __attribute__((always_inline))
 // Force a function to always be inlined.  Apply only to the prototype, not to the definition.
+#else
+#define CAPNPROTO_ALWAYS_INLINE(prototype) inline prototype
+// Don't force inline in debug mode.
+#endif
 
 #define CAPNPROTO_NORETURN __attribute__((noreturn));
 

c++/src/capnproto/test.capnp

@@ -354,3 +354,19 @@ struct TestLateUnion {
     grault @6 :Float32;
   }
 }
+
+struct TestOldVersion {
+  # A subset of TestNewVersion.
+  old1 @0 :Int32;
+  old2 @1 :Text;
+  old3 @2 :TestOldVersion;
+}
+
+struct TestNewVersion {
+  # A superset of TestOldVersion.
+  old1 @0 :Int32;
+  old2 @1 :Text;
+  old3 @2 :TestNewVersion;
+  new1 @3 :Int64 = 987;
+  new2 @4 :Text = "baz";
+}