address jhuber6's comments

Author: SchrodingerZhu

libc/src/__support/CMakeLists.txt

@@ -405,6 +405,7 @@ add_header_library(
   DEPENDS
     libc.hdr.types.size_t
     libc.src.__support.common
+    libc.src.__support.CPP.atomic
     libc.src.__support.threads.sleep
 )
 

libc/src/__support/aba_ptr.h

@@ -10,6 +10,7 @@
 #define LLVM_LIBC_SRC___SUPPORT_TAGGED_POINTER_H
 
 #include "hdr/types/size_t.h"
+#include "src/__support/CPP/atomic.h"
 #include "src/__support/common.h"
 #include "src/__support/threads/sleep.h"
 
@@ -40,24 +41,28 @@ template <class T, bool IsAtomic> struct AbaPtrImpl {
   template <class Op> LIBC_INLINE void transaction(Op &&op) {
     if constexpr (IsAtomic) {
       for (;;) {
+        cpp::AtomicRef<typename Impl::Atomic> ref(impl.atomic);
         typename Impl::Atomic snapshot, next;
-        __atomic_load(&impl.atomic, &snapshot, __ATOMIC_RELAXED);
+        snapshot = ref.load(cpp::MemoryOrder::RELAXED);
         next.ptr = op(snapshot.ptr);
         // Wrapping add for unsigned integers.
         next.tag = snapshot.tag + 1;
-        if (__atomic_compare_exchange(&impl.atomic, &snapshot, &next, true,
-                                      __ATOMIC_ACQ_REL, __ATOMIC_RELAXED))
+        // Redo transaction can be costly, so we use strong version.
+        if (ref.compare_exchange_strong(snapshot, next,
+                                        cpp::MemoryOrder::ACQ_REL,
+                                        cpp::MemoryOrder::RELAXED))
           return;
       }
     } else {
       // Acquire the lock.
-      while (__atomic_exchange_n(&impl.mutex.locked, true, __ATOMIC_ACQUIRE))
-        while (__atomic_load_n(&impl.mutex.locked, __ATOMIC_RELAXED))
+      cpp::AtomicRef<bool> ref(impl.mutex.locked);
+      while (ref.exchange(true, cpp::MemoryOrder::ACQUIRE))
+        while (ref.load(cpp::MemoryOrder::RELAXED))
           LIBC_NAMESPACE::sleep_briefly();
 
       impl.mutex.ptr = op(impl.mutex.ptr);
       // Release the lock.
-      __atomic_store_n(&impl.mutex.locked, false, __ATOMIC_RELEASE);
+      ref.store(false, cpp::MemoryOrder::RELEASE);
     }
   }
 
@@ -68,7 +73,8 @@ template <class T, bool IsAtomic> struct AbaPtrImpl {
       // 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);
+      cpp::AtomicRef<typename Impl::Atomic> ref(impl.atomic);
+      snapshot = ref.load(cpp::MemoryOrder::RELAXED);
       return snapshot.ptr;
     } else {
       return impl.mutex.ptr;

libc/src/__support/mpmc_stack.h

@@ -42,34 +42,43 @@ template <class T> class MPMCStack {
   }
   LIBC_INLINE bool push_all(T values[], size_t count) {
     struct Guard {
-      size_t count;
-      Node **allocated;
-      LIBC_INLINE Guard(Node *allocated[]) : count(0), allocated(allocated) {}
+      Node *cursor;
+      LIBC_INLINE Guard() : cursor(nullptr) {}
       LIBC_INLINE ~Guard() {
-        for (size_t i = 0; i < count; ++i)
-          delete allocated[i];
+        while (cursor) {
+          Node *next = cursor->next;
+          delete cursor;
+          cursor = next;
+        }
+      }
+      LIBC_INLINE void advance(Node *node) {
+        node->next = cursor;
+        cursor = node;
+      }
+      LIBC_INLINE Node *finish() {
+        Node *ret = cursor;
+        cursor = nullptr;
+        return ret;
       }
-      LIBC_INLINE void add(Node *node) { allocated[count++] = node; }
-      LIBC_INLINE void clear() { count = 0; }
     };
-    // Variable sized array is a GNU extension.
-    __extension__ Node *allocated[count];
+    Node *first = nullptr;
+    Node *last = nullptr;
     {
-      Guard guard(allocated);
+      Guard guard{};
       for (size_t i = 0; i < count; ++i) {
         AllocChecker ac;
         Node *new_node = new (ac) Node(values[i]);
         if (!ac)
           return false;
-        guard.add(new_node);
-        if (i != 0)
-          new_node->next = allocated[i - 1];
+        if (i == 0)
+          first = new_node;
+        guard.advance(new_node);
       }
-      guard.clear();
+      last = guard.finish();
     }
-    head.transaction([&allocated, count](Node *old_head) {
-      allocated[0]->next = old_head;
-      return allocated[count - 1];
+    head.transaction([first, last](Node *old_head) {
+      first->next = old_head;
+      return last;
     });
     return true;
   }
@@ -82,10 +91,10 @@ template <class T> class MPMCStack {
         return nullptr;
       }
       node = current_head;
-      node->visitor.fetch_add(1);
+      node->visitor.fetch_add(1, cpp::MemoryOrder::ACQUIRE);
       res = cpp::optional<T>{node->value};
       Node *next = node->next;
-      node->visitor.fetch_sub(1);
+      node->visitor.fetch_sub(1, cpp::MemoryOrder::RELEASE);
       return next;
     });
     // On a successful transaction, a node is popped by us. So we must delete
@@ -95,7 +104,7 @@ template <class T> class MPMCStack {
     // node.
     if (res) {
       // Spin until the node is no longer in use.
-      while (node->visitor.load() != 0)
+      while (node->visitor.load(cpp::MemoryOrder::RELAXED) != 0)
         LIBC_NAMESPACE::sleep_briefly();
       delete node;
     }