[libc++] Add __new_at_least

libcxx/include/CMakeLists.txt

@@ -650,11 +650,13 @@ set(files
   __mutex/unique_lock.h
   __new/align_val_t.h
   __new/allocate.h
+  __new/allocation_result.h
   __new/destroying_delete_t.h
   __new/exceptions.h
   __new/global_new_delete.h
   __new/interference_size.h
   __new/launder.h
+  __new/new_at_least.h
   __new/new_handler.h
   __new/nothrow_t.h
   __new/placement_new_delete.h

libcxx/include/__configuration/attributes.h

@@ -473,6 +473,12 @@
 #  define _LIBCPP_DISABLE_POINTER_FIELD_PROTECTION
 #endif
 
+#if __has_cpp_attribute(_Clang::__malloc_span__)
+#  define _LIBCPP_MALLOC_SPAN [[_Clang::__malloc_span__]]
+#else
+#  define _LIBCPP_MALLOC_SPAN
+#endif
+
 // TODO(LLVM 25): Remove this escape hatch
 #ifndef _LIBCPP_DISABLE_UNUSED_STRUCT_WARNINGS
 #  define _LIBCPP_WARN_UNUSED [[__gnu__::__warn_unused__]]

libcxx/include/__configuration/availability.h

@@ -39,6 +39,9 @@
 // in all versions of the library are available.
 #if !_LIBCPP_HAS_VENDOR_AVAILABILITY_ANNOTATIONS
 
+#  define _LIBCPP_INTRODUCED_IN_LLVM_24 1
+#  define _LIBCPP_INTRODUCED_IN_LLVM_24_ATTRIBUTE /* nothing */
+
 #  define _LIBCPP_INTRODUCED_IN_LLVM_23 1
 #  define _LIBCPP_INTRODUCED_IN_LLVM_23_ATTRIBUTE /* nothing */
 
@@ -73,6 +76,11 @@
 
 // clang-format off
 
+// LLVM 24
+// TODO: Fill this in
+#  define _LIBCPP_INTRODUCED_IN_LLVM_24 0
+#  define _LIBCPP_INTRODUCED_IN_LLVM_24_ATTRIBUTE __attribute__((unavailable))
+
 // LLVM 23
 // TODO: Fill this in
 #  define _LIBCPP_INTRODUCED_IN_LLVM_23 0
@@ -253,6 +261,10 @@
 
 #endif
 
+// This determines whether `std::__new_at_least` is available in the dylib, which
+// is used for std::allocator::allocate_at_least.
+#define _LIBCPP_AVAILABILITY_HAS_NEW_AT_LEAST _LIBCPP_INTRODUCED_IN_LLVM_24
+
 // This determines whether we assume that the internal std::__bad_variant_access_with_msg class
 // (which carries a message describing the cause of the failure in bad_variant_access::what())
 // provides a key function in the dylib. This allows centralizing its vtable and typeinfo instead

libcxx/include/__memory/allocate_at_least.h

@@ -11,26 +11,22 @@
 
 #include <__config>
 #include <__cstddef/size_t.h>
+#include <__fwd/memory.h>
 #include <__memory/allocator_traits.h>
+#include <__new/allocate.h>
+#include <__type_traits/is_constant_evaluated.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
 #endif
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
-template <class _Pointer, class _SizeT = size_t>
-struct __allocation_result {
-  _Pointer ptr;
-  _SizeT count;
-
-  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR __allocation_result(_Pointer __ptr, _SizeT __count)
-      : ptr(__ptr), count(__count) {}
-};
-_LIBCPP_CTAD_SUPPORTED_FOR_TYPE(__allocation_result);
-
 #if _LIBCPP_STD_VER >= 23
 
+// This function allocates memory using the allocator's allocate_at_least member if possible, and falls back the normal
+// allocate in older modes.
+
 template <class _Alloc>
 [[nodiscard]] _LIBCPP_HIDE_FROM_ABI constexpr auto __allocate_at_least(_Alloc& __alloc, size_t __n) {
   auto __res = std::allocator_traits<_Alloc>::allocate_at_least(__alloc, __n);
@@ -46,6 +42,17 @@ __allocate_at_least(_Alloc& __alloc, size_t __n) {
   return __allocation_result<typename _Traits::pointer, typename _Traits::size_type>(__alloc.allocate(__n), __n);
 }
 
+// Provide an efficient __allocate_at_least for std::allocator in all standard modes
+
+template <class _Tp>
+[[__nodiscard__]] _LIBCPP_CONSTEXPR __allocation_result<_Tp*> __allocate_at_least(allocator<_Tp>& __alloc, size_t __n) {
+  if (__libcpp_is_constant_evaluated()) {
+    return __allocation_result<_Tp*>(__alloc.allocate(__n), __n);
+  } else {
+    return std::__libcpp_allocate_at_least<_Tp>(__element_count(__n));
+  }
+}
+
 #endif // _LIBCPP_STD_VER >= 23
 
 _LIBCPP_END_NAMESPACE_STD

libcxx/include/__memory/allocator.h

@@ -96,7 +96,12 @@ class allocator
 #if _LIBCPP_STD_VER >= 23
   [[nodiscard]] _LIBCPP_HIDE_FROM_ABI constexpr allocation_result<_Tp*> allocate_at_least(size_t __n) {
     static_assert(sizeof(_Tp) >= 0, "cannot allocate memory for an incomplete type");
-    return {allocate(__n), __n};
+    if consteval {
+      return {allocate(__n), __n};
+    } else {
+      auto [__ptr, __count] = std::__libcpp_allocate_at_least<_Tp>(__element_count(__n));
+      return {__ptr, __count};
+    }
   }
 #endif
 

libcxx/include/__new/allocate.h

@@ -13,8 +13,11 @@
 #include <__cstddef/max_align_t.h>
 #include <__cstddef/size_t.h>
 #include <__new/align_val_t.h>
+#include <__new/allocation_result.h>
+#include <__new/new_at_least.h>
 #include <__type_traits/type_identity.h>
 #include <__utility/element_count.h>
+#include <cstdint>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
@@ -42,6 +45,34 @@ __libcpp_allocate(__element_count __n, [[__maybe_unused__]] size_t __align = _LI
   return static_cast<_Tp*>(__builtin_operator_new(__size));
 }
 
+template <class _Tp>
+_LIBCPP_NO_CFI __allocation_result<_Tp*>
+__libcpp_allocate_at_least(__element_count __n, [[__maybe_unused__]] size_t __align = _LIBCPP_ALIGNOF(_Tp)) {
+#if _LIBCPP_AVAILABILITY_HAS_NEW_AT_LEAST
+  size_t __size = static_cast<size_t>(__n) * sizeof(_Tp);
+#  if _LIBCPP_HAS_ALIGNED_ALLOCATION
+  if (__is_overaligned_for_new(__align)) {
+    auto [__ptr, __count] = std::__new_at_least(__size, static_cast<align_val_t>(__align));
+#    if __has_builtin(__builtin_assume)
+    __builtin_assume(__ptr);
+    __builtin_assume((reinterpret_cast<uintptr_t>(__ptr) & (__align - 1)) == 0);
+    __builtin_assume(__count >= __size);
+#    endif
+    return {static_cast<_Tp*>(__ptr), __count / sizeof(_Tp)};
+  }
+#  endif
+  auto [__ptr, __count] = std::__new_at_least(__size);
+#  if __has_builtin(__builtin_assume)
+  __builtin_assume(__ptr);
+  __builtin_assume((reinterpret_cast<uintptr_t>(__ptr) & (__align - 1)) == 0);
+  __builtin_assume(__count >= __size);
+#  endif
+  return __allocation_result<_Tp*>(static_cast<_Tp*>(__ptr), __count / sizeof(_Tp));
+#else
+  return __allocation_result<_Tp*>(std::__libcpp_allocate<_Tp>(__n, __align), __n);
+#endif
+}
+
 #if defined(__cpp_sized_deallocation) && __cpp_sized_deallocation >= 201309L
 #  define _LIBCPP_ONLY_IF_SIZED_DEALLOCATION(...) __VA_ARGS__
 #else

libcxx/include/__new/allocation_result.h

@@ -0,0 +1,33 @@
+//===----------------------------------------------------------------------===//
+//
+// 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___NEW_ALLOCATION_RESULT_H
+#define _LIBCPP___NEW_ALLOCATION_RESULT_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
+
+template <class _Pointer, class _SizeT = size_t>
+struct __allocation_result {
+  _Pointer ptr;
+  _SizeT count;
+
+  _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR __allocation_result(_Pointer __ptr, _SizeT __count)
+      : ptr(__ptr), count(__count) {}
+};
+_LIBCPP_CTAD_SUPPORTED_FOR_TYPE(__allocation_result);
+
+_LIBCPP_END_NAMESPACE_STD
+
+#endif // _LIBCPP___NEW_ALLOCATION_RESULT_H

libcxx/include/__new/new_at_least.h

@@ -0,0 +1,41 @@
+//===----------------------------------------------------------------------===//
+//
+// 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___NEW_NEW_AT_LEAST_H
+#define _LIBCPP___NEW_NEW_AT_LEAST_H
+
+#include <__config>
+#include <__cstddef/size_t.h>
+#include <__new/align_val_t.h>
+#include <__new/allocation_result.h>
+
+#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
+#  pragma GCC system_header
+#endif
+
+#if _LIBCPP_AVAILABILITY_HAS_NEW_AT_LEAST
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+_LIBCPP_BEGIN_EXPLICIT_ABI_ANNOTATIONS
+
+// `__new_at_least` acts like an overload of `operator new` which takes a size (and possibly alignment) and returns
+// a pointer as well as the actually allocated amount of memory. If the user replaces the relevant `operator new`
+// overload this will fall back to calling that. Otherwise it tries to allocate in a way to get the actually allocated
+// size, depending on what the platform provides.
+
+_LIBCPP_MALLOC_SPAN _LIBCPP_EXPORTED_FROM_ABI __allocation_result<void*> __new_at_least(size_t);
+#  if _LIBCPP_HAS_LIBRARY_ALIGNED_ALLOCATION
+_LIBCPP_MALLOC_SPAN _LIBCPP_EXPORTED_FROM_ABI __allocation_result<void*> __new_at_least(size_t, align_val_t);
+#  endif
+
+_LIBCPP_END_EXPLICIT_ABI_ANNOTATIONS
+_LIBCPP_END_NAMESPACE_STD
+
+#endif // _LIBCPP_AVAILABILITY_HAS_NEW_AT_LEAST
+
+#endif // _LIBCPP___NEW_NEW_AT_LEAST_H

libcxx/include/module.modulemap.in

@@ -1785,6 +1785,7 @@ module std {
       export std.utility.element_count // used as part of the API
       export * // TODO: Workaround for https://llvm.org/PR120108
     }
+    module allocation_result    { header "__new/allocation_result.h" }
     module destroying_delete_t  { header "__new/destroying_delete_t.h" }
     module exceptions           { header "__new/exceptions.h" }
     module global_new_delete    {
@@ -1794,6 +1795,7 @@ module std {
     }
     module interference_size    { header "__new/interference_size.h" }
     module launder              { header "__new/launder.h" }
+    module new_at_least         { header "__new/new_at_least.h" }
     module new_handler          { header "__new/new_handler.h" }
     module nothrow_t            { header "__new/nothrow_t.h" }
     module placement_new_delete { header "__new/placement_new_delete.h" }

libcxx/lib/abi/CHANGELOG.TXT

@@ -16,6 +16,12 @@ New entries should be added directly below the "Version" header.
 Version 24.0
 ------------
 
+* [libc++] Add __new_at_least
+
+  All platforms
+  -------------
+  Symbol added: _ZNSt3__114__new_at_leastEm
+
 * [libc++] Remove locale::id::__next_id from the ABI
 
   `__next_id` is only ever used inside the dylib, so we can avoid making it public.

libcxx/lib/abi/x86_64-unknown-linux-gnu.libcxxabi.v1.stable.exceptions.nonew.abilist

@@ -9,6 +9,7 @@
 {'is_defined': False, 'name': '_ZNSt16invalid_argumentD1Ev', 'type': 'FUNC'}
 {'is_defined': False, 'name': '_ZNSt20bad_array_new_lengthC1Ev', 'type': 'FUNC'}
 {'is_defined': False, 'name': '_ZNSt20bad_array_new_lengthD1Ev', 'type': 'FUNC'}
+{'is_defined': False, 'name': '_ZNSt3__114__new_at_leastEm', 'type': 'FUNC'}
 {'is_defined': False, 'name': '_ZNSt8bad_castC1Ev', 'type': 'FUNC'}
 {'is_defined': False, 'name': '_ZNSt8bad_castD1Ev', 'type': 'FUNC'}
 {'is_defined': False, 'name': '_ZNSt8bad_castD2Ev', 'type': 'FUNC'}

libcxx/src/support/new.ipp

@@ -6,6 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <__memory/allocate_at_least.h>
+#include <__new/new_at_least.h>
 #include <cstddef>
 #include <cstdlib>
 #include <new>
@@ -216,3 +218,90 @@ OVERRIDABLE_FUNCTION void* operator new[](size_t size, std::align_val_t alignmen
   ::operator delete[](ptr, alignment);
 }
 #endif // _LIBCPP_HAS_LIBRARY_ALIGNED_ALLOCATION
+
+// This part implements __new_at_least, a version of operator new that returns the actually allocated amount of memory
+// in addition to the pointer. Since users are allowed to replace operator new, we have to check whether it is replaced
+// and fall back to that. Otherwise we can use platform-specific APIs for an improved implementation.
+//
+// We do that check via `gnu::ifunc` if it's available. Otherwise we don't do anything and just unconditionally forward
+// to `operator new(size_t{, align_val_t})`. `gnu::ifunc` takes the mangled name of a resolver function. That resolver
+// function returns a pointer to the function that should be linked. We make use of `__is_function_overridden` to detect
+// whether we can use our own special implementation or have to fall back to a user-provided operator new. This approach
+// avoids any repeated checks in a very hot path.
+
+#ifdef __APPLE__
+#  include <malloc/malloc.h>
+#elifdef __FreeBSD__
+#  include <malloc_np.h>
+#endif
+
+// FIXME: Clang should really accept functions in [[gnu::ifunc]] (or possibly [[clang::ifunc]])
+using std::__allocation_result;
+
+using new_t         = void*(std::size_t);
+using new_aligned_t = void*(std::size_t, std::align_val_t);
+
+using new_at_least_t         = __allocation_result<void*>(std::size_t);
+using new_at_least_aligned_t = __allocation_result<void*>(std::size_t, std::align_val_t);
+
+[[maybe_unused]] static new_at_least_t* new_at_least_resolver() {
+  if (std::__is_function_overridden < new_t, operator new>()) {
+    return [](std::size_t size) -> __allocation_result<void*> { return {::operator new(size), size}; };
+  } else {
+    return [](std::size_t size) -> __allocation_result<void*> {
+#ifdef __APPLE__
+      auto good_size = ::malloc_good_size(size);
+      return {operator_new_impl<on_failure::throw_bad_alloc>(good_size), good_size};
+#elifdef __FreeBSD__
+      auto good_size = ::nallocx(size, 0);
+      return {operator_new_impl<on_failure::throw_bad_alloc>(good_size), good_size};
+#else
+      // Other platforms should specialize this for their system allocator
+      return {operator_new_impl<on_failure::throw_bad_alloc>(size), size};
+#endif
+    };
+  }
+}
+
+[[maybe_unused]] static new_at_least_aligned_t* new_at_least_aligned_resolver() {
+  if (std::__is_function_overridden < new_aligned_t, operator new>()) {
+    return [](std::size_t size, std::align_val_t align) -> __allocation_result<void*> {
+      return {::operator new(size, align), size};
+    };
+  } else {
+    return [](std::size_t size, std::align_val_t align) -> __allocation_result<void*> {
+#ifdef __APPLE__
+      auto good_size = ::malloc_good_size(size);
+      return {operator_new_aligned_impl<on_failure::throw_bad_alloc>(good_size, align), good_size};
+#elifdef __FreeBSD__
+      auto good_size = ::nallocx(size, MALLOCX_ALIGN(static_cast<size_t>(align)));
+      return {operator_new_aligned_impl<on_failure::throw_bad_alloc>(good_size), good_size};
+#else
+      return {operator_new_aligned_impl<on_failure::throw_bad_alloc>(size, align), size};
+#endif
+    };
+  }
+}
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+_LIBCPP_BEGIN_EXPLICIT_ABI_ANNOTATIONS
+
+#if __has_cpp_attribute(gnu::ifunc)
+
+[[gnu::ifunc("_ZL21new_at_least_resolverv")]] new_at_least_t __new_at_least;
+[[gnu::ifunc("_ZL29new_at_least_aligned_resolverv")]] new_at_least_aligned_t __new_at_least;
+
+#else
+
+std::__allocation_result<void*> __new_at_least(std::size_t size) { return {::operator new(size), size}; }
+
+#  if _LIBCPP_HAS_LIBRARY_ALIGNED_ALLOCATION
+std::__allocation_result<void*> __new_at_least(std::size_t size, std::align_val_t align) {
+  return {::operator new(size, align), size};
+}
+#  endif
+
+#endif
+
+_LIBCPP_END_EXPLICIT_ABI_ANNOTATIONS
+_LIBCPP_END_NAMESPACE_STD