diff --git a/src/util.h b/src/util.h index 20b2ffcdba454b..cf530a4da9146c 100644 --- a/src/util.h +++ b/src/util.h @@ -10,6 +10,7 @@ #include #include #include +#include #include // std::remove_reference @@ -304,29 +305,40 @@ 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(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. @@ -334,24 +346,35 @@ class MaybeStackBuffer { } // 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; } - bool IsAllocated() { - return 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), buf_(buf_st_) { + MaybeStackBuffer() : length_(0), capacity_(0), buf_(buf_st_) { // Default to a zero-length, null-terminated buffer. buf_[0] = T(); } @@ -361,12 +384,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/util.cc b/test/cctest/util.cc index f99a46cdf5353b..a6ece3c6f4d377 100644 --- a/test/cctest/util.cc +++ b/test/cctest/util.cc @@ -132,3 +132,127 @@ TEST(UtilTest, UncheckedCalloc) { TEST_AND_FREE(UncheckedCalloc(0)); TEST_AND_FREE(UncheckedCalloc(1)); } + +template +static void MaybeStackBufferBasic() { + using node::MaybeStackBuffer; + + MaybeStackBuffer 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(i); + for (size_t i = 0; i < buf.length(); i++) + EXPECT_EQ(static_cast(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(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(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(i); + buf.SetLength(10); + for (size_t i = 0; i < buf.length(); i++) + EXPECT_EQ(static_cast(i), buf[i]); + buf.SetLength(buf.capacity()); + for (size_t i = 0; i < buf.length(); i++) + EXPECT_EQ(static_cast(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(old_capacity * 1.5), buf.length()); + EXPECT_TRUE(buf.IsAllocated()); + for (size_t i = 0; i < old_length; i++) + EXPECT_EQ(static_cast(i), buf[i]); + + /* Basic I/O on Realloc'd buffer */ + for (size_t i = 0; i < buf.length(); i++) + buf[i] = static_cast(i); + for (size_t i = 0; i < buf.length(); i++) + EXPECT_EQ(static_cast(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(); + MaybeStackBufferBasic(); + + // Constructor with size parameter + { + MaybeStackBuffer 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(i); + for (size_t i = 0; i < buf.length(); i++) + EXPECT_EQ(static_cast(i), buf[i]); + + MaybeStackBuffer 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(i); + for (size_t i = 0; i < bigbuf.length(); i++) + EXPECT_EQ(static_cast(i), bigbuf[i]); + } + + // Invalidated buffer + { + MaybeStackBuffer 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()); + } +}