2 files changed, 5 insertions, 3 deletions
diff --git a/llvm/unittests/Support/AllocatorTest.cpp b/llvm/unittests/Support/AllocatorTest.cpp
index 4897c47eb28..74b394f1b17 100644
--- a/llvm/unittests/Support/AllocatorTest.cpp
+++ b/llvm/unittests/Support/AllocatorTest.cpp
@@ -147,7 +147,7 @@ public:
     // Allocate space for the alignment, the slab, and a void* that goes right
     // before the slab.
     size_t Alignment = 4096;
-    void *MemBase = malloc(Size + Alignment - 1 + sizeof(void*));
+    void *MemBase = safe_malloc(Size + Alignment - 1 + sizeof(void*));
 
     // Find the slab start.
     void *Slab = (void *)alignAddr((char*)MemBase + sizeof(void *), Alignment);
diff --git a/llvm/unittests/Support/ManagedStatic.cpp b/llvm/unittests/Support/ManagedStatic.cpp
index 07e324cdfb6..d3cc80cf5e9 100644
--- a/llvm/unittests/Support/ManagedStatic.cpp
+++ b/llvm/unittests/Support/ManagedStatic.cpp
@@ -6,6 +6,8 @@
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Config/config.h"
 #ifdef HAVE_PTHREAD_H
@@ -30,7 +32,7 @@ namespace test1 {
   // Valgrind's leak checker complains glibc's stack allocation.
   // To appease valgrind, we provide our own stack for each thread.
   void *allocate_stack(pthread_attr_t &a, size_t n = 65536) {
-    void *stack = malloc(n);
+    void *stack = safe_malloc(n);
     pthread_attr_init(&a);
 #if defined(__linux__)
     pthread_attr_setstack(&a, stack, n);
@@ -83,7 +85,7 @@ TEST(ManagedStaticTest, NestedStatics) {
 namespace CustomCreatorDeletor {
 struct CustomCreate {
   static void *call() {
-    void *Mem = std::malloc(sizeof(int));
+    void *Mem = safe_malloc(sizeof(int));
     *((int *)Mem) = 42;
     return Mem;
   }