Allow suggesting the initial block size without requiring an initial …

…block to be provided and managed separately. This can avoid small mallocs when the approximate arena size is known.

PiperOrigin-RevId: 713501871

upb/mem/arena.c

@@ -343,14 +343,16 @@ void* UPB_PRIVATE(_upb_Arena_SlowMalloc)(upb_Arena* a, size_t size) {
   return upb_Arena_Malloc(a, size - UPB_ASAN_GUARD_SIZE);
 }
 
-static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc) {
+static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc, size_t first_size) {
   const size_t first_block_overhead =
       sizeof(upb_ArenaState) + kUpb_MemblockReserve;
   upb_ArenaState* a;
 
   // We need to malloc the initial block.
   char* mem;
-  size_t block_size = first_block_overhead + 256;
+  size_t block_size =
+      first_block_overhead +
+      UPB_MAX(256, UPB_ALIGN_MALLOC(first_size) + UPB_ASAN_GUARD_SIZE);
   if (!alloc || !(mem = upb_malloc(alloc, block_size))) {
     return NULL;
   }
@@ -377,26 +379,25 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
   UPB_ASSERT(sizeof(void*) * UPB_ARENA_SIZE_HACK >= sizeof(upb_ArenaState));
   upb_ArenaState* a;
 
-  if (n) {
+  if (mem) {
     /* Align initial pointer up so that we return properly-aligned pointers. */
     void* aligned = (void*)UPB_ALIGN_UP((uintptr_t)mem, UPB_MALLOC_ALIGN);
     size_t delta = (uintptr_t)aligned - (uintptr_t)mem;
     n = delta <= n ? n - delta : 0;
     mem = aligned;
+    /* Round block size down to alignof(*a) since we will allocate the arena
+     * itself at the end. */
+    n = UPB_ALIGN_DOWN(n, UPB_ALIGN_OF(upb_ArenaState));
+  } else {
+    n = UPB_ALIGN_UP(n, UPB_ALIGN_OF(upb_ArenaState));
   }
 
-  /* Round block size down to alignof(*a) since we will allocate the arena
-   * itself at the end. */
-  n = UPB_ALIGN_DOWN(n, UPB_ALIGN_OF(upb_ArenaState));
-
-  if (UPB_UNLIKELY(n < sizeof(upb_ArenaState))) {
+  if (UPB_UNLIKELY(n < sizeof(upb_ArenaState) || !mem)) {
+    upb_Arena* ret = _upb_Arena_InitSlow(alloc, mem ? 0 : n);
 #ifdef UPB_TRACING_ENABLED
-    upb_Arena* ret = _upb_Arena_InitSlow(alloc);
     upb_Arena_LogInit(ret, n);
-    return ret;
-#else
-    return _upb_Arena_InitSlow(alloc);
 #endif
+    return ret;
   }
 
   a = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), upb_ArenaState);

upb/mem/arena.h

@@ -38,9 +38,14 @@ typedef void upb_AllocCleanupFunc(upb_alloc* alloc);
 extern "C" {
 #endif
 
-// Creates an arena from the given initial block (if any -- n may be 0).
-// Additional blocks will be allocated from |alloc|.  If |alloc| is NULL, this
-// is a fixed-size arena and cannot grow.
+// Creates an arena from the given initial block (if any -- mem may be NULL). If
+// an initial block is specified, the arena's lifetime cannot be extended by
+// |upb_Arena_IncRefFor| or |upb_Arena_Fuse|. Additional blocks will be
+// allocated from |alloc|. If |alloc| is NULL, this is a fixed-size arena and
+// cannot grow. If an initial block is specified, |n| is its length; if there is
+// no initial block, |n| is a hint of the size that should be allocated for the
+// first block of the arena, such that `upb_Arena_Malloc(hint)` will not require
+// another call to |alloc|.
 UPB_API upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc);
 
 UPB_API void upb_Arena_Free(upb_Arena* a);
@@ -76,6 +81,10 @@ UPB_API_INLINE upb_Arena* upb_Arena_New(void) {
   return upb_Arena_Init(NULL, 0, &upb_alloc_global);
 }
 
+UPB_API_INLINE upb_Arena* upb_Arena_NewSized(size_t size_hint) {
+  return upb_Arena_Init(NULL, size_hint, &upb_alloc_global);
+}
+
 UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size);
 
 UPB_API_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize,

upb/mem/arena.hpp

@@ -8,8 +8,10 @@
 #ifndef UPB_MEM_ARENA_HPP_
 #define UPB_MEM_ARENA_HPP_
 
+#include "upb/mem/alloc.h"
 #ifdef __cplusplus
 
+#include <cstddef>
 #include <memory>
 
 #include "upb/mem/arena.h"
@@ -23,6 +25,8 @@ class Arena {
   Arena(char* initial_block, size_t size)
       : ptr_(upb_Arena_Init(initial_block, size, &upb_alloc_global),
              upb_Arena_Free) {}
+  explicit Arena(size_t size)
+      : ptr_(upb_Arena_NewSized(size), upb_Arena_Free) {}
 
   upb_Arena* ptr() const { return ptr_.get(); }
 
@@ -35,26 +39,6 @@ class Arena {
  protected:
   std::unique_ptr<upb_Arena, decltype(&upb_Arena_Free)> ptr_;
 };
-
-// InlinedArena seeds the arenas with a predefined amount of memory. No heap
-// memory will be allocated until the initial block is exceeded.
-template <int N>
-class InlinedArena : public Arena {
- public:
-  InlinedArena() : Arena(initial_block_, N) {}
-  ~InlinedArena() {
-    // Explicitly destroy the arena now so that it does not outlive
-    // initial_block_.
-    ptr_.reset();
-  }
-
- private:
-  InlinedArena(const InlinedArena&) = delete;
-  InlinedArena& operator=(const InlinedArena&) = delete;
-
-  char initial_block_[N];
-};
-
 }  // namespace upb
 
 #endif  // __cplusplus

upb/mem/arena_test.cc

@@ -110,6 +110,23 @@ TEST(ArenaTest, SizedFree) {
   EXPECT_EQ(sizes.size(), 0);
 }
 
+TEST(ArenaTest, SizeHint) {
+  absl::flat_hash_map<void*, size_t> sizes;
+  SizeTracker alloc;
+  alloc.alloc.func = size_checking_allocfunc;
+  alloc.delegate_alloc = &upb_alloc_global;
+  alloc.sizes = &sizes;
+
+  upb_Arena* arena = upb_Arena_Init(nullptr, 2459, &alloc.alloc);
+  EXPECT_EQ(sizes.size(), 1);
+  EXPECT_NE(upb_Arena_Malloc(arena, 2459), nullptr);
+  EXPECT_EQ(sizes.size(), 1);
+  EXPECT_NE(upb_Arena_Malloc(arena, 500), nullptr);
+  EXPECT_EQ(sizes.size(), 2);
+  upb_Arena_Free(arena);
+  EXPECT_EQ(sizes.size(), 0);
+}
+
 TEST(ArenaTest, ArenaFuse) {
   upb_Arena* arena1 = upb_Arena_New();
   upb_Arena* arena2 = upb_Arena_New();