upb: implement _upb_Arena_Contains()

PiperOrigin-RevId: 606262924

upb/mem/arena.c

@@ -171,6 +171,21 @@ size_t upb_Arena_SpaceAllocated(upb_Arena* arena) {
   return memsize;
 }
 
+bool UPB_PRIVATE(_upb_Arena_Contains)(const upb_Arena* a, void* ptr) {
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  UPB_ASSERT(ai);
+
+  upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_relaxed);
+  while (block) {
+    uintptr_t beg = (uintptr_t)block;
+    uintptr_t end = beg + block->size;
+    if ((uintptr_t)ptr >= beg && (uintptr_t)ptr < end) return true;
+    block = upb_Atomic_Load(&block->next, memory_order_relaxed);
+  }
+
+  return false;
+}
+
 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,

upb/mem/arena_test.cc

@@ -7,20 +7,22 @@
 
 #include "upb/mem/arena.h"
 
+#include <stddef.h>
+
 #include <array>
 #include <atomic>
 #include <thread>
 #include <vector>
 
-#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/random/distributions.h"
 #include "absl/random/random.h"
 #include "absl/synchronization/notification.h"
-
-// Must be last.
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
+#include "upb/mem/alloc.h"
+
+// Must be last.
 #include "upb/port/def.inc"
 
 namespace {
@@ -35,7 +37,7 @@ TEST(ArenaTest, ArenaFuse) {
   upb_Arena_Free(arena2);
 }
 
-/* Do nothing allocator for testing */
+// Do-nothing allocator for testing.
 extern "C" void* TestAllocFunc(upb_alloc* alloc, void* ptr, size_t oldsize,
                                size_t size) {
   return upb_alloc_global.func(alloc, ptr, oldsize, size);
@@ -131,6 +133,44 @@ TEST(ArenaTest, FuzzSingleThreaded) {
   }
 }
 
+TEST(ArenaTest, Contains) {
+  upb_Arena* arena1 = upb_Arena_New();
+  upb_Arena* arena2 = upb_Arena_New();
+  void* ptr1a = upb_Arena_Malloc(arena1, 8);
+  void* ptr2a = upb_Arena_Malloc(arena2, 8);
+
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
+
+  void* ptr1b = upb_Arena_Malloc(arena1, 1000000);
+  void* ptr2b = upb_Arena_Malloc(arena2, 1000000);
+
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1b));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2b));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2b));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1b));
+
+  upb_Arena_Fuse(arena1, arena2);
+
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1b));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
+  EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2b));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2b));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
+  EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1b));
+
+  upb_Arena_Free(arena1);
+  upb_Arena_Free(arena2);
+}
+
 #ifdef UPB_USE_C11_ATOMICS
 
 TEST(ArenaTest, FuzzFuseFreeRace) {

upb/mem/internal/arena.h

@@ -40,6 +40,11 @@ void UPB_PRIVATE(_upb_Arena_SwapIn)(struct upb_Arena* des,
 void UPB_PRIVATE(_upb_Arena_SwapOut)(struct upb_Arena* des,
                                      const struct upb_Arena* src);
 
+// Returns whether |ptr| was allocated directly by |a| (so care must be used
+// with fused arenas).
+UPB_API bool UPB_ONLYBITS(_upb_Arena_Contains)(const struct upb_Arena* a,
+                                               void* ptr);
+
 UPB_INLINE size_t UPB_PRIVATE(_upb_ArenaHas)(const struct upb_Arena* a) {
   return (size_t)(a->UPB_ONLYBITS(end) - a->UPB_ONLYBITS(ptr));
 }