[libc++] Pass type information down to __libcpp_allocate (#118837)

ldionne · web-flow · commit cedb44af53f1 · 2025-01-13T09:10:36.000-05:00
Currently, places where we call __libcpp_allocate must drop type information on the ground even when they actually have such information available. That is unfortunate since some toolchains and system allocators are able to provide improved security when they know what type is being allocated. This is the purpose of http://wg21.link/p2719, where we introduce a new variant of `operator new` which takes a type in its interface. A different but related issue is that `std::allocator` does not honor any in-class `T::operator new` since it is specified to call the global `::operator new` instead. This patch closes the gap to make it trivial for implementations that provide typed memory allocators to actually benefit from that information in more contexts, and also makes libc++ forward-compatible with future proposals that would fix the existing defects in `std::allocator`. It also makes the internal allocation API higher level by operating on objects instead of operating on bytes of memory. Since this is a widely-used function and making this a template could have an impact on debug info sizes, I tried minimizing the number of templated layers by removing `__do_deallocate_handle_size`, which was easy to replace with a macro (and IMO this leads to cleaner code).
diff --git a/libcxx/include/CMakeLists.txt b/libcxx/include/CMakeLists.txt
@@ -547,7 +547,6 @@ set(files
   __memory/array_cookie.h
   __memory/assume_aligned.h
   __memory/auto_ptr.h
-  __memory/builtin_new_allocator.h
   __memory/compressed_pair.h
   __memory/concepts.h
   __memory/construct_at.h
@@ -880,6 +879,7 @@ set(files
   __utility/cmp.h
   __utility/convert_to_integral.h
   __utility/declval.h
+  __utility/element_count.h
   __utility/empty.h
   __utility/exception_guard.h
   __utility/exchange.h
diff --git a/libcxx/include/__functional/function.h b/libcxx/include/__functional/function.h
@@ -22,7 +22,6 @@
 #include <__memory/allocator.h>
 #include <__memory/allocator_destructor.h>
 #include <__memory/allocator_traits.h>
-#include <__memory/builtin_new_allocator.h>
 #include <__memory/compressed_pair.h>
 #include <__memory/unique_ptr.h>
 #include <__type_traits/aligned_storage.h>
@@ -193,6 +192,13 @@ class __alloc_func<_Fp, _Ap, _Rp(_ArgTypes...)> {
   }
 };
 
+template <class _Tp>
+struct __deallocating_deleter {
+  _LIBCPP_HIDE_FROM_ABI void operator()(void* __p) const {
+    std::__libcpp_deallocate<_Tp>(static_cast<_Tp*>(__p), __element_count(1));
+  }
+};
+
 template <class _Fp, class _Rp, class... _ArgTypes>
 class __default_alloc_func<_Fp, _Rp(_ArgTypes...)> {
   _Fp __f_;
@@ -212,8 +218,9 @@ class __default_alloc_func<_Fp, _Rp(_ArgTypes...)> {
   }
 
   _LIBCPP_HIDE_FROM_ABI __default_alloc_func* __clone() const {
-    __builtin_new_allocator::__holder_t __hold = __builtin_new_allocator::__allocate_type<__default_alloc_func>(1);
-    __default_alloc_func* __res                = ::new ((void*)__hold.get()) __default_alloc_func(__f_);
+    using _Self = __default_alloc_func;
+    unique_ptr<_Self, __deallocating_deleter<_Self>> __hold(std::__libcpp_allocate<_Self>(__element_count(1)));
+    _Self* __res = ::new ((void*)__hold.get()) _Self(__f_);
     (void)__hold.release();
     return __res;
   }
@@ -222,7 +229,7 @@ class __default_alloc_func<_Fp, _Rp(_ArgTypes...)> {
 
   _LIBCPP_HIDE_FROM_ABI static void __destroy_and_delete(__default_alloc_func* __f) {
     __f->destroy();
-    __builtin_new_allocator::__deallocate_type<__default_alloc_func>(__f, 1);
+    std::__libcpp_deallocate<__default_alloc_func>(__f, __element_count(1));
   }
 };
 
@@ -668,8 +675,8 @@ class __policy_func<_Rp(_ArgTypes...)> {
       if (__use_small_storage<_Fun>()) {
         ::new ((void*)&__buf_.__small) _Fun(std::move(__f));
       } else {
-        __builtin_new_allocator::__holder_t __hold = __builtin_new_allocator::__allocate_type<_Fun>(1);
-        __buf_.__large                             = ::new ((void*)__hold.get()) _Fun(std::move(__f));
+        unique_ptr<_Fun, __deallocating_deleter<_Fun>> __hold(std::__libcpp_allocate<_Fun>(__element_count(1)));
+        __buf_.__large = ::new ((void*)__hold.get()) _Fun(std::move(__f));
         (void)__hold.release();
       }
     }
diff --git a/libcxx/include/__memory/allocator.h b/libcxx/include/__memory/allocator.h
@@ -102,7 +102,7 @@ class _LIBCPP_TEMPLATE_VIS allocator : private __non_trivial_if<!is_void<_Tp>::v
     if (__libcpp_is_constant_evaluated()) {
       return static_cast<_Tp*>(::operator new(__n * sizeof(_Tp)));
     } else {
-      return static_cast<_Tp*>(std::__libcpp_allocate(__n * sizeof(_Tp), _LIBCPP_ALIGNOF(_Tp)));
+      return std::__libcpp_allocate<_Tp>(__element_count(__n));
     }
   }
 
@@ -117,7 +117,7 @@ class _LIBCPP_TEMPLATE_VIS allocator : private __non_trivial_if<!is_void<_Tp>::v
     if (__libcpp_is_constant_evaluated()) {
       ::operator delete(__p);
     } else {
-      std::__libcpp_deallocate((void*)__p, __n * sizeof(_Tp), _LIBCPP_ALIGNOF(_Tp));
+      std::__libcpp_deallocate<_Tp>(__p, __element_count(__n));
     }
   }
 
diff --git a/libcxx/include/__memory/builtin_new_allocator.h b/libcxx/include/__memory/builtin_new_allocator.h
diff --git a/libcxx/include/__memory/unique_temporary_buffer.h b/libcxx/include/__memory/unique_temporary_buffer.h
@@ -40,7 +40,7 @@ struct __temporary_buffer_deleter {
       return;
     }
 
-    std::__libcpp_deallocate_unsized((void*)__ptr, _LIBCPP_ALIGNOF(_Tp));
+    std::__libcpp_deallocate_unsized<_Tp>(__ptr);
   }
 };
 
diff --git a/libcxx/include/__new/allocate.h b/libcxx/include/__new/allocate.h
@@ -14,6 +14,8 @@
 #include <__cstddef/size_t.h>
 #include <__new/align_val_t.h>
 #include <__new/global_new_delete.h> // for _LIBCPP_HAS_SIZED_DEALLOCATION
+#include <__type_traits/type_identity.h>
+#include <__utility/element_count.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -47,52 +49,58 @@ _LIBCPP_HIDE_FROM_ABI void __libcpp_operator_delete(_Args... __args) _NOEXCEPT {
 #endif
 }
 
-inline _LIBCPP_HIDE_FROM_ABI void* __libcpp_allocate(size_t __size, size_t __align) {
+template <class _Tp>
+inline _LIBCPP_HIDE_FROM_ABI _Tp* __libcpp_allocate(__element_count __n, size_t __align = _LIBCPP_ALIGNOF(_Tp)) {
+  size_t __size = static_cast<size_t>(__n) * sizeof(_Tp);
 #if _LIBCPP_HAS_ALIGNED_ALLOCATION
   if (__is_overaligned_for_new(__align)) {
     const align_val_t __align_val = static_cast<align_val_t>(__align);
-    return __libcpp_operator_new(__size, __align_val);
+    return static_cast<_Tp*>(std::__libcpp_operator_new(__size, __align_val));
   }
 #endif
 
   (void)__align;
-  return __libcpp_operator_new(__size);
+  return static_cast<_Tp*>(std::__libcpp_operator_new(__size));
 }
 
-template <class... _Args>
-_LIBCPP_HIDE_FROM_ABI void __do_deallocate_handle_size(void* __ptr, size_t __size, _Args... __args) _NOEXCEPT {
-#if !_LIBCPP_HAS_SIZED_DEALLOCATION
-  (void)__size;
-  return std::__libcpp_operator_delete(__ptr, __args...);
+#if _LIBCPP_HAS_SIZED_DEALLOCATION
+#  define _LIBCPP_ONLY_IF_SIZED_DEALLOCATION(...) __VA_ARGS__
 #else
-  return std::__libcpp_operator_delete(__ptr, __size, __args...);
+#  define _LIBCPP_ONLY_IF_SIZED_DEALLOCATION(...) /* nothing */
 #endif
-}
 
-inline _LIBCPP_HIDE_FROM_ABI void __libcpp_deallocate(void* __ptr, size_t __size, size_t __align) _NOEXCEPT {
+template <class _Tp>
+inline _LIBCPP_HIDE_FROM_ABI void __libcpp_deallocate(
+    __type_identity_t<_Tp>* __ptr, __element_count __n, size_t __align = _LIBCPP_ALIGNOF(_Tp)) _NOEXCEPT {
+  size_t __size = static_cast<size_t>(__n) * sizeof(_Tp);
+  (void)__size;
 #if !_LIBCPP_HAS_ALIGNED_ALLOCATION
   (void)__align;
-  return __do_deallocate_handle_size(__ptr, __size);
+  return std::__libcpp_operator_delete(__ptr _LIBCPP_ONLY_IF_SIZED_DEALLOCATION(, __size));
 #else
   if (__is_overaligned_for_new(__align)) {
     const align_val_t __align_val = static_cast<align_val_t>(__align);
-    return __do_deallocate_handle_size(__ptr, __size, __align_val);
+    return std::__libcpp_operator_delete(__ptr _LIBCPP_ONLY_IF_SIZED_DEALLOCATION(, __size), __align_val);
   } else {
-    return __do_deallocate_handle_size(__ptr, __size);
+    return std::__libcpp_operator_delete(__ptr _LIBCPP_ONLY_IF_SIZED_DEALLOCATION(, __size));
   }
 #endif
 }
 
-inline _LIBCPP_HIDE_FROM_ABI void __libcpp_deallocate_unsized(void* __ptr, size_t __align) _NOEXCEPT {
+#undef _LIBCPP_ONLY_IF_SIZED_DEALLOCATION
+
+template <class _Tp>
+inline _LIBCPP_HIDE_FROM_ABI void
+__libcpp_deallocate_unsized(__type_identity_t<_Tp>* __ptr, size_t __align = _LIBCPP_ALIGNOF(_Tp)) _NOEXCEPT {
 #if !_LIBCPP_HAS_ALIGNED_ALLOCATION
   (void)__align;
-  return __libcpp_operator_delete(__ptr);
+  return std::__libcpp_operator_delete(__ptr);
 #else
   if (__is_overaligned_for_new(__align)) {
     const align_val_t __align_val = static_cast<align_val_t>(__align);
-    return __libcpp_operator_delete(__ptr, __align_val);
+    return std::__libcpp_operator_delete(__ptr, __align_val);
   } else {
-    return __libcpp_operator_delete(__ptr);
+    return std::__libcpp_operator_delete(__ptr);
   }
 #endif
 }
diff --git a/libcxx/include/__string/constexpr_c_functions.h b/libcxx/include/__string/constexpr_c_functions.h
@@ -25,6 +25,7 @@
 #include <__type_traits/is_trivially_copyable.h>
 #include <__type_traits/is_trivially_lexicographically_comparable.h>
 #include <__type_traits/remove_cv.h>
+#include <__utility/element_count.h>
 #include <__utility/is_pointer_in_range.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
@@ -33,10 +34,6 @@
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
-// Type used to encode that a function takes an integer that represents a number
-// of elements as opposed to a number of bytes.
-enum class __element_count : size_t {};
-
 template <class _Tp>
 inline const bool __is_char_type = false;
 
diff --git a/libcxx/include/__utility/element_count.h b/libcxx/include/__utility/element_count.h
@@ -0,0 +1,27 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBCPP___UTILITY_ELEMENT_COUNT_H
+#define _LIBCPP___UTILITY_ELEMENT_COUNT_H
+
+#include <__config>
+#include <__cstddef/size_t.h>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#  pragma GCC system_header
+#endif
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+
+// Type used to encode that a function takes an integer that represents a number
+// of elements as opposed to a number of bytes.
+enum class __element_count : size_t {};
+
+_LIBCPP_END_NAMESPACE_STD
+
+#endif // _LIBCPP___UTILITY_ELEMENT_COUNT_H
diff --git a/libcxx/include/__utility/small_buffer.h b/libcxx/include/__utility/small_buffer.h
@@ -68,7 +68,7 @@ class __small_buffer {
     if constexpr (__fits_in_buffer<_Stored>) {
       return std::launder(reinterpret_cast<_Stored*>(__buffer_));
     } else {
-      byte* __allocation = static_cast<byte*>(std::__libcpp_allocate(sizeof(_Stored), alignof(_Stored)));
+      byte* __allocation = reinterpret_cast<byte*>(std::__libcpp_allocate<_Stored>(__element_count(1)));
       std::construct_at(reinterpret_cast<byte**>(__buffer_), __allocation);
       return std::launder(reinterpret_cast<_Stored*>(__allocation));
     }
@@ -77,7 +77,7 @@ class __small_buffer {
   template <class _Stored>
   _LIBCPP_HIDE_FROM_ABI void __dealloc() noexcept {
     if constexpr (!__fits_in_buffer<_Stored>)
-      std::__libcpp_deallocate(*reinterpret_cast<void**>(__buffer_), sizeof(_Stored), alignof(_Stored));
+      std::__libcpp_deallocate<_Stored>(__get<_Stored>(), __element_count(1));
   }
 
   template <class _Stored, class... _Args>
diff --git a/libcxx/include/module.modulemap b/libcxx/include/module.modulemap
@@ -1527,14 +1527,16 @@ module std [system] {
     module aligned_alloc                      { header "__memory/aligned_alloc.h" }
     module allocate_at_least                  { header "__memory/allocate_at_least.h" }
     module allocation_guard                   { header "__memory/allocation_guard.h" }
-    module allocator                          { header "__memory/allocator.h" }
+    module allocator {
+      header "__memory/allocator.h"
+      export * // TODO: Workaround for https://github.com/llvm/llvm-project/issues/120108
+    }
     module allocator_arg_t                    { header "__memory/allocator_arg_t.h" }
     module allocator_destructor               { header "__memory/allocator_destructor.h" }
     module allocator_traits                   { header "__memory/allocator_traits.h" }
     module array_cookie                       { header "__memory/array_cookie.h" }
     module assume_aligned                     { header "__memory/assume_aligned.h" }
     module auto_ptr                           { header "__memory/auto_ptr.h" }
-    module builtin_new_allocator              { header "__memory/builtin_new_allocator.h" }
     module compressed_pair                    { header "__memory/compressed_pair.h" }
     module concepts                           { header "__memory/concepts.h" }
     module construct_at                       { header "__memory/construct_at.h" }
@@ -1569,6 +1571,7 @@ module std [system] {
       header "__memory/unique_temporary_buffer.h"
       export std.memory.unique_ptr
       export std_core.type_traits.is_constant_evaluated
+      export * // TODO: Workaround for https://github.com/llvm/llvm-project/issues/120108
     }
     module uses_allocator                     { header "__memory/uses_allocator.h" }
     module uses_allocator_construction        { header "__memory/uses_allocator_construction.h" }
@@ -1604,7 +1607,11 @@ module std [system] {
   module new {
     header "new"
     module align_val_t          { header "__new/align_val_t.h" }
-    module allocate             { header "__new/allocate.h" }
+    module allocate {
+      header "__new/allocate.h"
+      export std.utility.element_count // used as part of the API
+      export * // TODO: Workaround for https://github.com/llvm/llvm-project/issues/120108
+    }
     module destroying_delete_t  { header "__new/destroying_delete_t.h" }
     module exceptions           { header "__new/exceptions.h" }
     module global_new_delete    {
@@ -1911,7 +1918,10 @@ module std [system] {
 
   module string {
     module char_traits              { header "__string/char_traits.h" }
-    module constexpr_c_functions    { header "__string/constexpr_c_functions.h" }
+    module constexpr_c_functions    {
+      header "__string/constexpr_c_functions.h"
+      export std.utility.element_count // used as part of the constexpr C function's API
+    }
     module extern_template_lists    { header "__string/extern_template_lists.h" }
     module fwd                      {  header "__fwd/string.h" }
 
@@ -2021,6 +2031,7 @@ module std [system] {
     }
     module cmp                        { header "__utility/cmp.h" }
     module convert_to_integral        { header "__utility/convert_to_integral.h" }
+    module element_count              { header "__utility/element_count.h" }
     module exception_guard            { header "__utility/exception_guard.h" }
     module exchange                   { header "__utility/exchange.h" }
     module forward_like               { header "__utility/forward_like.h" }
diff --git a/libcxx/src/memory_resource.cpp b/libcxx/src/memory_resource.cpp
diff --git a/libcxx/test/libcxx/language.support/support.dynamic/libcpp_deallocate.sh.cpp b/libcxx/test/libcxx/language.support/support.dynamic/libcpp_deallocate.sh.cpp

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ class _LIBCPP_TEMPLATE_VIS allocator : private __non_trivial_if<!is_void<_Tp>::v`
`102`	`102`	`if (__libcpp_is_constant_evaluated()) {`
`103`	`103`	`return static_cast<_Tp>(::operator new(__n sizeof(_Tp)));`
`104`	`104`	`} else {`
`105`		`- return static_cast<_Tp>(std::__libcpp_allocate(__n sizeof(_Tp), _LIBCPP_ALIGNOF(_Tp)));`
	`105`	`+ return std::__libcpp_allocate<_Tp>(__element_count(__n));`
`106`	`106`	`}`
`107`	`107`	`}`
`108`	`108`
`@@ -117,7 +117,7 @@ class _LIBCPP_TEMPLATE_VIS allocator : private __non_trivial_if<!is_void<_Tp>::v`
`117`	`117`	`if (__libcpp_is_constant_evaluated()) {`
`118`	`118`	`::operator delete(__p);`
`119`	`119`	`} else {`
`120`		`- std::__libcpp_deallocate((void)__p, __n sizeof(_Tp), _LIBCPP_ALIGNOF(_Tp));`
	`120`	`+ std::__libcpp_deallocate<_Tp>(__p, __element_count(__n));`
`121`	`121`	`}`
`122`	`122`	`}`
`123`	`123`
Original file line number	Diff line number	Diff line change
`@@ -40,7 +40,7 @@ struct __temporary_buffer_deleter {`
`40`	`40`	`return;`
`41`	`41`	`}`
`42`	`42`
`43`		`- std::__libcpp_deallocate_unsized((void*)__ptr, _LIBCPP_ALIGNOF(_Tp));`
	`43`	`+ std::__libcpp_deallocate_unsized<_Tp>(__ptr);`
`44`	`44`	`}`
`45`	`45`	`};`
`46`	`46`