grpc 1.62.0 → 1.64.0.pre1

data/third_party/upb/upb/mem/arena.c

@@ -5,33 +5,124 @@
 // license that can be found in the LICENSE file or at
 // https://developers.google.com/open-source/licenses/bsd
 
-#include "upb/mem/internal/arena.h"
+#include "upb/mem/arena.h"
+
+#include <stddef.h>
+#include <stdint.h>
 
+#include "upb/mem/alloc.h"
+#include "upb/mem/internal/arena.h"
 #include "upb/port/atomic.h"
 
 // Must be last.
 #include "upb/port/def.inc"
 
-struct _upb_MemBlock {
+typedef struct upb_MemBlock {
   // Atomic only for the benefit of SpaceAllocated().
-  UPB_ATOMIC(_upb_MemBlock*) next;
+  UPB_ATOMIC(struct upb_MemBlock*) next;
   uint32_t size;
   // Data follows.
-};
+} upb_MemBlock;
+
+typedef struct upb_ArenaInternal {
+  // upb_alloc* together with a low bit which signals if there is an initial
+  // block.
+  uintptr_t block_alloc;
+
+  // When multiple arenas are fused together, each arena points to a parent
+  // arena (root points to itself). The root tracks how many live arenas
+  // reference it.
+
+  // The low bit is tagged:
+  //   0: pointer to parent
+  //   1: count, left shifted by one
+  UPB_ATOMIC(uintptr_t) parent_or_count;
+
+  // All nodes that are fused together are in a singly-linked list.
+  // == NULL at end of list.
+  UPB_ATOMIC(struct upb_ArenaInternal*) next;
+
+  // The last element of the linked list. This is present only as an
+  // optimization, so that we do not have to iterate over all members for every
+  // fuse.  Only significant for an arena root. In other cases it is ignored.
+  // == self when no other list members.
+  UPB_ATOMIC(struct upb_ArenaInternal*) tail;
+
+  // Linked list of blocks to free/cleanup. Atomic only for the benefit of
+  // upb_Arena_SpaceAllocated().
+  UPB_ATOMIC(upb_MemBlock*) blocks;
+} upb_ArenaInternal;
+
+// All public + private state for an arena.
+typedef struct {
+  upb_Arena head;
+  upb_ArenaInternal body;
+} upb_ArenaState;
+
+typedef struct {
+  upb_ArenaInternal* root;
+  uintptr_t tagged_count;
+} upb_ArenaRoot;
 
-static const size_t memblock_reserve =
-    UPB_ALIGN_UP(sizeof(_upb_MemBlock), UPB_MALLOC_ALIGN);
+static const size_t kUpb_MemblockReserve =
+    UPB_ALIGN_UP(sizeof(upb_MemBlock), UPB_MALLOC_ALIGN);
 
-typedef struct _upb_ArenaRoot {
-  upb_Arena* root;
-  uintptr_t tagged_count;
-} _upb_ArenaRoot;
+// Extracts the (upb_ArenaInternal*) from a (upb_Arena*)
+static upb_ArenaInternal* upb_Arena_Internal(const upb_Arena* a) {
+  return &((upb_ArenaState*)a)->body;
+}
+
+static bool _upb_Arena_IsTaggedRefcount(uintptr_t parent_or_count) {
+  return (parent_or_count & 1) == 1;
+}
+
+static bool _upb_Arena_IsTaggedPointer(uintptr_t parent_or_count) {
+  return (parent_or_count & 1) == 0;
+}
+
+static uintptr_t _upb_Arena_RefCountFromTagged(uintptr_t parent_or_count) {
+  UPB_ASSERT(_upb_Arena_IsTaggedRefcount(parent_or_count));
+  return parent_or_count >> 1;
+}
+
+static uintptr_t _upb_Arena_TaggedFromRefcount(uintptr_t refcount) {
+  uintptr_t parent_or_count = (refcount << 1) | 1;
+  UPB_ASSERT(_upb_Arena_IsTaggedRefcount(parent_or_count));
+  return parent_or_count;
+}
+
+static upb_ArenaInternal* _upb_Arena_PointerFromTagged(
+    uintptr_t parent_or_count) {
+  UPB_ASSERT(_upb_Arena_IsTaggedPointer(parent_or_count));
+  return (upb_ArenaInternal*)parent_or_count;
+}
+
+static uintptr_t _upb_Arena_TaggedFromPointer(upb_ArenaInternal* ai) {
+  uintptr_t parent_or_count = (uintptr_t)ai;
+  UPB_ASSERT(_upb_Arena_IsTaggedPointer(parent_or_count));
+  return parent_or_count;
+}
 
-static _upb_ArenaRoot _upb_Arena_FindRoot(upb_Arena* a) {
-  uintptr_t poc = upb_Atomic_Load(&a->parent_or_count, memory_order_acquire);
+static upb_alloc* _upb_ArenaInternal_BlockAlloc(upb_ArenaInternal* ai) {
+  return (upb_alloc*)(ai->block_alloc & ~0x1);
+}
+
+static uintptr_t _upb_Arena_MakeBlockAlloc(upb_alloc* alloc, bool has_initial) {
+  uintptr_t alloc_uint = (uintptr_t)alloc;
+  UPB_ASSERT((alloc_uint & 1) == 0);
+  return alloc_uint | (has_initial ? 1 : 0);
+}
+
+static bool _upb_ArenaInternal_HasInitialBlock(upb_ArenaInternal* ai) {
+  return ai->block_alloc & 0x1;
+}
+
+static upb_ArenaRoot _upb_Arena_FindRoot(upb_Arena* a) {
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  uintptr_t poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
   while (_upb_Arena_IsTaggedPointer(poc)) {
-    upb_Arena* next = _upb_Arena_PointerFromTagged(poc);
-    UPB_ASSERT(a != next);
+    upb_ArenaInternal* next = _upb_Arena_PointerFromTagged(poc);
+    UPB_ASSERT(ai != next);
     uintptr_t next_poc =
         upb_Atomic_Load(&next->parent_or_count, memory_order_acquire);
 
@@ -55,104 +146,111 @@ static _upb_ArenaRoot _upb_Arena_FindRoot(upb_Arena* a) {
       // further away over time, but the path towards that root will continue to
       // be valid and the creation of the path carries all the memory orderings
       // required.
-      UPB_ASSERT(a != _upb_Arena_PointerFromTagged(next_poc));
-      upb_Atomic_Store(&a->parent_or_count, next_poc, memory_order_relaxed);
+      UPB_ASSERT(ai != _upb_Arena_PointerFromTagged(next_poc));
+      upb_Atomic_Store(&ai->parent_or_count, next_poc, memory_order_relaxed);
     }
-    a = next;
+    ai = next;
     poc = next_poc;
   }
-  return (_upb_ArenaRoot){.root = a, .tagged_count = poc};
+  return (upb_ArenaRoot){.root = ai, .tagged_count = poc};
 }
 
-size_t upb_Arena_SpaceAllocated(upb_Arena* arena) {
-  arena = _upb_Arena_FindRoot(arena).root;
+size_t upb_Arena_SpaceAllocated(upb_Arena* arena, size_t* fused_count) {
+  upb_ArenaInternal* ai = _upb_Arena_FindRoot(arena).root;
   size_t memsize = 0;
+  size_t local_fused_count = 0;
 
-  while (arena != NULL) {
-    _upb_MemBlock* block =
-        upb_Atomic_Load(&arena->blocks, memory_order_relaxed);
+  while (ai != NULL) {
+    upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_relaxed);
     while (block != NULL) {
-      memsize += sizeof(_upb_MemBlock) + block->size;
+      memsize += sizeof(upb_MemBlock) + block->size;
       block = upb_Atomic_Load(&block->next, memory_order_relaxed);
     }
-    arena = upb_Atomic_Load(&arena->next, memory_order_relaxed);
+    ai = upb_Atomic_Load(&ai->next, memory_order_relaxed);
+    local_fused_count++;
   }
 
+  if (fused_count) *fused_count = local_fused_count;
   return memsize;
 }
 
 uint32_t upb_Arena_DebugRefCount(upb_Arena* a) {
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
   // These loads could probably be relaxed, but given that this is debug-only,
   // it's not worth introducing a new variant for it.
-  uintptr_t poc = upb_Atomic_Load(&a->parent_or_count, memory_order_acquire);
+  uintptr_t poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
   while (_upb_Arena_IsTaggedPointer(poc)) {
-    a = _upb_Arena_PointerFromTagged(poc);
-    poc = upb_Atomic_Load(&a->parent_or_count, memory_order_acquire);
+    ai = _upb_Arena_PointerFromTagged(poc);
+    poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
   }
   return _upb_Arena_RefCountFromTagged(poc);
 }
 
-static void upb_Arena_AddBlock(upb_Arena* a, void* ptr, size_t size) {
-  _upb_MemBlock* block = ptr;
+static void _upb_Arena_AddBlock(upb_Arena* a, void* ptr, size_t size) {
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  upb_MemBlock* block = ptr;
 
   // Insert into linked list.
   block->size = (uint32_t)size;
-  upb_Atomic_Init(&block->next, a->blocks);
-  upb_Atomic_Store(&a->blocks, block, memory_order_release);
+  upb_Atomic_Init(&block->next, ai->blocks);
+  upb_Atomic_Store(&ai->blocks, block, memory_order_release);
 
-  a->head.ptr = UPB_PTR_AT(block, memblock_reserve, char);
-  a->head.end = UPB_PTR_AT(block, size, char);
+  a->UPB_PRIVATE(ptr) = UPB_PTR_AT(block, kUpb_MemblockReserve, char);
+  a->UPB_PRIVATE(end) = UPB_PTR_AT(block, size, char);
 
-  UPB_POISON_MEMORY_REGION(a->head.ptr, a->head.end - a->head.ptr);
+  UPB_POISON_MEMORY_REGION(a->UPB_PRIVATE(ptr),
+                           a->UPB_PRIVATE(end) - a->UPB_PRIVATE(ptr));
 }
 
-static bool upb_Arena_AllocBlock(upb_Arena* a, size_t size) {
-  if (!a->block_alloc) return false;
-  _upb_MemBlock* last_block = upb_Atomic_Load(&a->blocks, memory_order_acquire);
+static bool _upb_Arena_AllocBlock(upb_Arena* a, size_t size) {
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  if (!ai->block_alloc) return false;
+  upb_MemBlock* last_block = upb_Atomic_Load(&ai->blocks, memory_order_acquire);
   size_t last_size = last_block != NULL ? last_block->size : 128;
-  size_t block_size = UPB_MAX(size, last_size * 2) + memblock_reserve;
-  _upb_MemBlock* block = upb_malloc(upb_Arena_BlockAlloc(a), block_size);
+  size_t block_size = UPB_MAX(size, last_size * 2) + kUpb_MemblockReserve;
+  upb_MemBlock* block =
+      upb_malloc(_upb_ArenaInternal_BlockAlloc(ai), block_size);
 
   if (!block) return false;
-  upb_Arena_AddBlock(a, block, block_size);
+  _upb_Arena_AddBlock(a, block, block_size);
+  UPB_ASSERT(UPB_PRIVATE(_upb_ArenaHas)(a) >= size);
   return true;
 }
 
-void* _upb_Arena_SlowMalloc(upb_Arena* a, size_t size) {
-  if (!upb_Arena_AllocBlock(a, size)) return NULL; /* Out of memory. */
-  UPB_ASSERT(_upb_ArenaHas(a) >= size);
-  return upb_Arena_Malloc(a, size);
+void* UPB_PRIVATE(_upb_Arena_SlowMalloc)(upb_Arena* a, size_t size) {
+  if (!_upb_Arena_AllocBlock(a, size)) return NULL;  // OOM
+  return upb_Arena_Malloc(a, size - UPB_ASAN_GUARD_SIZE);
 }
 
-/* Public Arena API ***********************************************************/
-
-static upb_Arena* upb_Arena_InitSlow(upb_alloc* alloc) {
-  const size_t first_block_overhead = sizeof(upb_Arena) + memblock_reserve;
-  upb_Arena* a;
+static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc) {
+  const size_t first_block_overhead =
+      sizeof(upb_ArenaState) + kUpb_MemblockReserve;
+  upb_ArenaState* a;
 
-  /* We need to malloc the initial block. */
+  // We need to malloc the initial block.
   char* mem;
   size_t n = first_block_overhead + 256;
   if (!alloc || !(mem = upb_malloc(alloc, n))) {
     return NULL;
   }
 
-  a = UPB_PTR_AT(mem, n - sizeof(*a), upb_Arena);
-  n -= sizeof(*a);
+  a = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), upb_ArenaState);
+  n -= sizeof(upb_ArenaState);
 
-  a->block_alloc = upb_Arena_MakeBlockAlloc(alloc, 0);
-  upb_Atomic_Init(&a->parent_or_count, _upb_Arena_TaggedFromRefcount(1));
-  upb_Atomic_Init(&a->next, NULL);
-  upb_Atomic_Init(&a->tail, a);
-  upb_Atomic_Init(&a->blocks, NULL);
+  a->body.block_alloc = _upb_Arena_MakeBlockAlloc(alloc, 0);
+  upb_Atomic_Init(&a->body.parent_or_count, _upb_Arena_TaggedFromRefcount(1));
+  upb_Atomic_Init(&a->body.next, NULL);
+  upb_Atomic_Init(&a->body.tail, &a->body);
+  upb_Atomic_Init(&a->body.blocks, NULL);
 
-  upb_Arena_AddBlock(a, mem, n);
+  _upb_Arena_AddBlock(&a->head, mem, n);
 
-  return a;
+  return &a->head;
 }
 
 upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
-  upb_Arena* a;
+  UPB_ASSERT(sizeof(void*) * UPB_ARENA_SIZE_HACK >= sizeof(upb_ArenaState));
+  upb_ArenaState* a;
 
   if (n) {
     /* Align initial pointer up so that we return properly-aligned pointers. */
@@ -164,63 +262,65 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
 
   /* 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_Arena));
+  n = UPB_ALIGN_DOWN(n, UPB_ALIGN_OF(upb_ArenaState));
 
-  if (UPB_UNLIKELY(n < sizeof(upb_Arena))) {
-    return upb_Arena_InitSlow(alloc);
+  if (UPB_UNLIKELY(n < sizeof(upb_ArenaState))) {
+    return _upb_Arena_InitSlow(alloc);
   }
 
-  a = UPB_PTR_AT(mem, n - sizeof(*a), upb_Arena);
+  a = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), upb_ArenaState);
 
-  upb_Atomic_Init(&a->parent_or_count, _upb_Arena_TaggedFromRefcount(1));
-  upb_Atomic_Init(&a->next, NULL);
-  upb_Atomic_Init(&a->tail, a);
-  upb_Atomic_Init(&a->blocks, NULL);
-  a->block_alloc = upb_Arena_MakeBlockAlloc(alloc, 1);
-  a->head.ptr = mem;
-  a->head.end = UPB_PTR_AT(mem, n - sizeof(*a), char);
+  upb_Atomic_Init(&a->body.parent_or_count, _upb_Arena_TaggedFromRefcount(1));
+  upb_Atomic_Init(&a->body.next, NULL);
+  upb_Atomic_Init(&a->body.tail, &a->body);
+  upb_Atomic_Init(&a->body.blocks, NULL);
 
-  return a;
+  a->body.block_alloc = _upb_Arena_MakeBlockAlloc(alloc, 1);
+  a->head.UPB_PRIVATE(ptr) = mem;
+  a->head.UPB_PRIVATE(end) = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), char);
+
+  return &a->head;
 }
 
-static void arena_dofree(upb_Arena* a) {
-  UPB_ASSERT(_upb_Arena_RefCountFromTagged(a->parent_or_count) == 1);
+static void _upb_Arena_DoFree(upb_ArenaInternal* ai) {
+  UPB_ASSERT(_upb_Arena_RefCountFromTagged(ai->parent_or_count) == 1);
 
-  while (a != NULL) {
+  while (ai != NULL) {
     // Load first since arena itself is likely from one of its blocks.
-    upb_Arena* next_arena =
-        (upb_Arena*)upb_Atomic_Load(&a->next, memory_order_acquire);
-    upb_alloc* block_alloc = upb_Arena_BlockAlloc(a);
-    _upb_MemBlock* block = upb_Atomic_Load(&a->blocks, memory_order_acquire);
+    upb_ArenaInternal* next_arena =
+        (upb_ArenaInternal*)upb_Atomic_Load(&ai->next, memory_order_acquire);
+    upb_alloc* block_alloc = _upb_ArenaInternal_BlockAlloc(ai);
+    upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_acquire);
     while (block != NULL) {
       // Load first since we are deleting block.
-      _upb_MemBlock* next_block =
+      upb_MemBlock* next_block =
           upb_Atomic_Load(&block->next, memory_order_acquire);
       upb_free(block_alloc, block);
       block = next_block;
     }
-    a = next_arena;
+    ai = next_arena;
   }
 }
 
 void upb_Arena_Free(upb_Arena* a) {
-  uintptr_t poc = upb_Atomic_Load(&a->parent_or_count, memory_order_acquire);
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  uintptr_t poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
 retry:
   while (_upb_Arena_IsTaggedPointer(poc)) {
-    a = _upb_Arena_PointerFromTagged(poc);
-    poc = upb_Atomic_Load(&a->parent_or_count, memory_order_acquire);
+    ai = _upb_Arena_PointerFromTagged(poc);
+    poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
   }
 
   // compare_exchange or fetch_sub are RMW operations, which are more
   // expensive then direct loads.  As an optimization, we only do RMW ops
   // when we need to update things for other threads to see.
   if (poc == _upb_Arena_TaggedFromRefcount(1)) {
-    arena_dofree(a);
+    _upb_Arena_DoFree(ai);
     return;
   }
 
   if (upb_Atomic_CompareExchangeWeak(
-          &a->parent_or_count, &poc,
+          &ai->parent_or_count, &poc,
           _upb_Arena_TaggedFromRefcount(_upb_Arena_RefCountFromTagged(poc) - 1),
           memory_order_release, memory_order_acquire)) {
     // We were >1 and we decremented it successfully, so we are done.
@@ -232,12 +332,14 @@ retry:
   goto retry;
 }
 
-static void _upb_Arena_DoFuseArenaLists(upb_Arena* const parent,
-                                        upb_Arena* child) {
-  upb_Arena* parent_tail = upb_Atomic_Load(&parent->tail, memory_order_relaxed);
+static void _upb_Arena_DoFuseArenaLists(upb_ArenaInternal* const parent,
+                                        upb_ArenaInternal* child) {
+  upb_ArenaInternal* parent_tail =
+      upb_Atomic_Load(&parent->tail, memory_order_relaxed);
+
   do {
     // Our tail might be stale, but it will always converge to the true tail.
-    upb_Arena* parent_tail_next =
+    upb_ArenaInternal* parent_tail_next =
         upb_Atomic_Load(&parent_tail->next, memory_order_relaxed);
     while (parent_tail_next != NULL) {
       parent_tail = parent_tail_next;
@@ -245,7 +347,7 @@ static void _upb_Arena_DoFuseArenaLists(upb_Arena* const parent,
           upb_Atomic_Load(&parent_tail->next, memory_order_relaxed);
     }
 
-    upb_Arena* displaced =
+    upb_ArenaInternal* displaced =
         upb_Atomic_Exchange(&parent_tail->next, child, memory_order_relaxed);
     parent_tail = upb_Atomic_Load(&child->tail, memory_order_relaxed);
 
@@ -257,8 +359,8 @@ static void _upb_Arena_DoFuseArenaLists(upb_Arena* const parent,
   upb_Atomic_Store(&parent->tail, parent_tail, memory_order_relaxed);
 }
 
-static upb_Arena* _upb_Arena_DoFuse(upb_Arena* a1, upb_Arena* a2,
-                                    uintptr_t* ref_delta) {
+static upb_ArenaInternal* _upb_Arena_DoFuse(upb_Arena* a1, upb_Arena* a2,
+                                            uintptr_t* ref_delta) {
   // `parent_or_count` has two disctint modes
   // -  parent pointer mode
   // -  refcount mode
@@ -266,14 +368,14 @@ static upb_Arena* _upb_Arena_DoFuse(upb_Arena* a1, upb_Arena* a2,
   // In parent pointer mode, it may change what pointer it refers to in the
   // tree, but it will always approach a root.  Any operation that walks the
   // tree to the root may collapse levels of the tree concurrently.
-  _upb_ArenaRoot r1 = _upb_Arena_FindRoot(a1);
-  _upb_ArenaRoot r2 = _upb_Arena_FindRoot(a2);
+  upb_ArenaRoot r1 = _upb_Arena_FindRoot(a1);
+  upb_ArenaRoot r2 = _upb_Arena_FindRoot(a2);
 
   if (r1.root == r2.root) return r1.root;  // Already fused.
 
   // Avoid cycles by always fusing into the root with the lower address.
   if ((uintptr_t)r1.root > (uintptr_t)r2.root) {
-    _upb_ArenaRoot tmp = r1;
+    upb_ArenaRoot tmp = r1;
     r1 = r2;
     r2 = tmp;
   }
@@ -316,7 +418,8 @@ static upb_Arena* _upb_Arena_DoFuse(upb_Arena* a1, upb_Arena* a2,
   return r1.root;
 }
 
-static bool _upb_Arena_FixupRefs(upb_Arena* new_root, uintptr_t ref_delta) {
+static bool _upb_Arena_FixupRefs(upb_ArenaInternal* new_root,
+                                 uintptr_t ref_delta) {
   if (ref_delta == 0) return true;  // No fixup required.
   uintptr_t poc =
       upb_Atomic_Load(&new_root->parent_or_count, memory_order_relaxed);
@@ -331,38 +434,62 @@ static bool _upb_Arena_FixupRefs(upb_Arena* new_root, uintptr_t ref_delta) {
 bool upb_Arena_Fuse(upb_Arena* a1, upb_Arena* a2) {
   if (a1 == a2) return true;  // trivial fuse
 
+  upb_ArenaInternal* ai1 = upb_Arena_Internal(a1);
+  upb_ArenaInternal* ai2 = upb_Arena_Internal(a2);
+
   // Do not fuse initial blocks since we cannot lifetime extend them.
   // Any other fuse scenario is allowed.
-  if (upb_Arena_HasInitialBlock(a1) || upb_Arena_HasInitialBlock(a2)) {
+  if (_upb_ArenaInternal_HasInitialBlock(ai1) ||
+      _upb_ArenaInternal_HasInitialBlock(ai2)) {
     return false;
   }
 
   // The number of refs we ultimately need to transfer to the new root.
   uintptr_t ref_delta = 0;
   while (true) {
-    upb_Arena* new_root = _upb_Arena_DoFuse(a1, a2, &ref_delta);
+    upb_ArenaInternal* new_root = _upb_Arena_DoFuse(a1, a2, &ref_delta);
     if (new_root != NULL && _upb_Arena_FixupRefs(new_root, ref_delta)) {
       return true;
     }
   }
 }
 
-void upb_Arena_IncRefFor(upb_Arena* arena, const void* owner) {
-  _upb_ArenaRoot r;
+bool upb_Arena_IncRefFor(upb_Arena* a, const void* owner) {
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  if (_upb_ArenaInternal_HasInitialBlock(ai)) return false;
+  upb_ArenaRoot r;
+
 retry:
-  r = _upb_Arena_FindRoot(arena);
+  r = _upb_Arena_FindRoot(a);
   if (upb_Atomic_CompareExchangeWeak(
           &r.root->parent_or_count, &r.tagged_count,
           _upb_Arena_TaggedFromRefcount(
               _upb_Arena_RefCountFromTagged(r.tagged_count) + 1),
           memory_order_release, memory_order_acquire)) {
     // We incremented it successfully, so we are done.
-    return;
+    return true;
   }
   // We failed update due to parent switching on the arena.
   goto retry;
 }
 
-void upb_Arena_DecRefFor(upb_Arena* arena, const void* owner) {
-  upb_Arena_Free(arena);
+void upb_Arena_DecRefFor(upb_Arena* a, const void* owner) { upb_Arena_Free(a); }
+
+void UPB_PRIVATE(_upb_Arena_SwapIn)(upb_Arena* des, const upb_Arena* src) {
+  upb_ArenaInternal* desi = upb_Arena_Internal(des);
+  upb_ArenaInternal* srci = upb_Arena_Internal(src);
+
+  *des = *src;
+  desi->block_alloc = srci->block_alloc;
+  upb_MemBlock* blocks = upb_Atomic_Load(&srci->blocks, memory_order_relaxed);
+  upb_Atomic_Init(&desi->blocks, blocks);
+}
+
+void UPB_PRIVATE(_upb_Arena_SwapOut)(upb_Arena* des, const upb_Arena* src) {
+  upb_ArenaInternal* desi = upb_Arena_Internal(des);
+  upb_ArenaInternal* srci = upb_Arena_Internal(src);
+
+  *des = *src;
+  upb_MemBlock* blocks = upb_Atomic_Load(&srci->blocks, memory_order_relaxed);
+  upb_Atomic_Store(&desi->blocks, blocks, memory_order_relaxed);
 }

data/third_party/upb/upb/mem/arena.h

@@ -22,19 +22,15 @@
 
 #include <stddef.h>
 #include <stdint.h>
-#include <string.h>
 
 #include "upb/mem/alloc.h"
+#include "upb/mem/internal/arena.h"
 
 // Must be last.
 #include "upb/port/def.inc"
 
 typedef struct upb_Arena upb_Arena;
 
-typedef struct {
-  char *ptr, *end;
-} _upb_ArenaHead;
-
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -47,35 +43,23 @@ UPB_API upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc);
 UPB_API void upb_Arena_Free(upb_Arena* a);
 UPB_API bool upb_Arena_Fuse(upb_Arena* a, upb_Arena* b);
 
-void upb_Arena_IncRefFor(upb_Arena* arena, const void* owner);
-void upb_Arena_DecRefFor(upb_Arena* arena, const void* owner);
+bool upb_Arena_IncRefFor(upb_Arena* a, const void* owner);
+void upb_Arena_DecRefFor(upb_Arena* a, const void* owner);
 
-void* _upb_Arena_SlowMalloc(upb_Arena* a, size_t size);
-size_t upb_Arena_SpaceAllocated(upb_Arena* arena);
-uint32_t upb_Arena_DebugRefCount(upb_Arena* arena);
+size_t upb_Arena_SpaceAllocated(upb_Arena* a, size_t* fused_count);
+uint32_t upb_Arena_DebugRefCount(upb_Arena* a);
 
-UPB_INLINE size_t _upb_ArenaHas(upb_Arena* a) {
-  _upb_ArenaHead* h = (_upb_ArenaHead*)a;
-  return (size_t)(h->end - h->ptr);
+UPB_API_INLINE upb_Arena* upb_Arena_New(void) {
+  return upb_Arena_Init(NULL, 0, &upb_alloc_global);
 }
 
-UPB_API_INLINE void* upb_Arena_Malloc(upb_Arena* a, size_t size) {
-  size = UPB_ALIGN_MALLOC(size);
-  size_t span = size + UPB_ASAN_GUARD_SIZE;
-  if (UPB_UNLIKELY(_upb_ArenaHas(a) < span)) {
-    return _upb_Arena_SlowMalloc(a, size);
-  }
-
-  // We have enough space to do a fast malloc.
-  _upb_ArenaHead* h = (_upb_ArenaHead*)a;
-  void* ret = h->ptr;
-  UPB_ASSERT(UPB_ALIGN_MALLOC((uintptr_t)ret) == (uintptr_t)ret);
-  UPB_ASSERT(UPB_ALIGN_MALLOC(size) == size);
-  UPB_UNPOISON_MEMORY_REGION(ret, size);
-
-  h->ptr += span;
+UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size) {
+  return UPB_PRIVATE(_upb_Arena_Malloc)(a, size);
+}
 
-  return ret;
+UPB_API_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize,
+                                       size_t size) {
+  return UPB_PRIVATE(_upb_Arena_Realloc)(a, ptr, oldsize, size);
 }
 
 // Shrinks the last alloc from arena.
@@ -84,43 +68,7 @@ UPB_API_INLINE void* upb_Arena_Malloc(upb_Arena* a, size_t size) {
 // this was not the last alloc.
 UPB_API_INLINE void upb_Arena_ShrinkLast(upb_Arena* a, void* ptr,
                                          size_t oldsize, size_t size) {
-  _upb_ArenaHead* h = (_upb_ArenaHead*)a;
-  oldsize = UPB_ALIGN_MALLOC(oldsize);
-  size = UPB_ALIGN_MALLOC(size);
-  // Must be the last alloc.
-  UPB_ASSERT((char*)ptr + oldsize == h->ptr - UPB_ASAN_GUARD_SIZE);
-  UPB_ASSERT(size <= oldsize);
-  h->ptr = (char*)ptr + size;
-}
-
-UPB_API_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize,
-                                       size_t size) {
-  _upb_ArenaHead* h = (_upb_ArenaHead*)a;
-  oldsize = UPB_ALIGN_MALLOC(oldsize);
-  size = UPB_ALIGN_MALLOC(size);
-  bool is_most_recent_alloc = (uintptr_t)ptr + oldsize == (uintptr_t)h->ptr;
-
-  if (is_most_recent_alloc) {
-    ptrdiff_t diff = size - oldsize;
-    if ((ptrdiff_t)_upb_ArenaHas(a) >= diff) {
-      h->ptr += diff;
-      return ptr;
-    }
-  } else if (size <= oldsize) {
-    return ptr;
-  }
-
-  void* ret = upb_Arena_Malloc(a, size);
-
-  if (ret && oldsize > 0) {
-    memcpy(ret, ptr, UPB_MIN(oldsize, size));
-  }
-
-  return ret;
-}
-
-UPB_API_INLINE upb_Arena* upb_Arena_New(void) {
-  return upb_Arena_Init(NULL, 0, &upb_alloc_global);
+  UPB_PRIVATE(_upb_Arena_ShrinkLast)(a, ptr, oldsize, size);
 }
 
 #ifdef __cplusplus

data/third_party/upb/upb/mem/arena.hpp

@@ -8,6 +8,8 @@
 #ifndef UPB_MEM_ARENA_HPP_
 #define UPB_MEM_ARENA_HPP_
 
+#ifdef __cplusplus
+
 #include <memory>
 
 #include "upb/mem/arena.h"
@@ -30,8 +32,8 @@ class Arena {
   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.
+// 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:
@@ -43,12 +45,14 @@ class InlinedArena : public Arena {
   }
 
  private:
-  InlinedArena(const InlinedArena*) = delete;
-  InlinedArena& operator=(const InlinedArena*) = delete;
+  InlinedArena(const InlinedArena&) = delete;
+  InlinedArena& operator=(const InlinedArena&) = delete;
 
   char initial_block_[N];
 };
 
 }  // namespace upb
 
+#endif  // __cplusplus
+
 #endif  // UPB_MEM_ARENA_HPP_