Bug 1264642 - Part 2. Add BufferList::Extract and Pickle::ExtractBuff…

…ers. r=billm

These methods allow us to move some buffers out of a pickle with minimum
copying. It's useful when the IPC deserialized type uses BufferList to
store data and we want to take the buffers from IPC directly.

Borrowing is not suitable to use for IPC to hand out data because we
often want to store the data somewhere for processing after IPC has
released the underlying buffers.

MozReview-Commit-ID: F1K2ZMkACqq

ipc/chromium/src/base/pickle.cc

@@ -421,6 +421,21 @@ bool Pickle::FlattenBytes(PickleIterator* iter, const char** data, uint32_t leng
   return iter->iter_.AdvanceAcrossSegments(buffers_, AlignInt(length) - length);
 }
 
+bool Pickle::ExtractBuffers(PickleIterator* iter, size_t length, BufferList* buffers) const
+{
+  if (AlignInt(length) < length) {
+    return false;
+  }
+
+  bool success;
+  *buffers = const_cast<BufferList*>(&buffers_)->Extract(iter->iter_, length, &success);
+  if (!success) {
+    return false;
+  }
+
+  return iter->iter_.AdvanceAcrossSegments(buffers_, AlignInt(length) - length);
+}
+
 bool Pickle::ReadBytesInto(PickleIterator* iter, void* data, uint32_t length) const {
   if (AlignInt(length) < length) {
     return false;

ipc/chromium/src/base/pickle.h

@@ -113,6 +113,7 @@ class Pickle {
   MOZ_MUST_USE bool ReadBytesInto(PickleIterator* iter, void* data, uint32_t length) const;
   MOZ_MUST_USE bool FlattenBytes(PickleIterator* iter, const char** data, uint32_t length,
                                  uint32_t alignment = sizeof(memberAlignmentType));
+  MOZ_MUST_USE bool ExtractBuffers(PickleIterator* iter, size_t length, BufferList* buffers) const;
 
   // Safer version of ReadInt() checks for the result not being negative.
   // Use it for reading the object sizes.

mfbt/BufferList.h

@@ -10,6 +10,7 @@
 #include <algorithm>
 #include "mozilla/AllocPolicy.h"
 #include "mozilla/Move.h"
+#include "mozilla/ScopeExit.h"
 #include "mozilla/Types.h"
 #include "mozilla/TypeTraits.h"
 #include "mozilla/Vector.h"
@@ -257,6 +258,16 @@ class BufferList : private AllocPolicy
   BufferList<BorrowingAllocPolicy> Borrow(IterImpl& aIter, size_t aSize, bool* aSuccess,
                                           BorrowingAllocPolicy aAP = BorrowingAllocPolicy());
 
+  // Return a new BufferList that adopts the byte range starting at Iter so that
+  // range [aIter, aIter + aSize) is transplanted to the returned BufferList.
+  // Contents of the buffer before aIter + aSize is left undefined.
+  // Extract can fail, in which case *aSuccess will be false upon return. The
+  // moved buffers are erased from the original BufferList. In case of extract
+  // fails, the original BufferList is intact.  All other iterators except aIter
+  // are invalidated.
+  // This method requires aIter and aSize to be 8-byte aligned.
+  BufferList Extract(IterImpl& aIter, size_t aSize, bool* aSuccess);
+
 private:
   explicit BufferList(AllocPolicy aAP)
    : AllocPolicy(aAP),
@@ -430,6 +441,74 @@ BufferList<AllocPolicy>::Borrow(IterImpl& aIter, size_t aSize, bool* aSuccess,
   return result;
 }
 
+template<typename AllocPolicy>
+BufferList<AllocPolicy>
+BufferList<AllocPolicy>::Extract(IterImpl& aIter, size_t aSize, bool* aSuccess)
+{
+  MOZ_RELEASE_ASSERT(aSize);
+  MOZ_RELEASE_ASSERT(mOwning);
+  MOZ_ASSERT(aSize % kSegmentAlignment == 0);
+  MOZ_ASSERT(intptr_t(aIter.mData) % kSegmentAlignment == 0);
+
+  IterImpl iter = aIter;
+  size_t size = aSize;
+  size_t toCopy = std::min(size, aIter.RemainingInSegment());
+  MOZ_ASSERT(toCopy % kSegmentAlignment == 0);
+
+  BufferList result(0, toCopy, mStandardCapacity);
+  BufferList error(0, 0, mStandardCapacity);
+
+  // Copy the head
+  if (!result.WriteBytes(aIter.mData, toCopy)) {
+    *aSuccess = false;
+    return error;
+  }
+  iter.Advance(*this, toCopy);
+  size -= toCopy;
+
+  // Move segments to result
+  auto resultGuard = MakeScopeExit([&] {
+    *aSuccess = false;
+    result.mSegments.erase(result.mSegments.begin()+1, result.mSegments.end());
+  });
+
+  size_t movedSize = 0;
+  uintptr_t toRemoveStart = iter.mSegment;
+  uintptr_t toRemoveEnd = iter.mSegment;
+  while (!iter.Done() &&
+         !iter.HasRoomFor(size)) {
+    if (!result.mSegments.append(Segment(mSegments[iter.mSegment].mData,
+                                         mSegments[iter.mSegment].mSize,
+                                         mSegments[iter.mSegment].mCapacity))) {
+      return error;
+    }
+    movedSize += iter.RemainingInSegment();
+    size -= iter.RemainingInSegment();
+    toRemoveEnd++;
+    iter.Advance(*this, iter.RemainingInSegment());
+  }
+
+  if (size)  {
+    if (!iter.HasRoomFor(size) ||
+        !result.WriteBytes(iter.Data(), size)) {
+      return error;
+    }
+    iter.Advance(*this, size);
+  }
+
+  mSegments.erase(mSegments.begin() + toRemoveStart, mSegments.begin() + toRemoveEnd);
+  mSize -= movedSize;
+  aIter.mSegment = iter.mSegment - (toRemoveEnd - toRemoveStart);
+  aIter.mData = iter.mData;
+  aIter.mDataEnd = iter.mDataEnd;
+  MOZ_ASSERT(aIter.mDataEnd == mSegments[aIter.mSegment].End());
+  result.mSize = aSize;
+
+  resultGuard.release();
+  *aSuccess = true;
+  return result;
+}
+
 } // namespace mozilla
 
 #endif /* mozilla_BufferList_h */

mfbt/tests/TestBufferList.cpp

@@ -239,5 +239,27 @@ int main(void)
   MOZ_RELEASE_ASSERT(iter2.AdvanceAcrossSegments(bl2, kBorrowSize - 5));
   MOZ_RELEASE_ASSERT(iter1.Data() == iter2.Data());
 
+  // Extracting.
+
+  const size_t kExtractStart = 8;
+  const size_t kExtractSize = 24;
+  const size_t kExtractOverSize = 1000;
+
+  iter = bl.Iter();
+  iter.Advance(bl, kExtractStart);
+  bl2 = bl.Extract(iter, kExtractSize, &success);
+  MOZ_RELEASE_ASSERT(success);
+  MOZ_RELEASE_ASSERT(bl2.Size() == kExtractSize);
+
+  BufferList bl3 = bl.Extract(iter, kExtractOverSize, &success);
+  MOZ_RELEASE_ASSERT(!success);
+
+  MOZ_RELEASE_ASSERT(iter.AdvanceAcrossSegments(bl, kSmallWrite * 3 - kExtractSize - kExtractStart));
+  MOZ_RELEASE_ASSERT(iter.Done());
+
+  iter = bl2.Iter();
+  MOZ_RELEASE_ASSERT(iter.AdvanceAcrossSegments(bl2, kExtractSize));
+  MOZ_RELEASE_ASSERT(iter.Done());
+
   return 0;
 }