diff --git a/stl/inc/functional b/stl/inc/functional index 65ad553b2da..2a71748f412 100644 --- a/stl/inc/functional +++ b/stl/inc/functional @@ -1629,15 +1629,37 @@ public: _STL_INTERNAL_STATIC_ASSERT(alignof(max_align_t) == alignof(void*)); // 64-bit target, can put small function into small move_only_function directly _Data._Impl = _Create_impl_ptr<_Impl_kind::_Old_fn_small, _Vt, void>(); - _Old_fn_impl->_Move(_Data._Buf_ptr()); - _Func._Tidy(); + if constexpr (is_lvalue_reference_v<_Fn>) { + _Old_fn_impl->_Copy(_Data._Buf_ptr()); + } else { + _Old_fn_impl->_Move(_Data._Buf_ptr()); + _Func._Tidy(); + } #else // ^^^ 64-bit / 32-bit vvv _STL_INTERNAL_STATIC_ASSERT(alignof(max_align_t) > alignof(void*)); // 32-bit target, cannot put small function into small move_only_function directly // due to potentially not enough alignment. Allocate large function void* _Where = ::operator new((_Small_object_num_ptrs - 1) * sizeof(void*)); - _Old_fn_impl->_Move(_Where); - _Func._Tidy(); + if constexpr (is_lvalue_reference_v<_Fn>) { + struct _NODISCARD _Guard_type { + void* _Ptr; + + ~_Guard_type() { + // _Ptr is not nullptr only if an exception is thrown as a result of _Copy. + // Check _Ptr before calling operator delete to save the call in the common case. + if (_Ptr) { + ::operator delete(_Ptr); + } + } + }; + + _Guard_type _Guard{_Where}; + _Old_fn_impl->_Copy(_Where); + _Guard._Ptr = nullptr; + } else { + _Old_fn_impl->_Move(_Where); + _Func._Tidy(); + } _Data._Impl = _Create_impl_ptr<_Impl_kind::_Old_fn_small_as_large, _Vt, void>(); _Data._Set_large_fn_ptr(_Where); @@ -1645,8 +1667,12 @@ public: } else { // Just take ownership of the inner impl pointer _Data._Impl = _Create_impl_ptr<_Impl_kind::_Old_fn_large, _Vt, void>(); - _Data._Set_large_fn_ptr(_Old_fn_impl); - _Func._Set(nullptr); + if constexpr (is_lvalue_reference_v<_Fn>) { + _Data._Set_large_fn_ptr(_Old_fn_impl->_Copy(nullptr)); + } else { + _Data._Set_large_fn_ptr(_Old_fn_impl); + _Func._Set(nullptr); + } } } diff --git a/tests/std/tests/GH_005504_avoid_function_call_wrapping/test.cpp b/tests/std/tests/GH_005504_avoid_function_call_wrapping/test.cpp index 44ec403e220..1102e74009b 100644 --- a/tests/std/tests/GH_005504_avoid_function_call_wrapping/test.cpp +++ b/tests/std/tests/GH_005504_avoid_function_call_wrapping/test.cpp @@ -104,6 +104,14 @@ void test_wrapped_call(const int expected_copies) { assert(outer(copy_counter{}) == expected_copies); } +template +void test_wrapped_copy_call(const int expected_copies) { + InnerWrapper inner{Callable{}}; + OuterWrapper outer{inner}; + assert(inner); + assert(outer(copy_counter{}) == expected_copies); +} + template void check_call_null(Wrapper& wrapper, const bool throws) { if (throws) { @@ -143,18 +151,22 @@ int main() { // Moves to the same alloc_checker{0}, test_wrapped_call, function, small_callable>(0); alloc_checker{1}, test_wrapped_call, function, large_callable>(0); + alloc_checker{0}, test_wrapped_copy_call, function, small_callable>(0); + alloc_checker{2}, test_wrapped_copy_call, function, large_callable>(0); alloc_checker{0}, test_wrapped_call, move_only_function, small_callable>(0); alloc_checker{1}, test_wrapped_call, move_only_function, large_callable>(0); + constexpr bool is_64_bit = sizeof(void*) > 4; + // Moves from function to move_only_function -#ifdef _WIN64 - alloc_checker{0}, -#else - alloc_checker{1}, -#endif + alloc_checker{is_64_bit ? 0 : 1}, test_wrapped_call, function, small_callable>(0); alloc_checker{1}, test_wrapped_call, function, large_callable>(0); + alloc_checker{is_64_bit ? 0 : 1}, + test_wrapped_copy_call, function, small_callable>(0); + alloc_checker{2}, test_wrapped_copy_call, function, large_callable>(0); + // nulls alloc_checker{0}, test_plain_null>(true); alloc_checker{0}, test_plain_null>(false);