diff --git a/toolbelt/payload_buffer.cc b/toolbelt/payload_buffer.cc index b852461..96df051 100644 --- a/toolbelt/payload_buffer.cc +++ b/toolbelt/payload_buffer.cc @@ -1,5 +1,6 @@ #include "toolbelt/payload_buffer.h" #include +#include namespace toolbelt { @@ -44,7 +45,8 @@ char *PayloadBuffer::SetString(PayloadBuffer **self, const char *s, size_t len, return reinterpret_cast(str); } -void PayloadBuffer::ClearString(PayloadBuffer **self, BufferOffset header_offset) { +void PayloadBuffer::ClearString(PayloadBuffer **self, + BufferOffset header_offset) { BufferOffset *hdr = (*self)->ToAddress(header_offset); if (*hdr != 0) { (*self)->Free((*self)->ToAddress(*hdr)); @@ -86,8 +88,10 @@ const char *PayloadBuffer::StringData(const StringHeader *addr) const { return reinterpret_cast(p + 1); } -absl::Span PayloadBuffer::AllocateString(PayloadBuffer **self, size_t len, BufferOffset header_offset, bool clear) { - // Get address of the string header +absl::Span PayloadBuffer::AllocateString(PayloadBuffer **self, size_t len, + BufferOffset header_offset, + bool clear) { + // Get address of the string header BufferOffset *hdr = (*self)->ToAddress(header_offset); void *str = nullptr; @@ -142,7 +146,7 @@ void PayloadBuffer::InitFreeList() { size_t header_size = sizeof(PayloadBuffer); if (magic == kMovableBufferMagic) { end_of_header += - sizeof(Resizer *); // Room for resizer function for movable buffers. + sizeof(Resizer *); // Room for resizer function for movable buffers. header_size += sizeof(Resizer *); } FreeBlockHeader *f = reinterpret_cast(end_of_header); @@ -193,7 +197,7 @@ inline uint32_t AlignSize(uint32_t s, void *PayloadBuffer::Allocate(PayloadBuffer **buffer, uint32_t n, uint32_t alignment, bool clear) { - n = AlignSize(n, alignment); // Aligned. + n = AlignSize(n, alignment); // Aligned. size_t full_length = n + sizeof(uint32_t); FreeBlockHeader *free_block = (*buffer)->FreeList(); FreeBlockHeader *prev = nullptr; @@ -203,8 +207,8 @@ void *PayloadBuffer::Allocate(PayloadBuffer **buffer, uint32_t n, // header, take the lower part of the free block and keep the remainder // in the free list. n = (*buffer)->TakeStartOfFreeBlock(free_block, n, full_length, prev); - size_t *newblock = (size_t *)free_block; // Start of new block. - *newblock = n; // Size of allocated block. + size_t *newblock = (size_t *)free_block; // Start of new block. + *newblock = n; // Size of allocated block. void *addr = reinterpret_cast(uintptr_t(free_block) + sizeof(uint32_t)); if (clear) { @@ -275,6 +279,32 @@ void *PayloadBuffer::Allocate(PayloadBuffer **buffer, uint32_t n, return nullptr; } +std::vector PayloadBuffer::AllocateMany(PayloadBuffer **buffer, + uint32_t size, uint32_t n, + uint32_t alignment, + bool clear) { + // Calculate space for the whole block. This is n*aligned(size) + n*sizeof(uint32_t). + size_t full_length = n * (AlignSize(size, alignment) + sizeof(uint32_t)); + void* start = Allocate(buffer, full_length, 8, clear); + if (start == nullptr) { + return {}; // No memory. + } + if (clear) { + memset(start, 0, full_length); + } + // Divide the block into freeable chunks, each of which is align(size) bytes + // long. The length of each block is stored immediately before the block. + std::vector blocks; + uint32_t* p = reinterpret_cast(start); + for (uint32_t i = 0; i < n; i++) { + uint32_t len = AlignSize(size, alignment); + p[0] = len; + blocks.push_back(reinterpret_cast(p + 1)); + p = reinterpret_cast(reinterpret_cast(p) + len + sizeof(uint32_t)); + } + return blocks; +} + void PayloadBuffer::MergeWithAboveIfPossible(FreeBlockHeader *alloc_block, FreeBlockHeader *alloc_header, FreeBlockHeader *free_block, @@ -325,7 +355,7 @@ void PayloadBuffer::InsertNewFreeBlockAtEnd(FreeBlockHeader *free_block, void PayloadBuffer::Free(void *p) { // An allocated block has its length immediately before its address. uint32_t alloc_length = - *(reinterpret_cast(p) - 1); // Length of allocated block. + *(reinterpret_cast(p) - 1); // Length of allocated block. // Point to real start of allocated block. FreeBlockHeader *alloc_header = @@ -389,10 +419,10 @@ void PayloadBuffer::ShrinkBlock(FreeBlockHeader *alloc_block, if (rem >= sizeof(FreeBlockHeader)) { // If we are freeing enough to make a free block, free it, otherwise // there's nothing we can do and we just keep the block the same size. - *len_ptr = new_length; // Change size of block. + *len_ptr = new_length; // Change size of block. uint32_t *newp = reinterpret_cast( reinterpret_cast(alloc_block) + sizeof(uint32_t) + new_length); - *newp = rem - sizeof(uint32_t); // Add header for free. + *newp = rem - sizeof(uint32_t); // Add header for free. Free(newp + 1); } } @@ -460,7 +490,7 @@ void *PayloadBuffer::Realloc(PayloadBuffer **buffer, void *p, uint32_t n, reinterpret_cast(uintptr_t(p) - sizeof(uint32_t)); uintptr_t alloc_addr = (uintptr_t)p; - n = AlignSize(n); // Aligned. + n = AlignSize(n); // Aligned. if (n == orig_length) { // Same size as current block, nothing to do. return p; @@ -533,4 +563,4 @@ void *PayloadBuffer::Realloc(PayloadBuffer **buffer, void *p, uint32_t n, return newp; } -} // namespace toolbelt +} // namespace toolbelt diff --git a/toolbelt/payload_buffer.h b/toolbelt/payload_buffer.h index e55111f..ca3f987 100644 --- a/toolbelt/payload_buffer.h +++ b/toolbelt/payload_buffer.h @@ -232,6 +232,14 @@ struct PayloadBuffer { static void *Realloc(PayloadBuffer **buffer, void *p, uint32_t n, uint32_t alignment, bool clear = true); + // Allocate 'n' items of size 'size' in the buffer. The buffer might move. + // Each allocation is aligned to 'alignment' and is capable of being freed + // individually. The addresses of the allocated items are returned in a + // vector. + static std::vector AllocateMany(PayloadBuffer **buffer, uint32_t size, + uint32_t n, uint32_t alignment, + bool clear = true); + bool IsValidMagic() const { return magic == kFixedBufferMagic || magic == kMovableBufferMagic; } @@ -387,7 +395,6 @@ inline void PayloadBuffer::VectorReserve(PayloadBuffer **self, VectorHeader *hdr, size_t n) { if (hdr->data == 0) { void *vecp = Allocate(self, n * sizeof(T), 8); - std::cout << "empty vector " << vecp << std::endl; hdr->data = (*self)->ToOffset(vecp); } else { // Vector has some values in it. Retrieve the total size from diff --git a/toolbelt/payload_buffer_test.cc b/toolbelt/payload_buffer_test.cc index 4ba4f9a..8a88caa 100644 --- a/toolbelt/payload_buffer_test.cc +++ b/toolbelt/payload_buffer_test.cc @@ -1,5 +1,5 @@ -#include "toolbelt/payload_buffer.h" #include "toolbelt/hexdump.h" +#include "toolbelt/payload_buffer.h" #include #include @@ -92,6 +92,36 @@ TEST(BufferTest, FreeThenAlloc) { free(buffer); } +TEST(BufferTest, Many) { + constexpr size_t kSize = 8192; + char *buffer = (char *)malloc(kSize); + PayloadBuffer *pb = new (buffer) PayloadBuffer(kSize); + + std::vector addrs = + PayloadBuffer::AllocateMany(&pb, 100, 10, 8, true); + ASSERT_EQ(10, addrs.size()); + // Print the addresses. + for (auto addr : addrs) { + std::cout << "Allocated " << addr << std::endl; + } + pb->Dump(std::cout); + toolbelt::Hexdump(pb, pb->hwm); + + // Make sure we can free them. + pb->Free(addrs[0]); + pb->Free(addrs[2]); + + pb->Dump(std::cout); + toolbelt::Hexdump(pb, pb->hwm); + + // Now allocate one more to take the first free block. + void *addr = PayloadBuffer::Allocate(&pb, 100, 8); + pb->Dump(std::cout); + toolbelt::Hexdump(pb, pb->hwm); + ASSERT_EQ(addrs[0], addr); + free(buffer); +} + TEST(BufferTest, String) { char *buffer = (char *)malloc(4096); PayloadBuffer *pb = new (buffer) PayloadBuffer(4096);