From 665a347447734252a090c8ea479f377122f0c068 Mon Sep 17 00:00:00 2001
From: mutouyun <orz@orzz.org>
Date: Sun, 17 Sep 2023 14:44:48 +0800
Subject: [PATCH] fix some bugs

---
 include/libpmr/monotonic_buffer_resource.h |  5 +-
 src/libpmr/monotonic_buffer_resource.cpp   | 61 +++++++++++++---------
 2 files changed, 40 insertions(+), 26 deletions(-)

diff --git a/include/libpmr/monotonic_buffer_resource.h b/include/libpmr/monotonic_buffer_resource.h
index fb8079b7..6b0d4dfd 100644
--- a/include/libpmr/monotonic_buffer_resource.h
+++ b/include/libpmr/monotonic_buffer_resource.h
@@ -34,9 +34,10 @@ class LIBIMP_EXPORT monotonic_buffer_resource {
 
   ::LIBIMP::byte *head_;
   ::LIBIMP::byte *tail_;
+  std::size_t next_size_;
 
-  ::LIBIMP::span<::LIBIMP::byte> const initial_buffer_;
-  std::size_t initial_size_;
+  ::LIBIMP::byte * const initial_buffer_;
+  std::size_t const initial_size_;
 
 public:
   monotonic_buffer_resource() noexcept;
diff --git a/src/libpmr/monotonic_buffer_resource.cpp b/src/libpmr/monotonic_buffer_resource.cpp
index 9d30f3ac..195aa862 100644
--- a/src/libpmr/monotonic_buffer_resource.cpp
+++ b/src/libpmr/monotonic_buffer_resource.cpp
@@ -1,5 +1,7 @@
 
 #include <utility>
+#include <memory>
+#include <algorithm>
 
 #include "libimp/log.h"
 #include "libimp/aligned.h"
@@ -23,38 +25,41 @@ Node *make_node(allocator const &upstream, std::size_t initial_size, std::size_t
   return node;
 }
 
+std::size_t next_buffer_size(std::size_t size) noexcept {
+  return size * 3 / 2;
+}
+
 } // namespace
 
 monotonic_buffer_resource::monotonic_buffer_resource() noexcept
   : monotonic_buffer_resource(allocator{}) {}
 
 monotonic_buffer_resource::monotonic_buffer_resource(allocator upstream) noexcept
-  : upstream_ (std::move(upstream))
-  , free_list_(nullptr)
-  , head_     (nullptr)
-  , tail_     (nullptr)
-  , initial_size_(0) {}
+  : monotonic_buffer_resource(0, std::move(upstream)) {}
 
 monotonic_buffer_resource::monotonic_buffer_resource(std::size_t initial_size)
   : monotonic_buffer_resource(initial_size, allocator{}) {}
 
 monotonic_buffer_resource::monotonic_buffer_resource(std::size_t initial_size, allocator upstream)
-  : upstream_ (std::move(upstream))
-  , free_list_(make_node<node>(upstream_, initial_size, alignof(std::max_align_t)))
-  , head_     (free_list_ == nullptr ? reinterpret_cast<::LIBIMP::byte *>(free_list_) + sizeof(node) : nullptr)
-  , tail_     (head_ + initial_size)
-  , initial_size_(initial_size) {}
+  : upstream_      (std::move(upstream))
+  , free_list_     (nullptr)
+  , head_          (nullptr)
+  , tail_          (nullptr)
+  , next_size_     (initial_size)
+  , initial_buffer_(nullptr)
+  , initial_size_  (initial_size) {}
 
 monotonic_buffer_resource::monotonic_buffer_resource(::LIBIMP::span<::LIBIMP::byte> buffer) noexcept
   : monotonic_buffer_resource(buffer, allocator{}) {}
 
 monotonic_buffer_resource::monotonic_buffer_resource(::LIBIMP::span<::LIBIMP::byte> buffer, allocator upstream) noexcept
-  : upstream_ (std::move(upstream))
-  , free_list_(nullptr)
-  , head_     (buffer.begin())
-  , tail_     (buffer.end())
-  , initial_buffer_(buffer)
-  , initial_size_(initial_buffer_.size()) {}
+  : upstream_      (std::move(upstream))
+  , free_list_     (nullptr)
+  , head_          (buffer.begin())
+  , tail_          (buffer.end())
+  , next_size_     (next_buffer_size(buffer.size()))
+  , initial_buffer_(buffer.begin())
+  , initial_size_  (buffer.size()) {}
 
 monotonic_buffer_resource::~monotonic_buffer_resource() {
   release();
@@ -70,8 +75,14 @@ void monotonic_buffer_resource::release() {
     upstream_.deallocate(free_list_, free_list_->size);
     free_list_ = next;
   }
-  head_ = initial_buffer_.begin();
-  tail_ = initial_buffer_.end();
+  // reset to initial state at contruction
+  if ((head_ = initial_buffer_) != nullptr) {
+    tail_ = head_ + initial_size_;
+    next_size_ = next_buffer_size(initial_size_);
+  } else {
+    tail_ = nullptr;
+    next_size_ = initial_size_;
+  }
 }
 
 void *monotonic_buffer_resource::allocate(std::size_t bytes, std::size_t alignment) noexcept {
@@ -81,19 +92,21 @@ void *monotonic_buffer_resource::allocate(std::size_t bytes, std::size_t alignme
     return nullptr;
   }
   if ((alignment & (alignment - 1)) != 0) {
-    log.error("failed: allocate alignment is not power of 2.");
+    log.error("failed: allocate alignment is not a power of 2.");
     return nullptr;
   }
-  auto *p = reinterpret_cast<::LIBIMP::byte *>(::LIBIMP::round_up(reinterpret_cast<std::size_t>(head_), alignment));
-  auto remain = static_cast<std::size_t>(tail_ - p);
-  if (remain < bytes) {
-    initial_size_ = (std::max)(initial_size_ * 2, bytes);
-    auto *node = make_node<monotonic_buffer_resource::node>(upstream_, initial_size_, alignment);
+  void *p = head_;
+  auto  s = static_cast<std::size_t>(tail_ - head_);
+  if (std::align(alignment, bytes, p, s) == nullptr) {
+    next_size_ = (std::max)(next_size_, bytes);
+    auto *node = make_node<monotonic_buffer_resource::node>(upstream_, next_size_, alignment);
     if (node == nullptr) return nullptr;
     node->next = free_list_;
     free_list_ = node;
+    next_size_ = next_buffer_size(next_size_);
     p = reinterpret_cast<::LIBIMP::byte *>(free_list_) + ::LIBIMP::round_up(sizeof(node), alignment);
   }
+  head_ = static_cast<::LIBIMP::byte *>(p) + bytes;
   return p;
 }