Merge pull request #610 from capnproto/list-ops

Add totalSize() to list readers, like struct readers.

c++/src/capnp/any.h

@@ -408,6 +408,10 @@ struct List<AnyPointer, Kind::OTHER> {
     inline Iterator begin() const { return Iterator(this, 0); }
     inline Iterator end() const { return Iterator(this, size()); }
 
+    inline MessageSize totalSize() const {
+      return reader.totalSize().asPublic();
+    }
+
   private:
     _::ListReader reader;
     template <typename U, Kind K>
@@ -585,6 +589,10 @@ public:
   inline Iterator begin() const { return Iterator(this, 0); }
   inline Iterator end() const { return Iterator(this, size()); }
 
+  inline MessageSize totalSize() const {
+    return reader.totalSize().asPublic();
+  }
+
 private:
   _::ListReader reader;
   template <typename U, Kind K>
@@ -650,6 +658,10 @@ public:
     return !(*this == right);
   }
 
+  inline MessageSize totalSize() const {
+    return _reader.totalSize().asPublic();
+  }
+
   template <typename T> ReaderFor<T> as() {
     // T must be List<U>.
     return ReaderFor<T>(_reader);

c++/src/capnp/capability.h

@@ -661,6 +661,10 @@ struct List<T, Kind::INTERFACE> {
     inline Iterator begin() const { return Iterator(this, 0); }
     inline Iterator end() const { return Iterator(this, size()); }
 
+    inline MessageSize totalSize() const {
+      return reader.totalSize().asPublic();
+    }
+
   private:
     _::ListReader reader;
     template <typename U, Kind K>

c++/src/capnp/common.h

@@ -322,9 +322,13 @@ struct PointerHelpers {};
 }  // namespace _ (private)
 
 struct MessageSize {
-  // Size of a message.  Every struct type has a method `.totalSize()` that returns this.
+  // Size of a message. Every struct and list type has a method `.totalSize()` that returns this.
   uint64_t wordCount;
   uint capCount;
+
+  inline constexpr MessageSize operator+(const MessageSize& other) const {
+    return { wordCount + other.wordCount, capCount + other.capCount };
+  }
 };
 
 // =======================================================================================

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

@@ -1929,6 +1929,32 @@ TEST(Encoding, DefaultListBuilder) {
   List<Text>::Builder(nullptr);
 }
 
+TEST(Encoding, ListSize) {
+  MallocMessageBuilder builder;
+  auto root = builder.initRoot<TestListDefaults>();
+  initTestMessage(root);
+
+  auto lists = root.asReader().getLists();
+
+  auto listSizes =
+      lists.getList0().totalSize() +
+      lists.getList1().totalSize() +
+      lists.getList8().totalSize() +
+      lists.getList16().totalSize() +
+      lists.getList32().totalSize() +
+      lists.getList64().totalSize() +
+      lists.getListP().totalSize() +
+      lists.getInt32ListList().totalSize() +
+      lists.getTextListList().totalSize() +
+      lists.getStructListList().totalSize();
+
+  auto structSize = lists.totalSize();
+
+  auto shallowSize = unbound(capnp::_::structSize<test::TestLists>().total() / WORDS);
+
+  EXPECT_EQ(structSize.wordCount - shallowSize, listSizes.wordCount);
+}
+
 }  // namespace
 }  // namespace _ (private)
 }  // namespace capnp

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

@@ -843,7 +843,7 @@ struct WireHelpers {
 
             // We count the actual size rather than the claimed word count because that's what
             // we'll end up with if we make a copy.
-            result.addWords(wordCount + POINTER_SIZE_IN_WORDS);
+            result.addWords(actualSize + POINTER_SIZE_IN_WORDS);
 
             WordCount dataSize = elementTag->structRef.dataSize.get();
             WirePointerCount pointerCount = elementTag->structRef.ptrCount.get();
@@ -3113,6 +3113,64 @@ StructReader ListReader::getStructElement(ElementCount index) const {
       nestingLimit - 1);
 }
 
+MessageSizeCounts ListReader::totalSize() const {
+  // TODO(cleanup): This is kind of a lot of logic duplicated from WireHelpers::totalSize(), but
+  //   it's unclear how to share it effectively.
+
+  MessageSizeCounts result = { ZERO * WORDS, 0 };
+
+  switch (elementSize) {
+    case ElementSize::VOID:
+      // Nothing.
+      break;
+    case ElementSize::BIT:
+    case ElementSize::BYTE:
+    case ElementSize::TWO_BYTES:
+    case ElementSize::FOUR_BYTES:
+    case ElementSize::EIGHT_BYTES:
+      result.addWords(WireHelpers::roundBitsUpToWords(
+          upgradeBound<uint64_t>(elementCount) * dataBitsPerElement(elementSize)));
+      break;
+    case ElementSize::POINTER: {
+      auto count = elementCount * (POINTERS / ELEMENTS);
+      result.addWords(count * WORDS_PER_POINTER);
+
+      for (auto i: kj::zeroTo(count)) {
+        result += WireHelpers::totalSize(segment, reinterpret_cast<const WirePointer*>(ptr) + i,
+                                         nestingLimit);
+      }
+      break;
+    }
+    case ElementSize::INLINE_COMPOSITE: {
+      // Don't forget to count the tag word.
+      auto wordSize = upgradeBound<uint64_t>(elementCount) * step / BITS_PER_WORD;
+      result.addWords(wordSize + POINTER_SIZE_IN_WORDS);
+
+      if (structPointerCount > ZERO * POINTERS) {
+        const word* pos = reinterpret_cast<const word*>(ptr);
+        for (auto i KJ_UNUSED: kj::zeroTo(elementCount)) {
+          pos += structDataSize / BITS_PER_WORD;
+
+          for (auto j KJ_UNUSED: kj::zeroTo(structPointerCount)) {
+            result += WireHelpers::totalSize(segment, reinterpret_cast<const WirePointer*>(pos),
+                                             nestingLimit);
+            pos += POINTER_SIZE_IN_WORDS;
+          }
+        }
+      }
+      break;
+    }
+  }
+
+  if (segment != nullptr) {
+    // This traversal should not count against the read limit, because it's highly likely that
+    // the caller is going to traverse the object again, e.g. to copy it.
+    segment->unread(result.wordCount);
+  }
+
+  return result;
+}
+
 CapTableReader* ListReader::getCapTable() {
   return capTable;
 }

c++/src/capnp/layout.h

@@ -625,9 +625,7 @@ public:
   MessageSizeCounts totalSize() const;
   // Return the total size of the struct and everything to which it points.  Does not count far
   // pointer overhead.  This is useful for deciding how much space is needed to copy the struct
-  // into a flat array.  However, the caller is advised NOT to treat this value as secure.  Instead,
-  // use the result as a hint for allocating the first segment, do the copy, and then throw an
-  // exception if it overruns.
+  // into a flat array.
 
   CapTableReader* getCapTable();
   // Gets the capability context in which this object is operating.
@@ -793,6 +791,9 @@ public:
 
   StructReader getStructElement(ElementCount index) const;
 
+  MessageSizeCounts totalSize() const;
+  // Like StructReader::totalSize(). Note that for struct lists, the size includes the list tag.
+
   CapTableReader* getCapTable();
   // Gets the capability context in which this object is operating.
 

c++/src/capnp/list.h

@@ -124,6 +124,10 @@ struct List<T, Kind::PRIMITIVE> {
     inline Iterator begin() const { return Iterator(this, 0); }
     inline Iterator end() const { return Iterator(this, size()); }
 
+    inline MessageSize totalSize() const {
+      return reader.totalSize().asPublic();
+    }
+
   private:
     _::ListReader reader;
     template <typename U, Kind K>
@@ -220,6 +224,10 @@ struct List<T, Kind::STRUCT> {
     inline Iterator begin() const { return Iterator(this, 0); }
     inline Iterator end() const { return Iterator(this, size()); }
 
+    inline MessageSize totalSize() const {
+      return reader.totalSize().asPublic();
+    }
+
   private:
     _::ListReader reader;
     template <typename U, Kind K>
@@ -343,6 +351,10 @@ struct List<List<T>, Kind::LIST> {
     inline Iterator begin() const { return Iterator(this, 0); }
     inline Iterator end() const { return Iterator(this, size()); }
 
+    inline MessageSize totalSize() const {
+      return reader.totalSize().asPublic();
+    }
+
   private:
     _::ListReader reader;
     template <typename U, Kind K>
@@ -452,6 +464,10 @@ struct List<T, Kind::BLOB> {
     inline Iterator begin() const { return Iterator(this, 0); }
     inline Iterator end() const { return Iterator(this, size()); }
 
+    inline MessageSize totalSize() const {
+      return reader.totalSize().asPublic();
+    }
+
   private:
     _::ListReader reader;
     template <typename U, Kind K>