add push all API

Author: SchrodingerZhu

libc/src/__support/mpmc_stack.h

@@ -29,7 +29,7 @@ template <class T> class MPMCStack {
 public:
   static_assert(cpp::is_copy_constructible<T>::value,
                 "T must be copy constructible");
-  LIBC_INLINE MPMCStack() : head(nullptr) {}
+  LIBC_INLINE constexpr MPMCStack() : head(nullptr) {}
   LIBC_INLINE bool push(T value) {
     AllocChecker ac;
     Node *new_node = new (ac) Node(value);
@@ -41,6 +41,39 @@ template <class T> class MPMCStack {
     });
     return true;
   }
+  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) {}
+      LIBC_INLINE ~Guard() {
+        for (size_t i = 0; i < count; ++i)
+          delete allocated[i];
+      }
+      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];
+    {
+      Guard guard(allocated);
+      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];
+      }
+      guard.clear();
+    }
+    head.transaction([&allocated, count](Node *old_head) {
+      allocated[0]->next = old_head;
+      return allocated[count - 1];
+    });
+    return true;
+  }
   LIBC_INLINE cpp::optional<T> pop() {
     cpp::optional<T> res = cpp::nullopt;
     Node *node = nullptr;

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

@@ -54,8 +54,66 @@ void multithread_test() {
       __builtin_trap();
 }
 
+void multithread_push_all_test() {
+  constexpr static size_t NUM_THREADS = 4;
+  constexpr static size_t BATCH_SIZE = 10;
+  constexpr static size_t NUM_BATCHES = 20;
+  struct State {
+    MPMCStack<size_t> stack;
+    cpp::Atomic<size_t> counter = 0;
+    cpp::Atomic<bool> flags[NUM_THREADS * BATCH_SIZE * NUM_BATCHES];
+  } 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);
+          size_t values[BATCH_SIZE];
+
+          for (size_t batch = 0; batch < NUM_BATCHES; ++batch) {
+            // Prepare batch of values
+            for (size_t j = 0; j < BATCH_SIZE; ++j) {
+              size_t current = state->counter.fetch_add(1);
+              values[j] = current;
+            }
+
+            // Push all values in batch
+            if (!state->stack.push_all(values, BATCH_SIZE))
+              __builtin_trap();
+          }
+
+          // Pop and mark all values
+          while (auto res = state->stack.pop()) {
+            size_t value = res.value();
+            if (value < NUM_THREADS * BATCH_SIZE * NUM_BATCHES)
+              state->flags[value].store(true);
+          }
+          return nullptr;
+        },
+        &state);
+  }
+
+  for (pthread_t thread : threads)
+    LIBC_NAMESPACE::pthread_join(thread, nullptr);
+
+  // Pop any remaining values
+  while (cpp::optional<size_t> res = state.stack.pop()) {
+    size_t value = res.value();
+    if (value < NUM_THREADS * BATCH_SIZE * NUM_BATCHES)
+      state.flags[value].store(true);
+  }
+
+  // Verify all values were processed
+  for (size_t i = 0; i < NUM_THREADS * BATCH_SIZE * NUM_BATCHES; ++i)
+    if (!state.flags[i].load())
+      __builtin_trap();
+}
+
 TEST_MAIN() {
   smoke_test();
   multithread_test();
+  multithread_push_all_test();
   return 0;
 }