src: clean up MaybeStackBuffer

- Add IsInvalidated() method - Add capacity() method for finding out the actual capacity, not the current size, of the buffer - Make IsAllocated() work for invalidated buffers - Allow multiple calls to AllocateSufficientStorage() and Invalidate() - Assert buffer is malloc'd in Release() - Assert buffer has not been invalidated in AllocateSufficientStorage() - Add more descriptive comments describing the purpose of the methods - Add cctest for MaybeStackBuffer PR-URL: #11464 Reviewed-By: Steven R Loomis <srloomis@us.ibm.com> Reviewed-By: Anna Henningsen <anna@addaleax.net> Reviewed-By: Ben Noordhuis <info@bnoordhuis.nl> Reviewed-By: James M Snell <jasnell@gmail.com>
nodejs · TimothyGu · Nov 28, 2017 · Oct 23, 2016 · Jun 1, 2017 · Nov 28, 2017
commit 2efd843cdd2a9735e643b624c087b3849762909a
diff --git a/src/util.h b/src/util.h
@@ -10,6 +10,7 @@
 #include <stddef.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 // OSX 10.9 defaults to libc++ which provides a C++11 <type_traits> header.
 #if defined(__APPLE__) && __ENVIRONMENT_MAC_OS_X_VERSION_MIN_REQUIRED__ < 1090
@@ -310,45 +311,76 @@ class MaybeStackBuffer {
     return length_;
   }
 
-  // Call to make sure enough space for `storage` entries is available.
-  // There can only be 1 call to AllocateSufficientStorage or Invalidate
-  // per instance.
+  // Current maximum capacity of the buffer with which SetLength() can be used
+  // without first calling AllocateSufficientStorage().
+  size_t capacity() const {
+    return IsAllocated() ? capacity_ :
+                           IsInvalidated() ? 0 : kStackStorageSize;
+  }
+
+  // Make sure enough space for `storage` entries is available.
+  // This method can be called multiple times throughout the lifetime of the
+  // buffer, but once this has been called Invalidate() cannot be used.
+  // Content of the buffer in the range [0, length()) is preserved.
   void AllocateSufficientStorage(size_t storage) {
-    if (storage <= kStackStorageSize) {
-      buf_ = buf_st_;
-    } else {
-      buf_ = Malloc<T>(storage);
+    CHECK(!IsInvalidated());
+    if (storage > capacity()) {
+      bool was_allocated = IsAllocated();
+      T* allocated_ptr = was_allocated ? buf_ : nullptr;
+      buf_ = Realloc(allocated_ptr, storage);
+      capacity_ = storage;
+      if (!was_allocated && length_ > 0)
+        memcpy(buf_, buf_st_, length_ * sizeof(buf_[0]));
     }
 
-    // Remember how much was allocated to check against that in SetLength().
     length_ = storage;
   }
 
   void SetLength(size_t length) {
-    // length_ stores how much memory was allocated.
-    CHECK_LE(length, length_);
+    // capacity() returns how much memory is actually available.
+    CHECK_LE(length, capacity());
     length_ = length;
   }
 
   void SetLengthAndZeroTerminate(size_t length) {
-    // length_ stores how much memory was allocated.
-    CHECK_LE(length + 1, length_);
+    // capacity() returns how much memory is actually available.
+    CHECK_LE(length + 1, capacity());
     SetLength(length);
 
     // T() is 0 for integer types, nullptr for pointers, etc.
     buf_[length] = T();
   }
 
   // Make derefencing this object return nullptr.
-  // Calling this is mutually exclusive with calling
-  // AllocateSufficientStorage.
+  // This method can be called multiple times throughout the lifetime of the
+  // buffer, but once this has been called AllocateSufficientStorage() cannot
+  // be used.
   void Invalidate() {
-    CHECK_EQ(buf_, buf_st_);
+    CHECK(!IsAllocated());
     length_ = 0;
     buf_ = nullptr;
   }
 
-  MaybeStackBuffer() : length_(0), buf_(buf_st_) {
+  // If the buffer is stored in the heap rather than on the stack.
+  bool IsAllocated() const {
+    return !IsInvalidated() && buf_ != buf_st_;
+  }
+
+  // If Invalidate() has been called.
+  bool IsInvalidated() const {
+    return buf_ == nullptr;
+  }
+
+  // Release ownership of the malloc'd buffer.
+  // Note: This does not free the buffer.
+  void Release() {
+    CHECK(IsAllocated());
+    buf_ = buf_st_;
+    length_ = 0;
+    capacity_ = 0;
+  }
+
+  MaybeStackBuffer() : length_(0), capacity_(0), buf_(buf_st_) {
     // Default to a zero-length, null-terminated buffer.
     buf_[0] = T();
   }
@@ -358,12 +390,14 @@ class MaybeStackBuffer {
   }
 
   ~MaybeStackBuffer() {
-    if (buf_ != buf_st_)
+    if (IsAllocated())
       free(buf_);
   }
 
  private:
   size_t length_;
+  // capacity of the malloc'ed buf_
+  size_t capacity_;
   T* buf_;
   T buf_st_[kStackStorageSize];
 };

diff --git a/test/cctest/test_util.cc b/test/cctest/test_util.cc
@@ -121,3 +121,127 @@ TEST(UtilTest, UncheckedCalloc) {
   EXPECT_NE(nullptr, UncheckedCalloc(0));
   EXPECT_NE(nullptr, UncheckedCalloc(1));
 }
+
+template <typename T>
+static void MaybeStackBufferBasic() {
+  using node::MaybeStackBuffer;
+
+  MaybeStackBuffer<T> buf;
+  size_t old_length;
+  size_t old_capacity;
+
+  /* Default constructor */
+  EXPECT_EQ(0U, buf.length());
+  EXPECT_FALSE(buf.IsAllocated());
+  EXPECT_GT(buf.capacity(), buf.length());
+
+  /* SetLength() expansion */
+  buf.SetLength(buf.capacity());
+  EXPECT_EQ(buf.capacity(), buf.length());
+  EXPECT_FALSE(buf.IsAllocated());
+
+  /* Means of accessing raw buffer */
+  EXPECT_EQ(buf.out(), *buf);
+  EXPECT_EQ(&buf[0], *buf);
+
+  /* Basic I/O */
+  for (size_t i = 0; i < buf.length(); i++)
+    buf[i] = static_cast<T>(i);
+  for (size_t i = 0; i < buf.length(); i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+
+  /* SetLengthAndZeroTerminate() */
+  buf.SetLengthAndZeroTerminate(buf.capacity() - 1);
+  EXPECT_EQ(buf.capacity() - 1, buf.length());
+  for (size_t i = 0; i < buf.length(); i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+  buf.SetLength(buf.capacity());
+  EXPECT_EQ(0, buf[buf.length() - 1]);
+
+  /* Initial Realloc */
+  old_length = buf.length() - 1;
+  old_capacity = buf.capacity();
+  buf.AllocateSufficientStorage(buf.capacity() * 2);
+  EXPECT_EQ(buf.capacity(), buf.length());
+  EXPECT_TRUE(buf.IsAllocated());
+  for (size_t i = 0; i < old_length; i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+  EXPECT_EQ(0, buf[old_length]);
+
+  /* SetLength() reduction and expansion */
+  for (size_t i = 0; i < buf.length(); i++)
+    buf[i] = static_cast<T>(i);
+  buf.SetLength(10);
+  for (size_t i = 0; i < buf.length(); i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+  buf.SetLength(buf.capacity());
+  for (size_t i = 0; i < buf.length(); i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+
+  /* Subsequent Realloc */
+  old_length = buf.length();
+  old_capacity = buf.capacity();
+  buf.AllocateSufficientStorage(old_capacity * 1.5);
+  EXPECT_EQ(buf.capacity(), buf.length());
+  EXPECT_EQ(static_cast<size_t>(old_capacity * 1.5), buf.length());
+  EXPECT_TRUE(buf.IsAllocated());
+  for (size_t i = 0; i < old_length; i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+
+  /* Basic I/O on Realloc'd buffer */
+  for (size_t i = 0; i < buf.length(); i++)
+    buf[i] = static_cast<T>(i);
+  for (size_t i = 0; i < buf.length(); i++)
+    EXPECT_EQ(static_cast<T>(i), buf[i]);
+
+  /* Release() */
+  T* rawbuf = buf.out();
+  buf.Release();
+  EXPECT_EQ(0U, buf.length());
+  EXPECT_FALSE(buf.IsAllocated());
+  EXPECT_GT(buf.capacity(), buf.length());
+  free(rawbuf);
+}
+
+TEST(UtilTest, MaybeStackBuffer) {
+  using node::MaybeStackBuffer;
+
+  MaybeStackBufferBasic<uint8_t>();
+  MaybeStackBufferBasic<uint16_t>();
+
+  // Constructor with size parameter
+  {
+    MaybeStackBuffer<unsigned char> buf(100);
+    EXPECT_EQ(100U, buf.length());
+    EXPECT_FALSE(buf.IsAllocated());
+    EXPECT_GT(buf.capacity(), buf.length());
+    buf.SetLength(buf.capacity());
+    EXPECT_EQ(buf.capacity(), buf.length());
+    EXPECT_FALSE(buf.IsAllocated());
+    for (size_t i = 0; i < buf.length(); i++)
+      buf[i] = static_cast<unsigned char>(i);
+    for (size_t i = 0; i < buf.length(); i++)
+      EXPECT_EQ(static_cast<unsigned char>(i), buf[i]);
+
+    MaybeStackBuffer<unsigned char> bigbuf(10000);
+    EXPECT_EQ(10000U, bigbuf.length());
+    EXPECT_TRUE(bigbuf.IsAllocated());
+    EXPECT_EQ(bigbuf.length(), bigbuf.capacity());
+    for (size_t i = 0; i < bigbuf.length(); i++)
+      bigbuf[i] = static_cast<unsigned char>(i);
+    for (size_t i = 0; i < bigbuf.length(); i++)
+      EXPECT_EQ(static_cast<unsigned char>(i), bigbuf[i]);
+  }
+
+  // Invalidated buffer
+  {
+    MaybeStackBuffer<char> buf;
+    buf.Invalidate();
+    EXPECT_TRUE(buf.IsInvalidated());
+    EXPECT_FALSE(buf.IsAllocated());
+    EXPECT_EQ(0U, buf.length());
+    EXPECT_EQ(0U, buf.capacity());
+    buf.Invalidate();
+    EXPECT_TRUE(buf.IsInvalidated());
+  }
+}