[V4] Half thread-local storage

benchmark/src/libfork_benchmark/fib/fib.hpp

@@ -59,32 +59,41 @@ struct tls_bump {
   }
 };
 
-constinit inline std::byte *bump_ptr = nullptr;
+// === Shared Context Logic ===
 
-struct global_bump {
+template <lf::stack_allocator Alloc>
+struct vector_ctx {
 
-  static auto operator new(std::size_t sz) -> void * {
-    auto *prev = bump_ptr;
-    bump_ptr += fib_align_size(sz);
-    return prev;
-  }
+  using handle_type = lf::frame_handle<vector_ctx>;
 
-  static auto operator delete(void *p, [[maybe_unused]] std::size_t sz) noexcept -> void {
-    bump_ptr = std::bit_cast<std::byte *>(p);
+  std::vector<handle_type> work;
+  Alloc allocator;
+
+  vector_ctx() { work.reserve(1024); }
+
+  auto alloc() noexcept -> Alloc & { return allocator; }
+
+  void push(handle_type handle) { work.push_back(handle); }
+
+  auto pop() noexcept -> handle_type {
+    auto handle = work.back();
+    work.pop_back();
+    return handle;
   }
 };
 
-// === Shared Context Logic ===
+template <lf::stack_allocator Alloc>
+struct poly_vector_ctx final : lf::polymorphic_context<Alloc> {
 
-struct vector_ctx final : lf::polymorphic_context {
+  using handle_type = lf::frame_handle<lf::polymorphic_context<Alloc>>;
 
-  std::vector<lf::work_handle> work;
+  std::vector<handle_type> work;
 
-  vector_ctx() { work.reserve(1024); }
+  poly_vector_ctx() { work.reserve(1024); }
 
-  void push(lf::work_handle handle) override { work.push_back(handle); }
+  void push(handle_type handle) override { work.push_back(handle); }
 
-  auto pop() noexcept -> lf::work_handle override {
+  auto pop() noexcept -> handle_type override {
     auto handle = work.back();
     work.pop_back();
     return handle;

benchmark/src/libfork_benchmark/fib/lf_parts.cpp

@@ -12,15 +12,41 @@ import libfork.core;
 
 namespace {
 
-struct stack_on_heap {
-  static constexpr auto operator new(std::size_t sz) -> void * { return ::operator new(sz); }
-  static constexpr auto operator delete(void *p, [[maybe_unused]] std::size_t sz) noexcept -> void {
-    ::operator delete(p, sz);
+struct global_allocator {
+
+  struct empty {};
+
+  constexpr static auto push(std::size_t sz) -> void * { return ::operator new(sz); }
+  constexpr static auto pop(void *p, std::size_t sz) noexcept -> void { ::operator delete(p, sz); }
+  constexpr static auto checkpoint() noexcept -> empty { return {}; }
+  constexpr static auto switch_to(empty) noexcept -> void {}
+};
+
+static_assert(lf::stack_allocator<global_allocator>);
+
+struct linear_allocator {
+
+  std::unique_ptr<std::byte[]> data = std::make_unique<std::byte[]>(1024 * 1024);
+  std::byte *ptr = data.get();
+
+  constexpr auto push(std::size_t sz) -> void * {
+    auto *prev = ptr;
+    ptr += fib_align_size(sz);
+    return prev;
   }
+  constexpr auto pop(void *p, std::size_t) noexcept -> void { ptr = static_cast<std::byte *>(p); }
+
+  constexpr auto checkpoint() noexcept -> std::byte * { return data.get(); }
+
+  constexpr static auto switch_to(std::byte *) noexcept -> void {}
 };
 
-template <lf::alloc_mixin Stack>
-constexpr auto no_await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, Stack> {
+static_assert(lf::stack_allocator<linear_allocator>);
+
+using lf::task;
+
+template <lf::worker_context T>
+constexpr auto no_await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> task<void, T> {
   if (n < 2) {
     *ret = n;
     co_return;
@@ -40,8 +66,8 @@ constexpr auto no_await = [](this auto fib, std::int64_t *ret, std::int64_t n) -
   *ret = lhs + rhs;
 };
 
-template <lf::alloc_mixin Stack>
-constexpr auto await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, Stack> {
+template <lf::worker_context T>
+constexpr auto await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, T> {
   if (n < 2) {
     *ret = n;
     co_return;
@@ -56,7 +82,8 @@ constexpr auto await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> l
   *ret = lhs + rhs;
 };
 
-constexpr auto ret = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, tls_bump> {
+template <lf::worker_context T>
+constexpr auto ret = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, T> {
   if (n < 2) {
     co_return n;
   }
@@ -70,8 +97,8 @@ constexpr auto ret = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t,
   co_return lhs + rhs;
 };
 
-template <typename Ctx, typename A = tls_bump>
-constexpr auto fork_call = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, A, Ctx> {
+template <typename T>
+constexpr auto fork_call = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, T> {
   if (n < 2) {
     co_return n;
   }
@@ -85,23 +112,24 @@ constexpr auto fork_call = [](this auto fib, std::int64_t n) -> lf::task<std::in
   co_return lhs + rhs;
 };
 
-template <auto Fn>
+using global_alloc = vector_ctx<global_allocator>;
+using linear_alloc = vector_ctx<linear_allocator>;
+
+template <auto Fn, typename T, typename U = T>
 void fib(benchmark::State &state) {
 
   std::int64_t n = state.range(0);
   std::int64_t expect = fib_ref(n);
 
   state.counters["n"] = static_cast<double>(n);
 
-  // Set bump allocator buffer
-  std::unique_ptr buf = std::make_unique<std::byte[]>(1024 * 1024);
-  tls_bump_ptr = buf.get();
-  bump_ptr = buf.get();
+  T context;
+
+  lf::thread_context<U> = static_cast<U *>(&context);
 
-  // Set both context and poly context
-  std::unique_ptr ctx = std::make_unique<vector_ctx>();
-  lf::thread_context<vector_ctx> = ctx.get();
-  lf::thread_context<lf::polymorphic_context> = ctx.get();
+  lf::defer _ = [] static noexcept {
+    lf::thread_context<U> = nullptr;
+  };
 
   for (auto _ : state) {
     benchmark::DoNotOptimize(n);
@@ -121,55 +149,44 @@ void fib(benchmark::State &state) {
     CHECK_RESULT(result, expect);
     benchmark::DoNotOptimize(result);
   }
-
-  if (tls_bump_ptr != buf.get() || bump_ptr != buf.get()) {
-    LF_TERMINATE("Stack leak detected");
-  }
-
-  tls_bump_ptr = nullptr;
-  bump_ptr = nullptr;
-  lf::thread_context<vector_ctx> = nullptr;
-  lf::thread_context<lf::polymorphic_context> = nullptr;
 }
 
 } // namespace
 
-// Return by ref-arg, test direct root, no co-await, direct resumes, uses new/delete for alloc
-BENCHMARK(fib<no_await<stack_on_heap>>)->Name("test/libfork/fib/heap/no_await")->Arg(fib_test);
-BENCHMARK(fib<no_await<stack_on_heap>>)->Name("base/libfork/fib/heap/no_await")->Arg(fib_base);
+static_assert(lf::worker_context<global_alloc>);
 
-// Same as above but uses tls bump allocator
-BENCHMARK(fib<no_await<tls_bump>>)->Name("test/libfork/fib/tls_bump/no_await")->Arg(fib_test);
-BENCHMARK(fib<no_await<tls_bump>>)->Name("base/libfork/fib/tls_bump/no_await")->Arg(fib_base);
+// Return by ref-arg, test direct root, no co-await, direct resumes, uses new/delete for alloc
+BENCHMARK(fib<no_await<global_alloc>, global_alloc>)->Name("test/libfork/fib/heap/no_await")->Arg(fib_test);
+BENCHMARK(fib<no_await<global_alloc>, global_alloc>)->Name("base/libfork/fib/heap/no_await")->Arg(fib_base);
 
-// Same as above but with global bump allocator
-BENCHMARK(fib<no_await<global_bump>>)->Name("test/libfork/fib/global_bump/no_await")->Arg(fib_test);
-BENCHMARK(fib<no_await<global_bump>>)->Name("base/libfork/fib/global_bump/no_await")->Arg(fib_base);
+// Same as above but uses bump allocator
+BENCHMARK(fib<no_await<linear_alloc>, linear_alloc>)->Name("test/libfork/fib/bump/no_await")->Arg(fib_test);
+BENCHMARK(fib<no_await<linear_alloc>, linear_alloc>)->Name("base/libfork/fib/bump/no_await")->Arg(fib_base);
 
 // TODO: no_await with segmented stack allocator?
 
 // Return by ref-arg, libfork call/call with co-await, uses new/delete for alloc
-BENCHMARK(fib<await<stack_on_heap>>)->Name("test/libfork/fib/heap/await")->Arg(fib_test);
-BENCHMARK(fib<await<stack_on_heap>>)->Name("base/libfork/fib/heap/await")->Arg(fib_base);
+BENCHMARK(fib<await<global_alloc>, global_alloc>)->Name("test/libfork/fib/heap/await")->Arg(fib_test);
+BENCHMARK(fib<await<global_alloc>, global_alloc>)->Name("base/libfork/fib/heap/await")->Arg(fib_base);
 
-// Same as above but uses tls bump allocator
-BENCHMARK(fib<await<tls_bump>>)->Name("test/libfork/fib/tls_bump/await")->Arg(fib_test);
-BENCHMARK(fib<await<tls_bump>>)->Name("base/libfork/fib/tls_bump/await")->Arg(fib_base);
-
-// Same as above but with global bump allocator
-BENCHMARK(fib<await<global_bump>>)->Name("test/libfork/fib/global_bump/await")->Arg(fib_test);
-BENCHMARK(fib<await<global_bump>>)->Name("base/libfork/fib/global_bump/await")->Arg(fib_base);
+// // Same as above but uses bump allocator
+BENCHMARK(fib<await<linear_alloc>, linear_alloc>)->Name("test/libfork/fib/bump/await")->Arg(fib_test);
+BENCHMARK(fib<await<linear_alloc>, linear_alloc>)->Name("base/libfork/fib/bump/await")->Arg(fib_base);
 
 // Return by value
 // libfork call/call with co-await
-BENCHMARK(fib<ret>)->Name("test/libfork/fib/tls_bump/return")->Arg(fib_test);
-BENCHMARK(fib<ret>)->Name("base/libfork/fib/tls_bump/return")->Arg(fib_base);
+BENCHMARK(fib<ret<linear_alloc>, linear_alloc>)->Name("test/libfork/fib/bump/return")->Arg(fib_test);
+BENCHMARK(fib<ret<linear_alloc>, linear_alloc>)->Name("base/libfork/fib/bump/return")->Arg(fib_base);
 
 // Return by value
 // libfork call/fork (no join)
 // Non-polymorphic vector-backed context
-BENCHMARK(fib<fork_call<vector_ctx>>)->Name("test/libfork/fib/vector_ctx")->Arg(fib_test);
-BENCHMARK(fib<fork_call<vector_ctx>>)->Name("base/libfork/fib/vector_ctx")->Arg(fib_base);
+BENCHMARK(fib<fork_call<linear_alloc>, linear_alloc>)->Name("test/libfork/fib/vector_ctx")->Arg(fib_test);
+BENCHMARK(fib<fork_call<linear_alloc>, linear_alloc>)->Name("base/libfork/fib/vector_ctx")->Arg(fib_base);
+
+using A = poly_vector_ctx<linear_allocator>;
+using B = lf::polymorphic_context<linear_allocator>;
 
-BENCHMARK(fib<fork_call<lf::polymorphic_context>>)->Name("test/libfork/fib/poly_vector_ctx")->Arg(fib_test);
-BENCHMARK(fib<fork_call<lf::polymorphic_context>>)->Name("base/libfork/fib/poly_vector_ctx")->Arg(fib_base);
+// Same as above but with polymorphic contexts.
+BENCHMARK(fib<fork_call<B>, A, B>)->Name("test/libfork/fib/poly_vector_ctx")->Arg(fib_test);
+BENCHMARK(fib<fork_call<B>, A, B>)->Name("base/libfork/fib/poly_vector_ctx")->Arg(fib_base);

src/core/concepts.cxx

@@ -5,6 +5,24 @@ import std;
 
 namespace lf {
 
+// ========== Specialization ========== //
+
+template <typename T, template <typename...> typename Template>
+struct is_specialization_of : std::false_type {};
+
+template <template <typename...> typename Template, typename... Args>
+struct is_specialization_of<Template<Args...>, Template> : std::true_type {};
+
+/**
+ * @brief Test if `T` is a specialization of the template `Template`.
+ */
+export template <typename T, template <typename...> typename Template>
+concept specialization_of = is_specialization_of<std::remove_cvref_t<T>, Template>::value;
+
+// ========== Task constraint related concepts ========== //
+
+// ==== Returnable
+
 /**
  * @brief A type returnable from libfork's async functions/coroutines.
  *
@@ -13,31 +31,78 @@ namespace lf {
 template <typename T>
 concept returnable = std::is_void_v<T> || std::movable<T>;
 
+// ==== Stack
+
 template <typename T>
-concept mixinable = std::is_empty_v<T> && !std::is_final_v<T>;
+  requires std::is_object_v<T>
+consteval auto constify(T &&x) noexcept -> std::add_const_t<T> &;
 
+/**
+ * @brief Defines the API for a libfork compatible stack allocator.
+ *
+ * - After construction push is valid.
+ * - Pop is valid provided the FILO order is respected.
+ * - Destruction is expected to only occur when the stack is empty.
+ * - Result of `.checkpoint()` is expected to be "cheap to copy".
+ * - Switch releases the current stack and resumes from the checkpoint:
+ *     - This is a noop if the checkpoint is from this stack.
+ *     - If the checkpoint is default-constructed it is expected to switch to a new stack.
+ *
+ * Fast-path operations: empty, push, pop, checkpoint
+ * Slow-path operations: switch
+ */
 export template <typename T>
-concept alloc_mixin = mixinable<T> && requires (std::size_t n, T *ptr) {
-  { T::operator new(n) } -> std::same_as<void *>;
-  { T::operator delete(ptr, n) } noexcept -> std::same_as<void>;
+concept stack_allocator = std::is_object_v<T> && requires (T alloc, std::size_t n, void *ptr) {
+  // { alloc.empty() } noexcept -> std::same_as<bool>;
+  { alloc.push(n) } -> std::same_as<void *>;
+  { alloc.pop(ptr, n) } noexcept -> std::same_as<void>;
+  { alloc.checkpoint() } noexcept -> std::semiregular;
+  { alloc.switch_to({}) } noexcept -> std::same_as<void>;
+  { alloc.switch_to(constify(alloc.checkpoint())) } noexcept -> std::same_as<void>;
 };
 
-template <typename T, template <typename...> typename Template>
-struct is_specialization_of : std::false_type {};
+/**
+ * @brief Fetch the checkpoint type of a stack allocator `T`.
+ */
+template <stack_allocator T>
+using checkpoint_t = decltype(std::declval<T &>().checkpoint());
 
-template <template <typename...> typename Template, typename... Args>
-struct is_specialization_of<Template<Args...>, Template> : std::true_type {};
+// ==== Context
+
+export template <typename T>
+class frame_handle;
+
+template <typename T>
+concept ref_to_stack_allocator = std::is_lvalue_reference_v<T> && stack_allocator<std::remove_reference_t<T>>;
 
 /**
- * @brief Test if `T` is a specialization of the template `Template`.
+ * @brief Defines the API for a libfork compatible worker context.
+ *
+ * This requires that `T` is an object type and supports the following operations:
+ *
+ * - Push/pop a frame handle onto the context in a LIFO manner.
+ * - Have a `stack_allocator` that can be accessed via `alloc()`.
  */
-template <typename T, template <typename...> typename Template>
-concept specialization_of = is_specialization_of<std::remove_cvref_t<T>, Template>::value;
+export template <typename T>
+concept worker_context = std::is_object_v<T> && requires (T ctx, frame_handle<T> handle) {
+  { ctx.alloc() } noexcept -> ref_to_stack_allocator;
+  { ctx.push(handle) } -> std::same_as<void>;
+  { ctx.pop() } noexcept -> std::same_as<frame_handle<T>>;
+};
+
+/**
+ * @brief Fetch the allocator type of a worker context `T`.
+ */
+template <worker_context T>
+using allocator_t = std::remove_reference_t<decltype(std::declval<T &>().alloc())>;
+
+// ==== Forward-decl
 
-// Forward-decl
-export template <returnable T, alloc_mixin Stack, typename Context>
+export template <returnable T, worker_context Context>
 struct task;
 
+// ========== Invocability ========== //
+
 /**
  * @brief Test if a callable `Fn` when invoked with `Args...` returns an `lf::task`.
  */
@@ -50,9 +115,12 @@ concept async_invocable =
  */
 export template <typename Fn, typename... Args>
   requires async_invocable<Fn, Args...>
-using async_result_t = std::invoke_result_t<Fn, Args...>::type;
+using async_result_t = std::invoke_result_t<Fn, Args...>::value_type;
 
-template <typename Fn, typename R, typename... Args>
-concept async_invocable_to = async_invocable<Fn, Args...> && std::same_as<async_result_t<Fn, Args...>, R>;
+/**
+ * @brief Subsumes `async_invocable` and checks the result type is `R`.
+ */
+export template <typename Fn, typename R, typename... Args>
+concept async_invocable_to = async_invocable<Fn, Args...> && std::same_as<R, async_result_t<Fn, Args...>>;
 
 } // namespace lf