add tests

Author: SchrodingerZhu

libc/src/__support/CMakeLists.txt

@@ -398,6 +398,30 @@ add_header_library(
     libc.src.__support.macros.attributes
 )
 
+add_header_library(
+  aba_ptr
+  HDRS
+    aba_ptr.h
+  DEPENDS
+    libc.hdr.types.size_t
+    libc.src.__support.common
+    libc.src.__support.threads.sleep
+)
+
+add_header_library(
+  mpmc_stack
+  HDRS
+    mpmc_stack.h
+  DEPENDS
+    libc.src.__support.aba_ptr
+    libc.src.__support.common
+    libc.src.__support.CPP.atomic
+    libc.src.__support.CPP.new
+    libc.src.__support.CPP.optional
+    libc.src.__support.CPP.type_traits
+)
+
+
 add_subdirectory(FPUtil)
 add_subdirectory(OSUtil)
 add_subdirectory(StringUtil)

libc/src/__support/aba_ptr.h

@@ -9,6 +9,7 @@
 #ifndef LLVM_LIBC_SRC___SUPPORT_TAGGED_POINTER_H
 #define LLVM_LIBC_SRC___SUPPORT_TAGGED_POINTER_H
 
+#include "hdr/types/size_t.h"
 #include "src/__support/common.h"
 #include "src/__support/threads/sleep.h"
 
@@ -24,7 +25,7 @@ template <class T, bool IsAtomic> struct AbaPtrImpl {
   union Impl {
     struct alignas(2 * alignof(void *)) Atomic {
       T *ptr;
-      __SIZE_TYPE__ tag;
+      size_t tag;
     } atomic;
     struct Mutex {
       T *ptr;
@@ -59,6 +60,20 @@ template <class T, bool IsAtomic> struct AbaPtrImpl {
       __atomic_store_n(&impl.mutex.locked, false, __ATOMIC_RELEASE);
     }
   }
+
+  LIBC_INLINE T *get() const {
+    if constexpr (IsAtomic) {
+      // Weak micro-architectures typically reguards simultaneous partial word
+      // loading and full word loading as a race condition. While there are
+      // implementations that uses racy read anyway, we still load the whole
+      // word to avoid any complications.
+      typename Impl::Atomic snapshot;
+      __atomic_load(&impl.atomic, &snapshot, __ATOMIC_RELAXED);
+      return snapshot.ptr;
+    } else {
+      return impl.mutex.ptr;
+    }
+  }
 };
 
 template <class T> using AbaPtr = AbaPtrImpl<T, LIBC_ABA_PTR_IS_ATOMIC>;

libc/src/__support/mpmc_stack.h

@@ -12,13 +12,14 @@
 #include "src/__support/CPP/atomic.h"
 #include "src/__support/CPP/new.h"
 #include "src/__support/CPP/optional.h"
+#include "src/__support/CPP/type_traits.h"
 #include "src/__support/aba_ptr.h"
 
 namespace LIBC_NAMESPACE_DECL {
 template <class T> class MPMCStack {
   struct Node {
     cpp::Atomic<size_t> visitor;
-    AbaPtr<Node> next;
+    Node *next;
     T value;
 
     LIBC_INLINE Node(T val) : visitor(0), next(nullptr), value(val) {}
@@ -31,7 +32,7 @@ template <class T> class MPMCStack {
   LIBC_INLINE MPMCStack() : head(nullptr) {}
   LIBC_INLINE bool push(T value) {
     AllocChecker ac;
-    Node *new_node = new Node(value, ac);
+    Node *new_node = new (ac) Node(value);
     if (!ac)
       return false;
     head.transaction([new_node](Node *old_head) {
@@ -41,17 +42,17 @@ template <class T> class MPMCStack {
     return true;
   }
   LIBC_INLINE cpp::optional<T> pop() {
-    cpp::optional<T> res;
-    Node *node;
-    head.transaction([&](Node *current_head) {
+    cpp::optional<T> res = cpp::nullopt;
+    Node *node = nullptr;
+    head.transaction([&](Node *current_head) -> Node * {
       if (!current_head) {
         res = cpp::nullopt;
         return nullptr;
       }
       node = current_head;
       node->visitor.fetch_add(1);
-      res = node->value;
-      auto next = node->next;
+      res = cpp::optional<T>{node->value};
+      Node *next = node->next;
       node->visitor.fetch_sub(1);
       return next;
     });

libc/test/integration/src/__support/CMakeLists.txt

@@ -2,3 +2,18 @@ add_subdirectory(threads)
 if(LIBC_TARGET_OS_IS_GPU)
   add_subdirectory(GPU)
 endif()
+
+add_libc_integration_test_suite(libc-support-integration-tests)
+
+add_integration_test(
+  mpmc_stack_test
+  SUITE
+    libc-support-integration-tests
+  SRCS
+    mpmc_stack_test.cpp
+  DEPENDS
+    libc.src.__support.mpmc_stack
+    libc.src.__support.threads.thread
+    libc.src.pthread.pthread_create
+    libc.src.pthread.pthread_join
+)

libc/test/integration/src/__support/mpmc_stack_test.cpp

@@ -0,0 +1,61 @@
+#include "src/__support/CPP/atomic.h"
+#include "src/__support/mpmc_stack.h"
+#include "src/pthread/pthread_create.h"
+#include "src/pthread/pthread_join.h"
+#include "test/IntegrationTest/test.h"
+
+using namespace LIBC_NAMESPACE;
+
+void smoke_test() {
+  MPMCStack<int> stack;
+  for (int i = 0; i <= 100; ++i)
+    if (!stack.push(i))
+      __builtin_trap();
+  for (int i = 100; i >= 0; --i)
+    if (*stack.pop() != i)
+      __builtin_trap();
+  if (stack.pop())
+    __builtin_trap(); // Should be empty now.
+}
+
+void multithread_test() {
+  constexpr static size_t NUM_THREADS = 5;
+  constexpr static size_t NUM_PUSHES = 100;
+  struct State {
+    MPMCStack<size_t> stack;
+    cpp::Atomic<size_t> counter = 0;
+    cpp::Atomic<bool> flags[NUM_PUSHES];
+  } state;
+  pthread_t threads[NUM_THREADS];
+  for (size_t i = 0; i < NUM_THREADS; ++i) {
+    LIBC_NAMESPACE::pthread_create(
+        &threads[i], nullptr,
+        [](void *arg) -> void * {
+          State *state = static_cast<State *>(arg);
+          for (;;) {
+            size_t current = state->counter.fetch_add(1);
+            if (current >= NUM_PUSHES)
+              break;
+            if (!state->stack.push(current))
+              __builtin_trap();
+          }
+          while (auto res = state->stack.pop())
+            state->flags[res.value()].store(true);
+          return nullptr;
+        },
+        &state);
+  }
+  for (pthread_t thread : threads)
+    LIBC_NAMESPACE::pthread_join(thread, nullptr);
+  while (cpp::optional<size_t> res = state.stack.pop())
+    state.flags[res.value()].store(true);
+  for (size_t i = 0; i < NUM_PUSHES; ++i)
+    if (!state.flags[i].load())
+      __builtin_trap();
+}
+
+TEST_MAIN() {
+  smoke_test();
+  multithread_test();
+  return 0;
+}