Add commonly-desired kj::VectorOutputStream.

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

@@ -62,5 +62,51 @@ KJ_TEST("stringify AutoCloseFd") {
   KJ_EXPECT(kj::str(in) == kj::str(fds[0]), in, fds[0]);
 }
 
+KJ_TEST("VectorOutputStream") {
+  VectorOutputStream output(16);
+  auto buf = output.getWriteBuffer();
+  KJ_ASSERT(buf.size() == 16);
+
+  for (auto i: kj::indices(buf)) {
+    buf[i] = 'a' + i;
+  }
+
+  output.write(buf.begin(), 4);
+  KJ_ASSERT(output.getArray().begin() == buf.begin());
+  KJ_ASSERT(output.getArray().size() == 4);
+
+  auto buf2 = output.getWriteBuffer();
+  KJ_ASSERT(buf2.end() == buf.end());
+  KJ_ASSERT(buf2.size() == 12);
+
+  output.write(buf2.begin(), buf2.size());
+  KJ_ASSERT(output.getArray().begin() == buf.begin());
+  KJ_ASSERT(output.getArray().size() == 16);
+
+  auto buf3 = output.getWriteBuffer();
+  KJ_ASSERT(buf3.size() == 16);
+  KJ_ASSERT(output.getArray().begin() != buf.begin());
+  KJ_ASSERT(output.getArray().end() == buf3.begin());
+  KJ_ASSERT(kj::str(output.getArray().asChars()) == "abcdefghijklmnop");
+
+  byte junk[24];
+  for (auto i: kj::indices(junk)) {
+    junk[i] = 'A' + i;
+  }
+
+  output.write(junk, 4);
+  KJ_ASSERT(output.getArray().begin() != buf.begin());
+  KJ_ASSERT(output.getArray().end() == buf3.begin() + 4);
+  KJ_ASSERT(kj::str(output.getArray().asChars()) == "abcdefghijklmnopABCD");
+
+  output.write(junk + 4, 20);
+  KJ_ASSERT(output.getArray().begin() != buf.begin());
+  KJ_ASSERT(output.getArray().end() != buf3.begin() + 24);
+  KJ_ASSERT(kj::str(output.getArray().asChars()) == "abcdefghijklmnopABCDEFGHIJKLMNOPQRSTUVWX");
+
+  KJ_ASSERT(output.getWriteBuffer().size() == 24);
+  KJ_ASSERT(output.getWriteBuffer().begin() == output.getArray().begin() + 40);
+}
+
 }  // namespace
 }  // namespace kj

c++/src/kj/io.c++

@@ -222,6 +222,7 @@ ArrayPtr<byte> ArrayOutputStream::getWriteBuffer() {
 void ArrayOutputStream::write(const void* src, size_t size) {
   if (src == fillPos) {
     // Oh goody, the caller wrote directly into our buffer.
+    KJ_REQUIRE(size <= array.end() - fillPos);
     fillPos += size;
   } else {
     KJ_REQUIRE(size <= (size_t)(array.end() - fillPos),
@@ -231,6 +232,45 @@ void ArrayOutputStream::write(const void* src, size_t size) {
   }
 }
 
+// -------------------------------------------------------------------
+
+VectorOutputStream::VectorOutputStream(size_t initialCapacity)
+    : vector(heapArray<byte>(initialCapacity)), fillPos(vector.begin()) {}
+VectorOutputStream::~VectorOutputStream() noexcept(false) {}
+
+ArrayPtr<byte> VectorOutputStream::getWriteBuffer() {
+  // Grow if needed.
+  if (fillPos == vector.end()) {
+    grow(vector.size() + 1);
+  }
+
+  return arrayPtr(fillPos, vector.end());
+}
+
+void VectorOutputStream::write(const void* src, size_t size) {
+  if (src == fillPos) {
+    // Oh goody, the caller wrote directly into our buffer.
+    KJ_REQUIRE(size <= vector.end() - fillPos);
+    fillPos += size;
+  } else {
+    if (vector.end() - fillPos < size) {
+      grow(fillPos - vector.begin() + size);
+    }
+
+    memcpy(fillPos, src, size);
+    fillPos += size;
+  }
+}
+
+void VectorOutputStream::grow(size_t minSize) {
+  size_t newSize = vector.size() * 2;
+  while (newSize < minSize) newSize *= 2;
+  auto newVector = heapArray<byte>(newSize);
+  memcpy(newVector.begin(), vector.begin(), fillPos - vector.begin());
+  fillPos = fillPos - vector.begin() + newVector.begin();
+  vector = kj::mv(newVector);
+}
+
 // =======================================================================================
 
 AutoCloseFd::~AutoCloseFd() noexcept(false) {

c++/src/kj/io.h

@@ -219,6 +219,28 @@ private:
   byte* fillPos;
 };
 
+class VectorOutputStream: public BufferedOutputStream {
+public:
+  explicit VectorOutputStream(size_t initialCapacity = 4096);
+  KJ_DISALLOW_COPY(VectorOutputStream);
+  ~VectorOutputStream() noexcept(false);
+
+  ArrayPtr<byte> getArray() {
+    // Get the portion of the array which has been filled in.
+    return arrayPtr(vector.begin(), fillPos);
+  }
+
+  // implements BufferedInputStream ----------------------------------
+  ArrayPtr<byte> getWriteBuffer() override;
+  void write(const void* buffer, size_t size) override;
+
+private:
+  Array<byte> vector;
+  byte* fillPos;
+
+  void grow(size_t minSize);
+};
+
 // =======================================================================================
 // File descriptor I/O