Skip to content

Commit 0a6bae5

Browse files
ericnieblerericniebler
committed
make task_scheduler allocator-aware
1 parent b5b9962 commit 0a6bae5

3 files changed

Lines changed: 186 additions & 45 deletions

File tree

‎include/stdexec/__detail/__any.hpp‎

Lines changed: 109 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
#include "__concepts.hpp"
1919
#include "__config.hpp"
20+
#include "__memory.hpp"
2021
#include "__type_traits.hpp"
2122
#include "__typeinfo.hpp"
2223
#include "__utility.hpp"
@@ -28,6 +29,7 @@
2829
#include <bit>
2930
#include <exception>
3031
#include <memory>
32+
#include <new>
3133
#include <span>
3234
#include <type_traits>
3335
#include <utility>
@@ -211,7 +213,7 @@ namespace STDEXEC::__any
211213
{
212214
template <class _Ty>
213215
STDEXEC_ATTRIBUTE(nodiscard, always_inline)
214-
constexpr auto &operator()(_Ty &&__arg) const noexcept
216+
constexpr auto &&operator()(_Ty &&__arg) const noexcept
215217
{
216218
return __arg.__value_(static_cast<_Ty &&>(__arg));
217219
}
@@ -395,6 +397,11 @@ namespace STDEXEC::__any
395397
// needed by MSVC for EBO to work for some reason:
396398
constexpr virtual ~__iroot() = default;
397399

400+
constexpr virtual void __delete_() noexcept
401+
{
402+
this->~__iroot();
403+
}
404+
398405
private:
399406
template <template <class> class, class, class, size_t, size_t>
400407
friend struct __interface_base;
@@ -506,8 +513,9 @@ namespace STDEXEC::__any
506513

507514
//////////////////////////////////////////////////////////////////////////////////////////
508515
// __emplace_into
509-
template <class _Model, class... _Args>
510-
constexpr _Model &__emplace_into([[maybe_unused]] __iroot *&__root_ptr,
516+
template <class _Model, class _Allocator, class... _Args>
517+
constexpr _Model &__emplace_into([[maybe_unused]] _Allocator const &__alloc,
518+
[[maybe_unused]] __iroot *&__root_ptr,
511519
[[maybe_unused]] std::span<std::byte> __buff,
512520
_Args &&...__args)
513521
{
@@ -526,20 +534,27 @@ namespace STDEXEC::__any
526534
}
527535
else
528536
{
529-
auto *const __model = ::new _Model(static_cast<_Args &&>(__args)...);
537+
auto __alloc2 = STDEXEC::__rebind_allocator<_Model>(__alloc);
538+
using __traits_t = std::allocator_traits<decltype(__alloc2)>;
539+
auto *const __model = __traits_t::allocate(__alloc2, 1);
540+
__scope_guard __guard{[&]() noexcept { __traits_t::deallocate(__alloc2, __model, 1); }};
541+
__traits_t::construct(__alloc2, __model, static_cast<_Args &&>(__args)...);
542+
__guard.__dismiss();
530543
*__std::start_lifetime_as<__tagged_ptr>(__buff.data()) = __tagged_ptr(__model);
531544
return *__model;
532545
}
533546
}
534547
}
535548

536-
template <int = 0, class _CvRefValue, class _Value = std::decay_t<_CvRefValue>>
549+
template <int = 0, class _CvRefValue, class _Allocator, class _Value = std::decay_t<_CvRefValue>>
537550
STDEXEC_ATTRIBUTE(always_inline)
538-
constexpr _Value &__emplace_into(__iroot *&__root_ptr,
551+
constexpr _Value &__emplace_into(_Allocator const &__alloc,
552+
__iroot *&__root_ptr,
539553
std::span<std::byte> __buff,
540554
_CvRefValue &&__value)
541555
{
542-
return STDEXEC::__any::__emplace_into<_Value>(__root_ptr,
556+
return STDEXEC::__any::__emplace_into<_Value>(__alloc,
557+
__root_ptr,
543558
__buff,
544559
static_cast<_CvRefValue &&>(__value));
545560
}
@@ -574,23 +589,39 @@ namespace STDEXEC::__any
574589
using __reference_proxy_model = __reference<_Interface>;
575590

576591
// __value
577-
template <template <class> class _Interface, class _Value>
578-
struct __value_root;
592+
template <template <class> class _Interface, class _Value, class _Allocator>
593+
struct __value_root_with_allocator;
579594

580-
template <template <class> class _Interface, class _Value>
595+
template <template <class> class _Interface, class _Value, class _Allocator>
581596
struct __value_model final
582-
: _Interface<__mcall1<__bases_of<_Interface>, __value_root<_Interface, _Value>>>
597+
: _Interface<__mcall1<__bases_of<_Interface>,
598+
__value_root_with_allocator<_Interface, _Value, _Allocator>>>
583599
{
584-
using __base_t = _Interface<__mcall1<__bases_of<_Interface>, __value_root<_Interface, _Value>>>;
600+
using __base_t =
601+
_Interface<__mcall1<__bases_of<_Interface>,
602+
__value_root_with_allocator<_Interface, _Value, _Allocator>>>;
585603
using __base_t::__base_t;
586604

605+
constexpr void __delete_() noexcept final
606+
{
607+
auto __alloc = STDEXEC::__rebind_allocator<__value_model>(this->__get_allocator());
608+
using __traits_t = std::allocator_traits<decltype(__alloc)>;
609+
__traits_t::destroy(__alloc, this);
610+
__traits_t::deallocate(__alloc, this, 1);
611+
}
612+
587613
// This is a virtual override if _Interface extends __imovable
588614
//! @pre __is_small<__value_model>(__buff.size())
589615
constexpr void __move_to(__iroot *&__ptr, std::span<std::byte> __buff) noexcept
590616
{
591617
static_assert(__extension_of<__iabstract<_Interface>, __imovable>);
592618
STDEXEC_ASSERT(STDEXEC::__any::__is_small<__value_model>(__buff.size()));
593-
STDEXEC::__any::__emplace_into(__ptr, __buff, std::move(*this));
619+
STDEXEC::__any::__emplace_into<__value_model>(this->__get_allocator(),
620+
__ptr,
621+
__buff,
622+
std::allocator_arg,
623+
this->__get_allocator(),
624+
__value(std::move(*this)));
594625
__reset(*this);
595626
}
596627

@@ -599,7 +630,12 @@ namespace STDEXEC::__any
599630
{
600631
static_assert(__extension_of<__iabstract<_Interface>, __icopyable>);
601632
STDEXEC_ASSERT(!__empty(*this));
602-
STDEXEC::__any::__emplace_into(__ptr, __buff, *this);
633+
STDEXEC::__any::__emplace_into<__value_model>(this->__get_allocator(),
634+
__ptr,
635+
__buff,
636+
std::allocator_arg,
637+
this->__get_allocator(),
638+
__value(*this));
603639
}
604640
};
605641

@@ -739,6 +775,34 @@ namespace STDEXEC::__any
739775
}
740776
};
741777

778+
//////////////////////////////////////////////////////////////////////////////////////////
779+
// __value_root_with_allocator
780+
template <template <class> class _Interface, class _Value, class _Allocator>
781+
struct STDEXEC_ATTRIBUTE(empty_bases) __value_root_with_allocator
782+
: __value_root<_Interface, _Value>
783+
, private _Allocator
784+
{
785+
constexpr explicit __value_root_with_allocator(_Value __value,
786+
_Allocator const &__alloc = _Allocator())
787+
: __value_root<_Interface, _Value>(std::move(__value))
788+
, _Allocator(__alloc)
789+
{}
790+
791+
template <class... _Args>
792+
constexpr explicit __value_root_with_allocator(std::allocator_arg_t,
793+
_Allocator const &__alloc,
794+
_Args &&...__args)
795+
: __value_root<_Interface, _Value>(static_cast<_Args &&>(__args)...)
796+
, _Allocator(__alloc)
797+
{}
798+
799+
[[nodiscard]]
800+
constexpr auto __get_allocator() const noexcept -> _Allocator const &
801+
{
802+
return *this;
803+
}
804+
};
805+
742806
//////////////////////////////////////////////////////////////////////////////////////////
743807
// __value_proxy_root
744808
template <template <class> class _Interface>
@@ -869,22 +933,27 @@ namespace STDEXEC::__any
869933
friend struct __any;
870934
friend struct __access;
871935

872-
template <class _Value, class... _Args>
873-
constexpr _Value &__emplace_(_Args &&...__args)
936+
template <class _Value, class _Allocator, class... _Args>
937+
constexpr _Value &__emplace_(std::allocator_arg_t, _Allocator const &__alloc, _Args &&...__args)
874938
{
875939
static_assert(__decays_to<_Value, _Value>, "Value must be an object type.");
876-
using __model_type = __value_model<_Interface, _Value>;
877-
auto &__model = STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
878-
__buff_,
879-
static_cast<_Args &&>(
880-
__args)...);
940+
static_assert(__simple_allocator<_Allocator>);
941+
using __model_type = __value_model<_Interface, _Value, _Allocator>;
942+
auto &__model = //
943+
STDEXEC::__any::__emplace_into<__model_type>(__alloc,
944+
__root_ptr_,
945+
__buff_,
946+
std::allocator_arg,
947+
__alloc,
948+
static_cast<_Args &&>(__args)...);
881949
return __value(__model);
882950
}
883951

884-
template <int = 0, class _CvRefValue, class _Value = std::decay_t<_CvRefValue>>
885-
constexpr _Value &__emplace_(_CvRefValue &&__value)
952+
template <class _Value, class... _Args>
953+
constexpr _Value &__emplace_(_Args &&...__args)
886954
{
887-
return __emplace_<_Value>(static_cast<_CvRefValue &&>(__value));
955+
std::allocator<_Value> const __alloc{};
956+
return __emplace_<_Value>(std::allocator_arg, __alloc, static_cast<_Args &&>(__args)...);
888957
}
889958

890959
template <class _Self>
@@ -938,7 +1007,7 @@ namespace STDEXEC::__any
9381007
else if (!__ptr.__is_tagged())
9391008
std::destroy_at(std::addressof(__value(*this)));
9401009
else
941-
delete std::addressof(__value(*this));
1010+
__value(*this).__delete_();
9421011

9431012
__ptr = __tagged_ptr();
9441013
}
@@ -1238,15 +1307,17 @@ namespace STDEXEC::__any
12381307
using __model_type = __reference_model<_Interface, __value_type, __extension_type>;
12391308
if constexpr (_CvModel::__root_kind == __root_kind::__reference)
12401309
{
1241-
STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
1310+
STDEXEC::__any::__emplace_into<__model_type>(std::allocator<std::byte>{}, // not used
1311+
__root_ptr_,
12421312
__buff_,
12431313
__model.__get_value_ptr_(),
12441314
__model.__get_root_ptr_());
12451315
}
12461316
else
12471317
{
12481318
__iroot *__root_ptr = std::addressof(STDEXEC::__unconst(__model));
1249-
STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
1319+
STDEXEC::__any::__emplace_into<__model_type>(std::allocator<std::byte>{}, // not used
1320+
__root_ptr_,
12501321
__buff_,
12511322
static_cast<__value_type *>(nullptr),
12521323
STDEXEC_DECAY_COPY(__root_ptr));
@@ -1258,7 +1329,8 @@ namespace STDEXEC::__any
12581329
{
12591330
static_assert(!__extension_of<_CvValue, _Interface>);
12601331
using __model_type = __reference_model<_Interface, _CvValue>;
1261-
STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
1332+
STDEXEC::__any::__emplace_into<__model_type>(std::allocator<std::byte>{}, // not used
1333+
__root_ptr_,
12621334
__buff_,
12631335
std::addressof(__value),
12641336
static_cast<__iroot *>(nullptr));
@@ -1597,18 +1669,19 @@ namespace STDEXEC::__any
15971669

15981670
// Construct from an object that implements the interface (and is not an any<>
15991671
// itself)
1600-
template <__model_of<_Interface> _Value>
1601-
constexpr __any(_Value __value)
1672+
template <__model_of<_Interface> _Value, class _Allocator = std::allocator<_Value>>
1673+
constexpr __any(_Value __value, _Allocator const &__alloc = _Allocator())
16021674
: __any()
16031675
{
1604-
(*this).__emplace_(std::move(__value));
1676+
static_assert(__simple_allocator<_Allocator>);
1677+
(*this).template __emplace_<_Value>(std::allocator_arg, __alloc, std::move(__value));
16051678
}
16061679

1607-
template <class _Type, class... _Args>
1608-
constexpr explicit __any(std::in_place_type_t<_Type>, _Args &&...__args)
1680+
template <class _Value, class... _Args>
1681+
constexpr explicit __any(std::in_place_type_t<_Value>, _Args &&...__args)
16091682
: __any()
16101683
{
1611-
(*this).template __emplace_<_Type>(static_cast<_Args &&>(__args)...);
1684+
(*this).template __emplace_<_Value>(static_cast<_Args &&>(__args)...);
16121685
}
16131686

16141687
// Implicit derived-to-base conversion constructor
@@ -1634,7 +1707,7 @@ namespace STDEXEC::__any
16341707
constexpr __any &operator=(_Value __value)
16351708
{
16361709
__reset(*this);
1637-
(*this).__emplace_(std::move(__value));
1710+
(*this).template __emplace_<_Value>(std::move(__value));
16381711
return *this;
16391712
}
16401713

@@ -1654,7 +1727,7 @@ namespace STDEXEC::__any
16541727
&& (_Other::__root_kind == __root_kind::__reference)
16551728
constexpr __any &operator=(_Interface<_Other> const &__other)
16561729
{
1657-
// Guard against __self-assignment when __other is a reference to *this
1730+
// Guard against self-assignment when __other is a reference to *this
16581731
if (__data(__other) == __data(*this))
16591732
return *this;
16601733

‎include/stdexec/__detail/__task_scheduler.hpp‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -227,10 +227,11 @@ namespace STDEXEC
227227
public:
228228
using scheduler_concept = scheduler_t;
229229

230-
template <__not_same_as<task_scheduler> _Sch, class _Alloc = std::allocator<std::byte>>
230+
template <__not_same_as<task_scheduler> _Sch, class _Alloc = std::allocator<_Sch>>
231231
requires __infallible_scheduler<_Sch, __task::__env_t<true>>
232-
constexpr explicit task_scheduler(_Sch __sch, _Alloc __alloc = {})
233-
: __backend_(__backend_for{std::move(__sch), __alloc})
232+
constexpr explicit task_scheduler(_Sch __sch, _Alloc const & __alloc = {})
233+
: __backend_(__backend_for{std::move(__sch), STDEXEC::__rebind_allocator<_Sch>(__alloc)},
234+
STDEXEC::__rebind_allocator<__backend_for<_Sch, _Alloc>>(__alloc))
234235
{
235236
using __opstate_t = __task::__opstate<_Alloc, schedule_result_t<_Sch>, __task::__env_t<true>>;
236237
static_assert(sizeof(__opstate_t) <= STDEXEC_TASK_SCHEDULE_OPSTATE_SIZE,
@@ -668,10 +669,11 @@ namespace STDEXEC
668669
// __emplace_into
669670
template <class _Ty, class _Alloc, class... _Args>
670671
constexpr auto
671-
__emplace_into(std::span<std::byte> __storage, _Alloc& __alloc, _Args&&... __args) -> _Ty&
672+
__emplace_into(std::span<std::byte> __storage, _Alloc const & __alloc, _Args&&... __args)
673+
-> _Ty&
672674
{
673-
using __traits_t = std::allocator_traits<_Alloc>::template rebind_traits<_Ty>;
674675
auto __alloc_copy = STDEXEC::__rebind_allocator<_Ty>(__alloc);
676+
using __traits_t = std::allocator_traits<decltype(__alloc_copy)>;
675677

676678
bool const __in_situ = __storage.size() >= sizeof(_Ty);
677679
auto* __ptr = __in_situ ? reinterpret_cast<_Ty*>(__storage.data())
@@ -756,10 +758,10 @@ namespace STDEXEC
756758
{
757759
STDEXEC_TRY
758760
{
759-
using __opstate_t = __task::__opstate<_Alloc, _Sndr, _Env>;
760-
bool const __in_situ = __storage.size() >= sizeof(__opstate_t);
761-
_Alloc& __alloc = *this;
762-
auto& __opstate = __task::__emplace_into<__opstate_t>(__storage,
761+
using __opstate_t = __task::__opstate<_Alloc, _Sndr, _Env>;
762+
bool const __in_situ = __storage.size() >= sizeof(__opstate_t);
763+
_Alloc const & __alloc = *this;
764+
auto& __opstate = __task::__emplace_into<__opstate_t>(__storage,
763765
__alloc,
764766
__alloc,
765767
static_cast<_Sndr&&>(__sndr),

0 commit comments

Comments
 (0)