[concurrency] Implement the Task allocator as bump-pointer allocator.

Use the StackAllocator as task allocator.

TODO: we could pass an initial pre-allocated first slab to the allocator, which is allocated on the stack or with the parent task's allocator.

rdar://problem/71157018

Author: eeckstein

stdlib/public/Concurrency/Task.cpp

@@ -311,6 +311,8 @@ AsyncTaskAndContext swift::swift_task_create_future_f(
   initialContext->Flags.setShouldNotDeallocateInCallee(true);
 
   // Initialize the task-local allocator.
+  // TODO: consider providing an initial pre-allocated first slab to the
+  //       allocator.
   _swift_task_alloc_initialize(task);
 
   return {task, initialContext};

stdlib/public/Concurrency/TaskAlloc.cpp

@@ -19,40 +19,27 @@
 
 #include "TaskPrivate.h"
 #include "swift/Runtime/Concurrency.h"
-#include "swift/Runtime/Debug.h"
+#include "../runtime/StackAllocator.h"
 #include <stdlib.h>
-#include <vector>
 
 using namespace swift;
 
 namespace {
 
-class TaskAllocator {
-  // Just keep track of all allocations in a vector so that we can
-  // verify stack discipline.  We should make sure the allocator
-  // implementation strictly verifies allocation order at least
-  // until we've stabilized the compiler implementation.
-  std::vector<void*> Allocations;
+/// The size of an allocator slab.
+///
+/// TODO: find the optimal value by experiment.
+static constexpr size_t SlabCapacity = 1024;
 
-public:
-  void *alloc(size_t size) {
-    void *ptr = malloc(size);
-    Allocations.push_back(ptr);
-    return ptr;
-  }
+using TaskAllocator = StackAllocator<SlabCapacity>;
 
-  void dealloc(void *ptr) {
-    if (Allocations.empty() || Allocations.back() != ptr)
-      fatalError(0, "pointer was not the last allocation on this task");
+struct GlobalAllocator {
+  TaskAllocator allocator;
+  void *spaceForFirstSlab[64];
 
-    Allocations.pop_back();
-    free(ptr);
-  }
+  GlobalAllocator() : allocator(spaceForFirstSlab, sizeof(spaceForFirstSlab)) {}
 };
 
-static_assert(sizeof(TaskAllocator) <= sizeof(AsyncTask::AllocatorPrivate),
-              "task allocator must fit in allocator-private slot");
-
 static_assert(alignof(TaskAllocator) <= alignof(decltype(AsyncTask::AllocatorPrivate)),
               "task allocator must not be more aligned than "
               "allocator-private slot");
@@ -70,8 +57,8 @@ static TaskAllocator &allocator(AsyncTask *task) {
   // FIXME: this fall-back shouldn't be necessary, but it's useful
   // for now, since the current execution tests aren't setting up a task
   // properly.
-  static TaskAllocator global;
-  return global;
+  static GlobalAllocator global;
+  return global.allocator;
 }
 
 void swift::_swift_task_alloc_destroy(AsyncTask *task) {

stdlib/public/runtime/StackAllocator.h

@@ -60,6 +60,9 @@ class StackAllocator {
   /// True if the first slab is pre-allocated.
   bool firstSlabIsPreallocated;
 
+  /// The minimal alignment of allocated memory.
+  static constexpr size_t alignment = alignof(std::max_align_t);
+  
   /// If set to true, memory allocations are checked for buffer overflows and
   /// use-after-free, similar to guard-malloc.
   static constexpr bool guardAllocations =
@@ -90,11 +93,21 @@ class StackAllocator {
       assert((size_t)capacity == newCapacity && "capacity overflow");
     }
 
+    /// The size of the slab header.
+    static size_t headerSize() {
+      return llvm::alignTo(sizeof(Slab), llvm::Align(alignment));
+    }
+
+    /// Return \p size with the added overhead of the slab header.
+    static size_t includingHeader(size_t size) {
+      return headerSize() + size;
+    }
+
     /// Return the payload buffer address at \p atOffset.
     ///
     /// Note: it's valid to call this function on a not-yet-constructed slab.
     char *getAddr(size_t atOffset) {
-      return (char *)(this + 1) + atOffset;
+      return (char *)this + headerSize() + atOffset;
     }
 
     /// Return true if this slab can fit an allocation of \p size.
@@ -162,23 +175,20 @@ class StackAllocator {
       previous(previous), slab(slab) {}
 
     void *getAllocatedMemory() {
-      return (void *)(this + 1);
+      return (char *)this + headerSize();
+    }
+
+    /// The size of the allocation header.
+    static size_t headerSize() {
+      return llvm::alignTo(sizeof(Allocation), llvm::Align(alignment));
     }
 
     /// Return \p size with the added overhead of the allocation header.
     static size_t includingHeader(size_t size) {
-      return size + sizeof(Allocation);
+      return headerSize() + size;
     }
   };
 
-  static constexpr size_t alignment = alignof(std::max_align_t);
-  
-  static_assert(sizeof(Slab) % StackAllocator::alignment == 0,
-      "Slab size must be a multiple of the max allocation alignment");
-
-  static_assert(sizeof(Allocation) % StackAllocator::alignment == 0,
-      "Allocation size must be a multiple of the max allocation alignment");
-
   // Return a slab which is suitable to allocate \p size memory.
   Slab *getSlabForAllocation(size_t size) {
     Slab *slab = (lastAllocation ? lastAllocation->slab : firstSlab);
@@ -204,7 +214,7 @@ class StackAllocator {
     }
     size_t capacity = std::max(SlabCapacity,
                                Allocation::includingHeader(size));
-    void *slabBuffer = malloc(sizeof(Slab) + capacity);
+    void *slabBuffer = malloc(Slab::includingHeader(capacity));
     Slab *newSlab = new (slabBuffer) Slab(capacity);
     if (slab)
       slab->next = newSlab;
@@ -235,10 +245,12 @@ class StackAllocator {
 
   /// Construct a StackAllocator with a pre-allocated first slab.
   StackAllocator(void *firstSlabBuffer, size_t bufferCapacity) {
-    char *start = (char *)llvm::alignAddr(firstSlabBuffer, llvm::Align(alignment));
+    char *start = (char *)llvm::alignAddr(firstSlabBuffer,
+                                          llvm::Align(alignment));
     char *end = (char *)firstSlabBuffer + bufferCapacity;
-    assert(start + sizeof(Slab) <= end && "buffer for first slab too small");
-    firstSlab = new (start) Slab(end - start - sizeof(Slab));
+    assert(start + Slab::headerSize() <= end &&
+           "buffer for first slab too small");
+    firstSlab = new (start) Slab(end - start - Slab::headerSize());
     firstSlabIsPreallocated = true;
   }