[libc-commits] [libc] [libc] make integration test malloc work properly when threaded (PR #151622)

[Schrodinger ZHU Yifan via libc-commits]

https://github.com/SchrodingerZhu updated https://github.com/llvm/llvm-project/pull/151622

>From 0ffe7de92c7302e5a110ec55c56316017d418223 Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <i at zhuyi.fan>
Date: Thu, 31 Jul 2025 21:56:23 -0400
Subject: [PATCH 1/3] [libc] make integration test malloc work properly when
 threaded

---
 libc/test/IntegrationTest/CMakeLists.txt |  1 +
 libc/test/IntegrationTest/test.cpp       | 20 ++++++++++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/libc/test/IntegrationTest/CMakeLists.txt b/libc/test/IntegrationTest/CMakeLists.txt
index 4a999407d48d7..1763ed8b932fe 100644
--- a/libc/test/IntegrationTest/CMakeLists.txt
+++ b/libc/test/IntegrationTest/CMakeLists.txt
@@ -12,5 +12,6 @@ add_object_library(
     test.h
   DEPENDS
     libc.src.__support.OSUtil.osutil
+    libc.src.__support.CPP.atomic
     ${arch_specific_deps}
 )
diff --git a/libc/test/IntegrationTest/test.cpp b/libc/test/IntegrationTest/test.cpp
index 871bdf0dc562b..1a489403f37b5 100644
--- a/libc/test/IntegrationTest/test.cpp
+++ b/libc/test/IntegrationTest/test.cpp
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "src/__support/CPP/atomic.h"
 #include "src/__support/common.h"
 #include "src/__support/macros/config.h"
 #include <stddef.h>
@@ -65,14 +66,25 @@ int atexit(void (*func)(void)) { return LIBC_NAMESPACE::atexit(func); }
 
 static constexpr uint64_t MEMORY_SIZE = 16384;
 static uint8_t memory[MEMORY_SIZE];
-static uint8_t *ptr = memory;
+static LIBC_NAMESPACE::cpp::Atomic<size_t> used = 0;
 
 extern "C" {
 
 void *malloc(size_t s) {
-  void *mem = ptr;
-  ptr += s;
-  return static_cast<uint64_t>(ptr - memory) >= MEMORY_SIZE ? nullptr : mem;
+  constexpr size_t MAX_ALIGNMENT = alignof(long double);
+  if (s > MEMORY_SIZE)
+    return nullptr; // Not enough memory.
+  s = s + MAX_ALIGNMENT -
+      (s % MAX_ALIGNMENT); // Align the size to the max alignment.
+  for (;;) {
+    size_t current_used = used.load();
+    if (current_used + s > MEMORY_SIZE)
+      return nullptr; // Not enough memory.
+
+    // Try to reserve the memory.
+    if (used.compare_exchange_strong(current_used, current_used + s))
+      return &memory[current_used];
+  }
 }
 
 void free(void *) {}

>From 9894fbbd2697f4c8d4d5d61a59f30adda3f93786 Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <i at zhuyi.fan>
Date: Thu, 31 Jul 2025 22:00:58 -0400
Subject: [PATCH 2/3] fix copilot comment

---
 libc/test/IntegrationTest/test.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/libc/test/IntegrationTest/test.cpp b/libc/test/IntegrationTest/test.cpp
index 1a489403f37b5..e754cdf9f5f9b 100644
--- a/libc/test/IntegrationTest/test.cpp
+++ b/libc/test/IntegrationTest/test.cpp
@@ -74,8 +74,10 @@ void *malloc(size_t s) {
   constexpr size_t MAX_ALIGNMENT = alignof(long double);
   if (s > MEMORY_SIZE)
     return nullptr; // Not enough memory.
-  s = s + MAX_ALIGNMENT -
-      (s % MAX_ALIGNMENT); // Align the size to the max alignment.
+  size_t offset = s % MAX_ALIGNMENT;
+  s = offset == 0
+          ? s
+          : s + MAX_ALIGNMENT - offset; // Align the size to the max alignment.
   for (;;) {
     size_t current_used = used.load();
     if (current_used + s > MEMORY_SIZE)

>From e347a01e3220f491ff0700ff916d3a258ca23d66 Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <yifanzhu at rochester.edu>
Date: Fri, 1 Aug 2025 10:09:00 -0400
Subject: [PATCH 3/3] simplify bump allocation

---
 libc/test/IntegrationTest/test.cpp | 23 +++++++----------------
 1 file changed, 7 insertions(+), 16 deletions(-)

diff --git a/libc/test/IntegrationTest/test.cpp b/libc/test/IntegrationTest/test.cpp
index bf9eb1290f5cc..2b8d1b733eb9a 100644
--- a/libc/test/IntegrationTest/test.cpp
+++ b/libc/test/IntegrationTest/test.cpp
@@ -5,8 +5,8 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-#include "src/__support/CPP/atomic.h"
 #include "hdr/stdint_proxy.h"
+#include "src/__support/CPP/atomic.h"
 #include "src/__support/common.h"
 #include "src/__support/macros/config.h"
 #include <stddef.h>
@@ -69,23 +69,14 @@ static LIBC_NAMESPACE::cpp::Atomic<size_t> used = 0;
 
 extern "C" {
 
+// For simple test purposes.
 void *malloc(size_t s) {
   constexpr size_t MAX_ALIGNMENT = alignof(long double);
-  if (s > MEMORY_SIZE)
-    return nullptr; // Not enough memory.
-  size_t offset = s % MAX_ALIGNMENT;
-  s = offset == 0
-          ? s
-          : s + MAX_ALIGNMENT - offset; // Align the size to the max alignment.
-  for (;;) {
-    size_t current_used = used.load();
-    if (current_used + s > MEMORY_SIZE)
-      return nullptr; // Not enough memory.
-
-    // Try to reserve the memory.
-    if (used.compare_exchange_strong(current_used, current_used + s))
-      return &memory[current_used];
-  }
+  s += (-s) & (MAX_ALIGNMENT - 1); // Align to max alignment.
+  auto res = used.fetch_add(s);
+  if (res + s > MEMORY_SIZE)
+    return nullptr; // Out of memory.
+  return &memory[res];
 }
 
 void free(void *) {}