Add --segment-size option for testing purposes.

c++/src/capnp/compiler/capnp.c++

@@ -177,6 +177,11 @@ public:
            .addOption({'p', "packed"}, KJ_BIND_METHOD(*this, codePacked),
                       "Pack the output message with standard Cap'n Proto packing, which "
                       "deflates zero-valued bytes.")
+           .addOptionWithArg({"segment-size"}, KJ_BIND_METHOD(*this, setSegmentSize), "<n>",
+                             "Sets the preferred segment size on the MallocMessageBuilder to <n> "
+                             "words and turns off heuristic growth.  This flag is mainly useful "
+                             "for testing.  Without it, each message will be written as a single "
+                             "segment.")
            .expectArg("<schema-file>", KJ_BIND_METHOD(*this, addSource))
            .expectArg("<type>", KJ_BIND_METHOD(*this, setRootType))
            .callAfterParsing(KJ_BIND_METHOD(*this, encode));
@@ -434,6 +439,15 @@ public:
     pretty = false;
     return true;
   }
+  kj::MainBuilder::Validity setSegmentSize(kj::StringPtr size) {
+    if (flat) return "cannot be used with --flat";
+    char* end;
+    segmentSize = strtol(size.cStr(), &end, 0);
+    if (size.size() == 0 || *end != '\0') {
+      return "not an integer";
+    }
+    return true;
+  }
 
   kj::MainBuilder::Validity setRootType(kj::StringPtr type) {
     KJ_ASSERT(sourceFiles.size() == 1);
@@ -582,10 +596,6 @@ public:
     auto tokens = lexedTokens.asReader().getTokens();
     CapnpParser::ParserInput parserInput(tokens.begin(), tokens.end());
 
-    // Allocate some scratch space.
-    kj::Array<word> scratch = kj::heapArray<word>(8192);
-    memset(scratch.begin(), 0, scratch.size() * sizeof(word));
-
     // Set up stuff for the ValueTranslator.
     ValueResolverGlue resolver(compiler->getLoader(), errorReporter);
     auto type = arena.getOrphanage().newOrphan<schema::Type>();
@@ -597,11 +607,22 @@ public:
 
     while (parserInput.getPosition() != tokens.end()) {
       KJ_IF_MAYBE(expression, parser.getParsers().parenthesizedValueExpression(parserInput)) {
-        MallocMessageBuilder item(scratch);
+        MallocMessageBuilder item(
+            segmentSize == 0 ? SUGGESTED_FIRST_SEGMENT_WORDS : segmentSize,
+            segmentSize == 0 ? SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy::FIXED_SIZE);
         ValueTranslator translator(resolver, errorReporter, item.getOrphanage());
 
         KJ_IF_MAYBE(value, translator.compileValue(expression->getReader(), type.getReader())) {
-          writeEncoded(value->getReader().as<DynamicStruct>(), output);
+          if (segmentSize == 0) {
+            writeFlat(value->getReader().as<DynamicStruct>(), output);
+          } else {
+            item.adoptRoot(value->releaseAs<DynamicStruct>());
+            if (packed) {
+              writePackedMessage(output, item);
+            } else {
+              writeMessage(output, item);
+            }
+          }
         } else {
           // Errors were reported, so we'll exit with a failure status later.
         }
@@ -620,7 +641,7 @@ public:
   }
 
 private:
-  void writeEncoded(DynamicStruct::Reader value, kj::BufferedOutputStream& output) {
+  void writeFlat(DynamicStruct::Reader value, kj::BufferedOutputStream& output) {
     // Always copy the message to a flat array so that the output is predictable (one segment,
     // in canonical order).
     size_t size = value.totalSizeInWords() + 1;
@@ -744,8 +765,9 @@ private:
   bool flat = false;
   bool packed = false;
   bool pretty = true;
+  uint segmentSize = 0;
   StructSchema rootType;
-  // For the "decode" command.
+  // For the "decode" and "encode" commands.
 
   struct SourceFile {
     uint64_t id;

c++/src/capnp/dynamic.h

@@ -685,6 +685,7 @@ private:
   friend class Orphanage;
   friend class Orphan<DynamicObject>;
   friend class Orphan<DynamicValue>;
+  friend class MessageBuilder;
 };
 
 template <>
@@ -936,6 +937,10 @@ template <>
 void MessageBuilder::setRoot<const DynamicStruct::Reader&>(const DynamicStruct::Reader& value);
 template <>
 void MessageBuilder::setRoot<DynamicStruct::Reader&>(DynamicStruct::Reader& value);
+template <>
+inline void MessageBuilder::adoptRoot<DynamicStruct>(Orphan<DynamicStruct>&& orphan) {
+  adoptRootInternal(kj::mv(orphan.builder));
+}
 
 namespace _ {  // private
 

c++/src/capnp/layout.c++

@@ -2056,6 +2056,10 @@ StructBuilder StructBuilder::getRoot(
       reinterpret_cast<WirePointer*>(location), segment, size, nullptr);
 }
 
+void StructBuilder::adoptRoot(SegmentBuilder* segment, word* location, OrphanBuilder orphan) {
+  WireHelpers::adopt(segment, reinterpret_cast<WirePointer*>(location), kj::mv(orphan));
+}
+
 StructBuilder StructBuilder::initStructField(
     WirePointerCount ptrIndex, StructSize size) {
   return WireHelpers::initStructPointer(pointers + ptrIndex, segment, size);

c++/src/capnp/layout.h

@@ -277,6 +277,7 @@ public:
   static StructBuilder initRoot(SegmentBuilder* segment, word* location, StructSize size);
   static void setRoot(SegmentBuilder* segment, word* location, StructReader value);
   static StructBuilder getRoot(SegmentBuilder* segment, word* location, StructSize size);
+  static void adoptRoot(SegmentBuilder* segment, word* location, OrphanBuilder orphan);
 
   inline BitCount getDataSectionSize() const { return dataSize; }
   inline WirePointerCount getPointerSectionSize() const { return pointerCount; }

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

@@ -107,6 +107,12 @@ _::StructBuilder MessageBuilder::getRoot(_::StructSize size) {
       rootSegment, rootSegment->getPtrUnchecked(0 * WORDS), size);
 }
 
+void MessageBuilder::adoptRootInternal(_::OrphanBuilder orphan) {
+  _::SegmentBuilder* rootSegment = getRootSegment();
+  _::StructBuilder::adoptRoot(
+      rootSegment, rootSegment->getPtrUnchecked(0 * WORDS), kj::mv(orphan));
+}
+
 kj::ArrayPtr<const kj::ArrayPtr<const word>> MessageBuilder::getSegmentsForOutput() {
   if (allocatedArena) {
     return arena()->getSegmentsForOutput();

c++/src/capnp/message.h

@@ -39,6 +39,8 @@ namespace _ {  // private
 
 class StructSchema;
 class Orphanage;
+template <typename T>
+class Orphan;
 
 // =======================================================================================
 
@@ -178,6 +180,10 @@ public:
   // RootType in this case must be DynamicStruct, and you must #include <capnp/dynamic.h> to
   // use this.
 
+  template <typename T>
+  void adoptRoot(Orphan<T>&& orphan);
+  // Like setRoot() but adopts the orphan without copying.
+
   kj::ArrayPtr<const kj::ArrayPtr<const word>> getSegmentsForOutput();
 
   Orphanage getOrphanage();
@@ -201,6 +207,7 @@ private:
   _::StructBuilder initRoot(_::StructSize size);
   void setRootInternal(_::StructReader reader);
   _::StructBuilder getRoot(_::StructSize size);
+  void adoptRootInternal(_::OrphanBuilder orphan);
 };
 
 template <typename RootType>
@@ -380,6 +387,12 @@ inline typename RootType::Builder MessageBuilder::getRoot() {
   return typename RootType::Builder(getRoot(_::structSize<RootType>()));
 }
 
+template <typename T>
+void MessageBuilder::adoptRoot(Orphan<T>&& orphan) {
+  static_assert(kind<T>() == Kind::STRUCT, "Root type must be a Cap'n Proto struct type.");
+  adoptRootInternal(kj::mv(orphan.builder));
+}
+
 template <typename RootType>
 typename RootType::Reader readMessageUnchecked(const word* data) {
   return typename RootType::Reader(_::StructReader::readRootUnchecked(data));

c++/src/capnp/orphan.h

@@ -70,6 +70,7 @@ private:
   template <typename U>
   friend class Orphan;
   friend class Orphanage;
+  friend class MessageBuilder;
 };
 
 class Orphanage: private kj::DisallowConstCopy {