Google Git
Sign in
chromium/chromium/src/base/allocator/partition_allocator/refs/heads/main^!/.
commit
407e9f8b2be6bf2e738c2a87806e10ca4e17e909
[log]
author
Arthur Sonzogni <arthursonzogni@chromium.org>
Fri Jul 10 17:00:06 2026
committer
Copybara-Service <copybara-worker@google.com>
Fri Jul 10 17:01:55 2026
tree
8dc2c7098cf0ba2aaf0461cc085c23bd15a431d2
parent
f36d5a2c8460f1c82885165601a942717a39c0f0 [diff]
partition_alloc: Migrate MEMORY_TOOL_REPLACES_ALLOCATOR to PA buildflag

Move MEMORY_TOOL_REPLACES_ALLOCATOR from a global //build define to a
local partition_alloc buildflag.

This enables PartitionAlloc to compile standalone when embedded in other
projects without depending on //build. Using a buildflag also prevents
silent errors if the defining config is missing.

We double-checked that MEMORY_TOOL_REPLACES_ALLOCATOR is not used in
build embedders: v8, Dawn, ANGLE, PDFium, WebRTC, and SwiftShader.

Additionally, the legacy macro has been banned in PartitionAlloc by
removing it from the PRESUBMIT.py allowed list.

Bug: 351867706
Change-Id: I2fc46a66ab4ff4ccfae36a057e7977a3e3b21f48
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7871106
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Reviewed-by: Hans Wennborg <hans@chromium.org>
Reviewed-by: Titouan Rigoudy <titouan@chromium.org>
Auto-Submit: Arthur Sonzogni <arthursonzogni@chromium.org>
Reviewed-by: Stephen Nusko <nuskos@chromium.org>
Commit-Queue: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1660347}
NOKEYCHECK=True
GitOrigin-RevId: 1c41747a9ae220994d4a97bcd9b29d6a85d0949f

diff --git a/PRESUBMIT.py b/PRESUBMIT.py
index 8726217..3f34239 100644
--- a/PRESUBMIT.py
+++ b/PRESUBMIT.py

@@ -361,12 +361,9 @@
         'HAS_HW_CAPS',
         'HAVE_BACKTRACE',
         'LINUX_NAME_REGION',
-        'MEMORY_TOOL_REPLACES_ALLOCATOR',
         'NDEBUG',
         'NEEDS_HANDLING_OF_HW_CAPABILITIES',
     }
-
-
     target_path_prefix = (
         'base/allocator/partition_allocator/src/partition_alloc/')
 

diff --git a/src/partition_alloc/BUILD.gn b/src/partition_alloc/BUILD.gn
index eea905b..c25f2d5 100644
--- a/src/partition_alloc/BUILD.gn
+++ b/src/partition_alloc/BUILD.gn

@@ -140,6 +140,10 @@
   partition_alloc_realloc_growth_factor_mitigation = false
 }
 
+# Sanitizers replace the allocator, don't use our own.
+memory_tool_replaces_allocator =
+    is_asan || is_hwasan || is_lsan || is_tsan || is_msan
+
 pa_buildflag_header("buildflags") {
   header = "buildflags.h"
   flags = [
@@ -182,6 +186,7 @@
     "IS_COMPONENT_BUILD=$is_component_build",
     "IS_DEBUG=$partition_alloc_is_debug",
     "IS_HWASAN=$is_hwasan",
+    "MEMORY_TOOL_REPLACES_ALLOCATOR=$memory_tool_replaces_allocator",
     "OFFICIAL=$partition_alloc_is_official_build",
     "RAW_PTR_ZERO_ON_CONSTRUCT=$raw_ptr_zero_on_construct",
     "RAW_PTR_ZERO_ON_DESTRUCT=$raw_ptr_zero_on_destruct",

diff --git a/src/partition_alloc/address_space_randomization.h b/src/partition_alloc/address_space_randomization.h
index 7b27d8e..2042dba 100644
--- a/src/partition_alloc/address_space_randomization.h
+++ b/src/partition_alloc/address_space_randomization.h

@@ -38,7 +38,7 @@
 
 #if PA_BUILDFLAG(PA_ARCH_CPU_64_BITS)
 
-  #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+  #if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
     // We shouldn't allocate system pages at all for sanitizer builds. However,
     // we do, and if random hint addresses interfere with address ranges

diff --git a/src/partition_alloc/hardening_unittest.cc b/src/partition_alloc/hardening_unittest.cc
index cdceced..2b12b63 100644
--- a/src/partition_alloc/hardening_unittest.cc
+++ b/src/partition_alloc/hardening_unittest.cc

@@ -20,7 +20,7 @@
 #include "testing/gtest/include/gtest/gtest.h"
 
 // With *SAN, PartitionAlloc is rerouted to malloc().
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 namespace partition_alloc::internal {
 namespace {
@@ -248,4 +248,4 @@
 }  // namespace
 }  // namespace partition_alloc::internal
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)

diff --git a/src/partition_alloc/internal/partition_root_internal.h b/src/partition_alloc/internal/partition_root_internal.h
index 5c5dae6..9215733 100644
--- a/src/partition_alloc/internal/partition_root_internal.h
+++ b/src/partition_alloc/internal/partition_root_internal.h

@@ -43,13 +43,13 @@
 
 // When a memory tool is replacing malloc to keep aligned behaviour working we
 // use window's aligned_malloc and aligned_free, but otherwise we need memalign.
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #if PA_BUILDFLAG(PA_COMPILER_MSVC)
 #include <malloc.h>
 #else
 #include <stdlib.h>
 #endif  // PA_BUILDFLAG(PA_COMPILER_MSVC)
-#endif  // defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 namespace partition_alloc {
 
@@ -414,7 +414,7 @@
 PA_ALWAYS_INLINE bool PartitionRoot::FreeProlog(void* object,
                                                 const PartitionRoot* root) {
   static_assert(AreValidFlags(flags));
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
   if constexpr (!ContainsFlags(flags, FreeFlags::kNoMemoryToolOverride)) {
 #if PA_BUILDFLAG(PA_COMPILER_MSVC)
     if (ContainsFlags(flags, FreeFlags::kAlignedFreeForMemoryTool)) {
@@ -427,11 +427,11 @@
 #endif  // PA_BUILDFLAG(PA_COMPILER_MSVC)
     return true;
   }
-#else   // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#else   // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
   // If the memory tool is not replacing the allocator, then the
   // kAlignedFreeForMemoryTool flag is unused and should not be passed.
   static_assert(!ContainsFlags(flags, FreeFlags::kAlignedFreeForMemoryTool));
-#endif  // defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
   if (!object) [[unlikely]] {
     return true;
@@ -1275,7 +1275,7 @@
   static_assert(!ContainsFlags(
       flags, AllocFlags::kMemoryShouldBeTaggedForMte));  // Internal only.
 
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
   if constexpr (!ContainsFlags(flags, AllocFlags::kNoMemoryToolOverride)) {
     if (!PartitionRoot::AllocWithMemoryToolProlog<flags>(requested_size)) {
       // Early return if AllocWithMemoryToolProlog returns false
@@ -1317,7 +1317,7 @@
     }
     return result;
   }
-#endif  // defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
   constexpr bool no_hooks = ContainsFlags(flags, AllocFlags::kNoHooks);
   bool hooks_enabled;
@@ -1652,7 +1652,7 @@
                                    size_t new_size,
                                    const char* type_name) {
   static_assert(!ContainsFlags(alloc_flags, AllocFlags::kAlignedAlloc));
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
   if (!PartitionRoot::AllocWithMemoryToolProlog<alloc_flags>(new_size)) {
     // Early return if AllocWithMemoryToolProlog returns false
     return nullptr;

diff --git a/src/partition_alloc/memory_reclaimer_unittest.cc b/src/partition_alloc/memory_reclaimer_unittest.cc
index 3dd2aa0..a4fc13b 100644
--- a/src/partition_alloc/memory_reclaimer_unittest.cc
+++ b/src/partition_alloc/memory_reclaimer_unittest.cc

@@ -25,7 +25,7 @@
 
 // Otherwise, PartitionAlloc doesn't allocate any memory, and the tests are
 // meaningless.
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 namespace partition_alloc {
 
@@ -150,4 +150,4 @@
 
 }  // namespace partition_alloc
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)

diff --git a/src/partition_alloc/page_allocator_unittest.cc b/src/partition_alloc/page_allocator_unittest.cc
index d56a409..bf7045b 100644
--- a/src/partition_alloc/page_allocator_unittest.cc
+++ b/src/partition_alloc/page_allocator_unittest.cc

@@ -46,7 +46,7 @@
 #endif
 #endif  // PA_BUILDFLAG(HAS_MEMORY_TAGGING)
 
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 #if PA_BUILDFLAG(IS_IOS)
 #include <sys/sysctl.h>
@@ -760,4 +760,4 @@
 
 }  // namespace partition_alloc::internal
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)

diff --git a/src/partition_alloc/partition_alloc_unittest.cc b/src/partition_alloc/partition_alloc_unittest.cc
index 5cdd8f2..26a03c3 100644
--- a/src/partition_alloc/partition_alloc_unittest.cc
+++ b/src/partition_alloc/partition_alloc_unittest.cc

@@ -6,7 +6,7 @@
 
 #include "partition_alloc/bucket_lookup.h"
 #include "partition_alloc/slot_start.h"
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 #include <algorithm>
 #include <cstddef>
@@ -6632,4 +6632,4 @@
 }
 }  // namespace partition_alloc::internal
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)

diff --git a/src/partition_alloc/partition_root.cc b/src/partition_alloc/partition_root.cc
index f12dc32..6a613a3 100644
--- a/src/partition_alloc/partition_root.cc
+++ b/src/partition_alloc/partition_root.cc

@@ -1445,7 +1445,7 @@
 // returned value, it'd use the same amount of underlying memory as the
 // allocation with |size|.
 size_t PartitionRoot::AllocationCapacityFromRequestedSize(size_t size) const {
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
   return size;
 #else
   PA_DCHECK(PartitionRoot::initialized_);
@@ -2114,11 +2114,11 @@
   // Normally kAlignedFree is a no-op call into Free, but with memory tools it
   // will instead remap to the appropriate system aligned free call.
   constexpr FreeFlags kMaybeAlignedFreeForMemoryTool =
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
       FreeFlags::kAlignedFreeForMemoryTool;
 #else
       FreeFlags::kNone;
-#endif  // defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
   FreeInline<flags | kMaybeAlignedFreeForMemoryTool>(object);
 }
 

diff --git a/src/partition_alloc/partition_root.h b/src/partition_alloc/partition_root.h
index 28e7c21..b778b44 100644
--- a/src/partition_alloc/partition_root.h
+++ b/src/partition_alloc/partition_root.h

@@ -56,13 +56,13 @@
 
 // When a memory tool is replacing malloc to keep aligned behaviour working we
 // use window's aligned_malloc and aligned_free, but otherwise we need memalign.
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #if PA_BUILDFLAG(PA_COMPILER_MSVC)
 #include <malloc.h>
 #else
 #include <stdlib.h>
 #endif  // PA_BUILDFLAG(PA_COMPILER_MSVC)
-#endif  // defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 namespace partition_alloc::internal {
 

diff --git a/src/partition_alloc/pointers/raw_ptr_unittest.cc b/src/partition_alloc/pointers/raw_ptr_unittest.cc
index c814faf..dcec891 100644
--- a/src/partition_alloc/pointers/raw_ptr_unittest.cc
+++ b/src/partition_alloc/pointers/raw_ptr_unittest.cc

@@ -1631,7 +1631,7 @@
 namespace base::internal {
 
 #if PA_BUILDFLAG(USE_RAW_PTR_BACKUP_REF_IMPL) && \
-    !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+    !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 void HandleOOM(size_t unused_size) {
   PA_LOG(FATAL) << "Out of memory";
@@ -2460,7 +2460,7 @@
 }
 
 #endif  // PA_BUILDFLAG(USE_RAW_PTR_BACKUP_REF_IMPL) &&
-        // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+        // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 #if PA_BUILDFLAG(USE_RAW_PTR_HOOKABLE_IMPL)
 

diff --git a/src/partition_alloc/scheduler_loop_quarantine_unittest.cc b/src/partition_alloc/scheduler_loop_quarantine_unittest.cc
index 2fb29bb..93f601e 100644
--- a/src/partition_alloc/scheduler_loop_quarantine_unittest.cc
+++ b/src/partition_alloc/scheduler_loop_quarantine_unittest.cc

@@ -17,7 +17,7 @@
 
 namespace partition_alloc {
 
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 namespace {
 
@@ -318,6 +318,6 @@
 
 }  // namespace
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 }  // namespace partition_alloc

diff --git a/src/partition_alloc/shim/allocator_interception_apple_unittest.mm b/src/partition_alloc/shim/allocator_interception_apple_unittest.mm
index 948b72f..a58ee83 100644
--- a/src/partition_alloc/shim/allocator_interception_apple_unittest.mm
+++ b/src/partition_alloc/shim/allocator_interception_apple_unittest.mm

@@ -50,7 +50,7 @@
   }
 };
 
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 TEST_F(AllocatorInterceptionTest, ShimNewMallocZones) {
   InitializeAllocatorShim();
   ChromeMallocZone* default_malloc_zone =

diff --git a/src/partition_alloc/shim/allocator_shim.cc b/src/partition_alloc/shim/allocator_shim.cc
index b4bcb3c..3ca168d 100644
--- a/src/partition_alloc/shim/allocator_shim.cc
+++ b/src/partition_alloc/shim/allocator_shim.cc

@@ -56,7 +56,7 @@
 
 // Cross-checks.
 
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #error The allocator shim should not be compiled when building for memory tools.
 #endif
 

diff --git a/src/partition_alloc/shim/allocator_shim_android.cc b/src/partition_alloc/shim/allocator_shim_android.cc
index 26f2db6..48cbcb7 100644
--- a/src/partition_alloc/shim/allocator_shim_android.cc
+++ b/src/partition_alloc/shim/allocator_shim_android.cc

@@ -41,7 +41,7 @@
 
 // Cross-checks.
 
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #error The allocator shim should not be compiled when building for memory tools.
 #endif
 

diff --git a/src/partition_alloc/shim/allocator_shim_apple.cc b/src/partition_alloc/shim/allocator_shim_apple.cc
index f812a9b..5766259 100644
--- a/src/partition_alloc/shim/allocator_shim_apple.cc
+++ b/src/partition_alloc/shim/allocator_shim_apple.cc

@@ -109,7 +109,7 @@
 
 // Cross-checks.
 
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #error The allocator shim should not be compiled when building for memory tools.
 #endif
 

diff --git a/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.cc b/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.cc
index f62baf2..ad4e069 100644
--- a/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.cc
+++ b/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.cc

@@ -510,7 +510,7 @@
     void* object,
     void* context) {
   constexpr partition_alloc::FreeFlags kMaybeAlignedFreeForMemoryTool =
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
       partition_alloc::FreeFlags::kAlignedFreeForMemoryTool;
 #else
       partition_alloc::FreeFlags::kNone;

diff --git a/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc_unittest.cc b/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc_unittest.cc
index 17d0b28..b1d3e43 100644
--- a/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc_unittest.cc
+++ b/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc_unittest.cc

@@ -22,7 +22,7 @@
 #include <malloc/malloc.h>
 #endif
 
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR) && \
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR) && \
     PA_BUILDFLAG(USE_PARTITION_ALLOC)
 
 namespace allocator_shim::internal {
@@ -242,5 +242,5 @@
 #endif  // PA_BUILDFLAG(IS_APPLE)
 
 }  // namespace allocator_shim::internal
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR) &&
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR) &&
         // PA_BUILDFLAG(USE_PARTITION_ALLOC)

diff --git a/src/partition_alloc/shim/allocator_shim_win_component.cc b/src/partition_alloc/shim/allocator_shim_win_component.cc
index 59a675b..e0706ae 100644
--- a/src/partition_alloc/shim/allocator_shim_win_component.cc
+++ b/src/partition_alloc/shim/allocator_shim_win_component.cc

@@ -41,7 +41,7 @@
 #error This code is only for Windows component build.
 #endif
 
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #error The allocator shim should not be compiled when building for memory tools.
 #endif
 

diff --git a/src/partition_alloc/shim/allocator_shim_win_static.cc b/src/partition_alloc/shim/allocator_shim_win_static.cc
index 9596826..e2fa7aa 100644
--- a/src/partition_alloc/shim/allocator_shim_win_static.cc
+++ b/src/partition_alloc/shim/allocator_shim_win_static.cc

@@ -26,7 +26,7 @@
 #error This code is only for Windows static build.
 #endif
 
-#if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 #error The allocator shim should not be compiled when building for memory tools.
 #endif
 

diff --git a/src/partition_alloc/slot_start_unittest.cc b/src/partition_alloc/slot_start_unittest.cc
index c6662ee..86dfb64 100644
--- a/src/partition_alloc/slot_start_unittest.cc
+++ b/src/partition_alloc/slot_start_unittest.cc

@@ -11,7 +11,7 @@
 // These tests are default-disabled when PA passes through to a
 // sanitizer, in which case the values returned from `Alloc()` are not
 // managed by PartitionAlloc.
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 namespace partition_alloc::internal {
 namespace {
@@ -77,4 +77,4 @@
 }  // namespace
 }  // namespace partition_alloc::internal
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)

diff --git a/src/partition_alloc/stack/stack_unittest.cc b/src/partition_alloc/stack/stack_unittest.cc
index adcb216..ff3290c 100644
--- a/src/partition_alloc/stack/stack_unittest.cc
+++ b/src/partition_alloc/stack/stack_unittest.cc

@@ -11,7 +11,7 @@
 #include "partition_alloc/partition_alloc_base/compiler_specific.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)
 
 #if PA_BUILDFLAG(IS_LINUX) && \
     (PA_BUILDFLAG(PA_ARCH_CPU_X86) || PA_BUILDFLAG(PA_ARCH_CPU_X86_64))
@@ -351,4 +351,4 @@
 
 }  // namespace partition_alloc::internal
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR)

diff --git a/src/partition_alloc/thread_cache_unittest.cc b/src/partition_alloc/thread_cache_unittest.cc
index 618bf57..304adb5 100644
--- a/src/partition_alloc/thread_cache_unittest.cc
+++ b/src/partition_alloc/thread_cache_unittest.cc

@@ -28,7 +28,7 @@
 // cannot test the thread cache.
 //
 // Finally, the thread cache is not supported on all platforms.
-#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR) && \
+#if !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR) && \
     PA_CONFIG(THREAD_CACHE_SUPPORTED)
 
 namespace partition_alloc::internal {
@@ -1424,5 +1424,5 @@
 
 }  // namespace partition_alloc::internal
 
-#endif  // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR) &&
+#endif  // !PA_BUILDFLAG(MEMORY_TOOL_REPLACES_ALLOCATOR) &&
         // PA_CONFIG(THREAD_CACHE_SUPPORTED)