P0339R6 polymorphic_allocator<> (#1311)

Co-authored-by: Casey Carter <cartec69@gmail.com>
Co-authored-by: Stephan T. Lavavej <stl@microsoft.com>

stl/inc/xpolymorphic_allocator.h

@@ -186,7 +186,11 @@ namespace pmr {
     }
 
     // CLASS TEMPLATE polymorphic_allocator
+#if _HAS_CXX20 && defined(__cpp_lib_byte)
+    template <class _Ty = byte>
+#else
     template <class _Ty>
+#endif // _HAS_CXX20 && defined(__cpp_lib_byte)
     class polymorphic_allocator {
     public:
         template <class>
@@ -222,6 +226,47 @@ namespace pmr {
             _Resource->deallocate(_Ptr, _Count * sizeof(_Ty), alignof(_Ty));
         }
 
+#if _HAS_CXX20
+        _NODISCARD __declspec(allocator) void* allocate_bytes(
+            const size_t _Bytes, const size_t _Align = alignof(max_align_t)) {
+            return _Resource->allocate(_Bytes, _Align);
+        }
+
+        void deallocate_bytes(void* const _Ptr, const size_t _Bytes,
+            const size_t _Align = alignof(max_align_t)) noexcept /* strengthened */ {
+            _Resource->deallocate(_Ptr, _Bytes, _Align);
+        }
+
+        template <class _Uty>
+        _NODISCARD __declspec(allocator) _Uty* allocate_object(_CRT_GUARDOVERFLOW const size_t _Count = 1) {
+            void* const _Vp = allocate_bytes(_Get_size_of_n<sizeof(_Uty)>(_Count), alignof(_Uty));
+            return static_cast<_Uty*>(_Vp);
+        }
+
+        template <class _Uty>
+        void deallocate_object(_Uty* const _Ptr, const size_t _Count = 1) noexcept /* strengthened */ {
+            deallocate_bytes(_Ptr, _Count * sizeof(_Uty), alignof(_Uty));
+        }
+
+        template <class _Uty, class... _Types>
+        _NODISCARD __declspec(allocator) _Uty* new_object(_Types&&... _Args) {
+            _Uty* const _Ptr = allocate_object<_Uty>();
+            _TRY_BEGIN
+            construct(_Ptr, _STD forward<_Types>(_Args)...);
+            _CATCH_ALL
+            deallocate_object(_Ptr);
+            _RERAISE;
+            _CATCH_END
+            return _Ptr;
+        }
+
+        template <class _Uty>
+        void delete_object(_Uty* const _Ptr) noexcept /* strengthened */ {
+            _Destroy_in_place(*_Ptr);
+            deallocate_object(_Ptr);
+        }
+#endif // _HAS_CXX20
+
         template <class _Uty, class... _Types>
         void construct(_Uty* const _Ptr, _Types&&... _Args) {
             // propagate allocator *this if uses_allocator_v<_Uty, polymorphic_allocator>

stl/inc/yvals_core.h

@@ -137,6 +137,7 @@
 // P0202R3 constexpr For <algorithm> And exchange()
 // P0318R1 unwrap_reference, unwrap_ref_decay
 // P0325R4 to_array()
+// P0339R6 polymorphic_allocator<>
 // P0356R5 bind_front()
 // P0357R3 Supporting Incomplete Types In reference_wrapper
 // P0415R1 constexpr For <complex> (Again)
@@ -1203,6 +1204,7 @@
 #define __cpp_lib_latch                        201907L
 #define __cpp_lib_list_remove_return_type      201806L
 #define __cpp_lib_math_constants               201907L
+#define __cpp_lib_polymorphic_allocator        201902L
 #define __cpp_lib_remove_cvref                 201711L
 #define __cpp_lib_semaphore                    201907L
 #define __cpp_lib_shift                        201806L

tests/std/test.lst

@@ -218,6 +218,7 @@ tests\P0220R1_sample
 tests\P0220R1_searchers
 tests\P0220R1_string_view
 tests\P0325R4_to_array
+tests\P0339R6_polymorphic_allocator
 tests\P0356R5_bind_front
 tests\P0357R3_supporting_incomplete_types_in_reference_wrapper
 tests\P0414R2_shared_ptr_for_arrays

tests/std/tests/P0339R6_polymorphic_allocator/env.lst

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+RUNALL_INCLUDE ..\usual_latest_matrix.lst

tests/std/tests/P0339R6_polymorphic_allocator/test.cpp

@@ -0,0 +1,106 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <memory_resource>
+#include <new>
+
+using std::pmr::polymorphic_allocator;
+
+void allocate_bytes_test() {
+    constexpr int N = 5;
+
+    polymorphic_allocator<> alloc{};
+
+    void* vp = alloc.allocate_bytes(sizeof(int) * N, alignof(int));
+
+    int* arr = static_cast<int*>(vp);
+
+    for (int i = 0; i < N; ++i) {
+        alloc.construct(arr + i, i);
+    }
+
+    for (int i = 0; i < N; ++i) {
+        assert(arr[i] == i);
+    }
+
+    std::destroy(arr, arr + N);
+
+    alloc.deallocate_bytes(vp, sizeof(int) * N, alignof(int));
+
+    void* vp2 = alloc.allocate_bytes(sizeof(int));
+    assert(reinterpret_cast<std::uintptr_t>(vp2) % alignof(std::max_align_t) == 0);
+    alloc.deallocate_bytes(vp2, sizeof(int));
+}
+
+void allocate_object_test() {
+    constexpr int N = 10;
+
+    polymorphic_allocator<> alloc{};
+
+    int* arr = alloc.allocate_object<int>(N);
+
+    for (int i = 0; i < N; ++i) {
+        alloc.construct(arr + i, i);
+    }
+
+    for (int i = 0; i < N; ++i) {
+        assert(arr[i] == i);
+    }
+
+    std::destroy(arr, arr + N);
+
+    alloc.deallocate_object(arr, N);
+
+    // N = 1
+    int* p = alloc.allocate_object<int>();
+    alloc.construct(p, 20);
+    assert(*p == 20);
+    std::destroy_at(p);
+    alloc.deallocate_object(p);
+}
+
+void allocate_object_overflow_test() {
+    constexpr auto threshold = std::numeric_limits<std::size_t>::max() / sizeof(int);
+
+    polymorphic_allocator<> alloc{};
+
+    try {
+        int* vp = alloc.allocate_object<int>(threshold);
+        alloc.deallocate_object(vp, threshold);
+    } catch (const std::bad_alloc&) {
+    } catch (...) {
+        assert(false);
+    }
+
+    try {
+        [[maybe_unused]] int* vp = alloc.allocate_object<int>(threshold + 1);
+    } catch (const std::bad_array_new_length&) {
+        return;
+    } catch (...) {
+        assert(false);
+    }
+    assert(false);
+}
+
+void new_object_test() {
+    polymorphic_allocator<> alloc{};
+
+    int* p = alloc.new_object<int>(20);
+
+    assert(*p == 20);
+
+    alloc.delete_object(p);
+}
+
+
+int main() {
+    allocate_bytes_test();
+    allocate_object_test();
+    allocate_object_overflow_test();
+    new_object_test();
+}

tests/std/tests/VSO_0157762_feature_test_macros/test.cpp

@@ -1092,6 +1092,20 @@ STATIC_ASSERT(__cpp_lib_parallel_algorithm == 201603L);
 #endif
 #endif
 
+#if _HAS_CXX20
+#ifndef __cpp_lib_polymorphic_allocator
+#error __cpp_lib_polymorphic_allocator is not defined
+#elif __cpp_lib_polymorphic_allocator != 201902L
+#error __cpp_lib_polymorphic_allocator is not 201902L
+#else
+STATIC_ASSERT(__cpp_lib_polymorphic_allocator == 201902L);
+#endif
+#else
+#ifdef __cpp_lib_polymorphic_allocator
+#error __cpp_lib_polymorphic_allocator is defined
+#endif
+#endif
+
 #ifndef __cpp_lib_quoted_string_io
 #error __cpp_lib_quoted_string_io is not defined
 #elif __cpp_lib_quoted_string_io != 201304L