SafeFunction.h 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  1. /*
  2. * Copyright (c) 2016 Apple Inc. All rights reserved.
  3. * Copyright (c) 2021, Gunnar Beutner <gbeutner@serenityos.org>
  4. * Copyright (c) 2022, Andreas Kling <kling@serenityos.org>
  5. *
  6. * SPDX-License-Identifier: BSD-2-Clause
  7. */
  8. #pragma once
  9. #include <AK/Function.h>
  10. #include <AK/SourceLocation.h>
  11. namespace JS {
  12. void register_safe_function_closure(void*, size_t, SourceLocation*);
  13. void unregister_safe_function_closure(void*, size_t, SourceLocation*);
  14. template<typename>
  15. class SafeFunction;
  16. template<typename Out, typename... In>
  17. class SafeFunction<Out(In...)> {
  18. AK_MAKE_NONCOPYABLE(SafeFunction);
  19. public:
  20. SafeFunction() = default;
  21. SafeFunction(nullptr_t)
  22. {
  23. }
  24. ~SafeFunction()
  25. {
  26. clear(false);
  27. }
  28. void register_closure()
  29. {
  30. if (!m_size)
  31. return;
  32. if (auto* wrapper = callable_wrapper())
  33. register_safe_function_closure(wrapper, m_size, &m_location);
  34. }
  35. void unregister_closure()
  36. {
  37. if (!m_size)
  38. return;
  39. if (auto* wrapper = callable_wrapper())
  40. unregister_safe_function_closure(wrapper, m_size, &m_location);
  41. }
  42. template<typename CallableType>
  43. SafeFunction(CallableType&& callable, SourceLocation location = SourceLocation::current())
  44. requires((AK::IsFunctionObject<CallableType> && IsCallableWithArguments<CallableType, Out, In...> && !IsSame<RemoveCVReference<CallableType>, SafeFunction>))
  45. : m_location(location)
  46. {
  47. init_with_callable(forward<CallableType>(callable), CallableKind::FunctionObject);
  48. }
  49. template<typename FunctionType>
  50. SafeFunction(FunctionType f, SourceLocation location = SourceLocation::current())
  51. requires((AK::IsFunctionPointer<FunctionType> && IsCallableWithArguments<RemovePointer<FunctionType>, Out, In...> && !IsSame<RemoveCVReference<FunctionType>, SafeFunction>))
  52. : m_location(location)
  53. {
  54. init_with_callable(move(f), CallableKind::FunctionPointer);
  55. }
  56. SafeFunction(SafeFunction&& other)
  57. : m_location(move(other.m_location))
  58. {
  59. move_from(move(other));
  60. }
  61. // Note: Despite this method being const, a mutable lambda _may_ modify its own captures.
  62. Out operator()(In... in) const
  63. {
  64. auto* wrapper = callable_wrapper();
  65. VERIFY(wrapper);
  66. ++m_call_nesting_level;
  67. ScopeGuard guard([this] {
  68. if (--m_call_nesting_level == 0 && m_deferred_clear)
  69. const_cast<SafeFunction*>(this)->clear(false);
  70. });
  71. return wrapper->call(forward<In>(in)...);
  72. }
  73. explicit operator bool() const { return !!callable_wrapper(); }
  74. SafeFunction& operator=(nullptr_t)
  75. {
  76. clear();
  77. return *this;
  78. }
  79. SafeFunction& operator=(SafeFunction&& other)
  80. {
  81. if (this != &other) {
  82. clear();
  83. move_from(move(other));
  84. }
  85. return *this;
  86. }
  87. private:
  88. enum class CallableKind {
  89. FunctionPointer,
  90. FunctionObject,
  91. };
  92. class CallableWrapperBase {
  93. public:
  94. virtual ~CallableWrapperBase() = default;
  95. // Note: This is not const to allow storing mutable lambdas.
  96. virtual Out call(In...) = 0;
  97. virtual void destroy() = 0;
  98. virtual void init_and_swap(u8*, size_t) = 0;
  99. };
  100. template<typename CallableType>
  101. class CallableWrapper final : public CallableWrapperBase {
  102. AK_MAKE_NONMOVABLE(CallableWrapper);
  103. AK_MAKE_NONCOPYABLE(CallableWrapper);
  104. public:
  105. explicit CallableWrapper(CallableType&& callable)
  106. : m_callable(move(callable))
  107. {
  108. }
  109. Out call(In... in) final override
  110. {
  111. return m_callable(forward<In>(in)...);
  112. }
  113. void destroy() final override
  114. {
  115. delete this;
  116. }
  117. // NOLINTNEXTLINE(readability-non-const-parameter) False positive; destination is used in a placement new expression
  118. void init_and_swap(u8* destination, size_t size) final override
  119. {
  120. VERIFY(size >= sizeof(CallableWrapper));
  121. new (destination) CallableWrapper { move(m_callable) };
  122. }
  123. private:
  124. CallableType m_callable;
  125. };
  126. enum class FunctionKind {
  127. NullPointer,
  128. Inline,
  129. Outline,
  130. };
  131. CallableWrapperBase* callable_wrapper() const
  132. {
  133. switch (m_kind) {
  134. case FunctionKind::NullPointer:
  135. return nullptr;
  136. case FunctionKind::Inline:
  137. return bit_cast<CallableWrapperBase*>(&m_storage);
  138. case FunctionKind::Outline:
  139. return *bit_cast<CallableWrapperBase**>(&m_storage);
  140. default:
  141. VERIFY_NOT_REACHED();
  142. }
  143. }
  144. void clear(bool may_defer = true)
  145. {
  146. bool called_from_inside_function = m_call_nesting_level > 0;
  147. // NOTE: This VERIFY could fail because a Function is destroyed from within itself.
  148. VERIFY(may_defer || !called_from_inside_function);
  149. if (called_from_inside_function && may_defer) {
  150. m_deferred_clear = true;
  151. return;
  152. }
  153. m_deferred_clear = false;
  154. auto* wrapper = callable_wrapper();
  155. if (m_kind == FunctionKind::Inline) {
  156. VERIFY(wrapper);
  157. wrapper->~CallableWrapperBase();
  158. unregister_closure();
  159. } else if (m_kind == FunctionKind::Outline) {
  160. VERIFY(wrapper);
  161. wrapper->destroy();
  162. unregister_closure();
  163. }
  164. m_kind = FunctionKind::NullPointer;
  165. }
  166. template<typename Callable>
  167. void init_with_callable(Callable&& callable, CallableKind kind)
  168. {
  169. VERIFY(m_call_nesting_level == 0);
  170. VERIFY(m_kind == FunctionKind::NullPointer);
  171. using WrapperType = CallableWrapper<Callable>;
  172. if constexpr (sizeof(WrapperType) > inline_capacity) {
  173. *bit_cast<CallableWrapperBase**>(&m_storage) = new WrapperType(forward<Callable>(callable));
  174. m_kind = FunctionKind::Outline;
  175. } else {
  176. new (m_storage) WrapperType(forward<Callable>(callable));
  177. m_kind = FunctionKind::Inline;
  178. }
  179. if (kind == CallableKind::FunctionObject)
  180. m_size = sizeof(WrapperType);
  181. else
  182. m_size = 0;
  183. register_closure();
  184. }
  185. void move_from(SafeFunction&& other)
  186. {
  187. VERIFY(m_call_nesting_level == 0);
  188. VERIFY(other.m_call_nesting_level == 0);
  189. VERIFY(m_kind == FunctionKind::NullPointer);
  190. auto* other_wrapper = other.callable_wrapper();
  191. m_size = other.m_size;
  192. AK::TypedTransfer<SourceLocation>::move(&m_location, &other.m_location, 1);
  193. switch (other.m_kind) {
  194. case FunctionKind::NullPointer:
  195. break;
  196. case FunctionKind::Inline:
  197. other.unregister_closure();
  198. other_wrapper->init_and_swap(m_storage, inline_capacity);
  199. m_kind = FunctionKind::Inline;
  200. register_closure();
  201. break;
  202. case FunctionKind::Outline:
  203. other.unregister_closure();
  204. *bit_cast<CallableWrapperBase**>(&m_storage) = other_wrapper;
  205. m_kind = FunctionKind::Outline;
  206. register_closure();
  207. break;
  208. default:
  209. VERIFY_NOT_REACHED();
  210. }
  211. other.m_kind = FunctionKind::NullPointer;
  212. }
  213. FunctionKind m_kind { FunctionKind::NullPointer };
  214. bool m_deferred_clear { false };
  215. mutable Atomic<u16> m_call_nesting_level { 0 };
  216. size_t m_size { 0 };
  217. SourceLocation m_location;
  218. // Empirically determined to fit most lambdas and functions.
  219. static constexpr size_t inline_capacity = 4 * sizeof(void*);
  220. alignas(max(alignof(CallableWrapperBase), alignof(CallableWrapperBase*))) u8 m_storage[inline_capacity];
  221. };
  222. }