StringBuilder: allow dynamic resize

GitOrigin-RevId: cd59f013d10f4d8e5911933eeb75384d675fda94

tdutils/td/utils/StringBuilder.cpp

@@ -61,8 +61,34 @@ static char *print_int(char *current_ptr, T x) {
   return print_uint(current_ptr, x);
 }
 
+bool StringBuilder::reserve_inner(size_t size) {
+  if (!use_buffer_) {
+    return false;
+  }
+  //TODO: check size_t oveflow
+  size_t old_data_size = current_ptr_ - begin_ptr_;
+  size_t need_data_size = old_data_size + size;
+  size_t old_buffer_size = end_ptr_ - begin_ptr_;
+  size_t new_buffer_size = (old_buffer_size + 1) * 2;
+  if (new_buffer_size < need_data_size) {
+    new_buffer_size = need_data_size;
+  }
+  if (new_buffer_size < 100) {
+    new_buffer_size = 100;
+  }
+  new_buffer_size += reserved_size;
+  auto new_buffer_ = std::make_unique<char[]>(new_buffer_size);
+  std::memcpy(new_buffer_.get(), begin_ptr_, old_data_size);
+  buffer_ = std::move(new_buffer_);
+  begin_ptr_ = buffer_.get();
+  current_ptr_ = begin_ptr_ + old_data_size;
+  end_ptr_ = begin_ptr_ + new_buffer_size - reserved_size;
+  CHECK(end_ptr_ > begin_ptr_);
+  return true;
+}
+
 StringBuilder &StringBuilder::operator<<(int x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ = print_int(current_ptr_, x);
@@ -70,7 +96,7 @@ StringBuilder &StringBuilder::operator<<(int x) {
 }
 
 StringBuilder &StringBuilder::operator<<(unsigned int x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ = print_uint(current_ptr_, x);
@@ -78,7 +104,7 @@ StringBuilder &StringBuilder::operator<<(unsigned int x) {
 }
 
 StringBuilder &StringBuilder::operator<<(long int x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ = print_int(current_ptr_, x);
@@ -86,7 +112,7 @@ StringBuilder &StringBuilder::operator<<(long int x) {
 }
 
 StringBuilder &StringBuilder::operator<<(long unsigned int x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ = print_uint(current_ptr_, x);
@@ -94,7 +120,7 @@ StringBuilder &StringBuilder::operator<<(long unsigned int x) {
 }
 
 StringBuilder &StringBuilder::operator<<(long long int x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ = print_int(current_ptr_, x);
@@ -102,7 +128,7 @@ StringBuilder &StringBuilder::operator<<(long long int x) {
 }
 
 StringBuilder &StringBuilder::operator<<(long long unsigned int x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ = print_uint(current_ptr_, x);
@@ -110,7 +136,7 @@ StringBuilder &StringBuilder::operator<<(long long unsigned int x) {
 }
 
 StringBuilder &StringBuilder::operator<<(FixedDouble x) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
 
@@ -137,7 +163,7 @@ StringBuilder &StringBuilder::operator<<(FixedDouble x) {
 }
 
 StringBuilder &StringBuilder::operator<<(const void *ptr) {
-  if (unlikely(end_ptr_ < current_ptr_)) {
+  if (unlikely(!reserve())) {
     return on_error();
   }
   current_ptr_ += std::snprintf(current_ptr_, reserved_size, "%p", ptr);

tdutils/td/utils/StringBuilder.h

@@ -18,10 +18,17 @@ namespace td {
 
 class StringBuilder {
  public:
-  explicit StringBuilder(MutableSlice slice)
-      : begin_ptr_(slice.begin()), current_ptr_(begin_ptr_), end_ptr_(slice.end() - reserved_size) {
+  explicit StringBuilder(MutableSlice slice, bool use_buffer = false)
+      : begin_ptr_(slice.begin())
+      , current_ptr_(begin_ptr_)
+      , end_ptr_(slice.end() - reserved_size)
+      , use_buffer_(use_buffer) {
     if (slice.size() <= reserved_size) {
-      std::abort();  // shouldn't happen
+      auto buffer_size = reserved_size + 100;
+      buffer_ = std::make_unique<char[]>(buffer_size);
+      begin_ptr_ = buffer_.get();
+      current_ptr_ = begin_ptr_;
+      end_ptr_ = begin_ptr_ + buffer_size - reserved_size;
     }
   }
 
@@ -49,15 +56,18 @@ class StringBuilder {
   StringBuilder &operator<<(const wchar_t *str) = delete;
 
   StringBuilder &operator<<(Slice slice) {
-    if (unlikely(end_ptr_ < current_ptr_)) {
-      return on_error();
-    }
-    auto size = static_cast<size_t>(end_ptr_ + reserved_size - 1 - current_ptr_);
-    if (unlikely(slice.size() > size)) {
-      error_flag_ = true;
-    } else {
-      size = slice.size();
+    size_t size = slice.size();
+    if (unlikely(!reserve(size))) {
+      if (end_ptr_ < current_ptr_) {
+        return on_error();
+      }
+      auto available_size = static_cast<size_t>(end_ptr_ + reserved_size - 1 - current_ptr_);
+      if (size > available_size) {
+        error_flag_ = true;
+        size = available_size;
+      }
     }
+
     std::memcpy(current_ptr_, slice.begin(), size);
     current_ptr_ += size;
     return *this;
@@ -68,7 +78,7 @@ class StringBuilder {
   }
 
   StringBuilder &operator<<(char c) {
-    if (unlikely(end_ptr_ < current_ptr_)) {
+    if (unlikely(!reserve())) {
       return on_error();
     }
     *current_ptr_++ = c;
@@ -120,12 +130,28 @@ class StringBuilder {
   char *current_ptr_;
   char *end_ptr_;
   bool error_flag_ = false;
+  bool use_buffer_ = false;
+  std::unique_ptr<char[]> buffer_;
   static constexpr size_t reserved_size = 30;
 
   StringBuilder &on_error() {
     error_flag_ = true;
     return *this;
   }
+
+  bool reserve() {
+    if (end_ptr_ > current_ptr_) {
+      return true;
+    }
+    return reserve_inner(reserved_size);
+  }
+  bool reserve(size_t size) {
+    if (end_ptr_ > current_ptr_ && static_cast<size_t>(end_ptr_ - current_ptr_) >= size) {
+      return true;
+    }
+    return reserve_inner(size);
+  }
+  bool reserve_inner(size_t size);
 };
 
 template <class T>

tdutils/test/misc.cpp

@@ -492,3 +492,29 @@ TEST(Misc, full_split) {
   test_full_split(" ab cd ef ", {"", "ab", "cd", "ef", ""});
   test_full_split("  ab  cd  ef  ", {"", "", "ab", "", "cd", "", "ef", "", ""});
 }
+
+TEST(Misc, StringBuilder) {
+  auto small = std::string{"abcdefghij"};
+  auto big = std::string(1000, 'a');
+  using V = std::vector<std::string>;
+  for (auto use_buf : {false, true}) {
+    for (size_t initial_buffer_size : {0, 1, 5, 10, 100, 1000, 2000}) {
+      for (auto test : {V{small}, {small, big, big, small}, {big, small, big}}) {
+        std::string buf(initial_buffer_size, '\0');
+        td::StringBuilder sb(buf, use_buf);
+        std::string res;
+        for (auto x : test) {
+          res += x;
+          sb << x;
+        }
+        if (use_buf) {
+          ASSERT_EQ(res, sb.as_cslice());
+        } else {
+          auto got = sb.as_cslice();
+          res.resize(got.size());
+          ASSERT_EQ(res, got);
+        }
+      }
+    }
+  }
+}