diff --git a/upb/mem/arena.c b/upb/mem/arena.c index a01421e0d7..e26a0972e0 100644 --- a/upb/mem/arena.c +++ b/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, diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc index 444bd0d901..a7889a069b 100644 --- a/upb/mem/arena_test.cc +++ b/upb/mem/arena_test.cc @@ -7,20 +7,22 @@ #include "upb/mem/arena.h" +#include + #include #include #include #include -#include #include #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) { diff --git a/upb/mem/internal/arena.h b/upb/mem/internal/arena.h index cc97be54de..9a6469edc1 100644 --- a/upb/mem/internal/arena.h +++ b/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)); }