1 | //===- llvm/ADT/STLExtras.h - Useful STL related functions ------*- C++ -*-===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This file contains some templates that are useful if you are working with the |
10 | // STL at all. |
11 | // |
12 | // No library is required when using these functions. |
13 | // |
14 | //===----------------------------------------------------------------------===// |
15 | |
16 | // c10: modified from llvm::function_ref |
17 | // c10: added more SFINAE to enable use in overloaded functions |
18 | |
19 | #pragma once |
20 | |
21 | #include <cstdint> |
22 | #include <type_traits> |
23 | #include <utility> |
24 | |
25 | namespace c10 { |
26 | |
27 | /// An efficient, type-erasing, non-owning reference to a callable. This is |
28 | /// intended for use as the type of a function parameter that is not used |
29 | /// after the function in question returns. |
30 | /// |
31 | /// This class does not own the callable, so it is not in general safe to store |
32 | /// a function_ref. |
33 | template <typename Fn> |
34 | class function_ref; |
35 | |
36 | template <typename Ret, typename... Params> |
37 | class function_ref<Ret(Params...)> { |
38 | Ret (*callback)(intptr_t callable, Params... params) = nullptr; |
39 | intptr_t callable; |
40 | |
41 | template <typename Callable> |
42 | static Ret callback_fn(intptr_t callable, Params... params) { |
43 | return (*reinterpret_cast<Callable*>(callable))(std::forward<Params>( |
44 | params)...); |
45 | } |
46 | |
47 | public: |
48 | function_ref() = default; |
49 | function_ref(std::nullptr_t) {} |
50 | |
51 | template <typename Callable> |
52 | function_ref( |
53 | Callable&& callable, |
54 | typename std::enable_if<!std::is_same< |
55 | typename std::remove_reference<Callable>::type, |
56 | function_ref>::value>::type* = nullptr, |
57 | typename std::enable_if<std::is_convertible< |
58 | typename c10::invoke_result_t<Callable, Params...>, |
59 | Ret>::value>::type* = nullptr) |
60 | : callback(callback_fn<typename std::remove_reference<Callable>::type>), |
61 | callable(reinterpret_cast<intptr_t>(&callable)) {} |
62 | |
63 | Ret operator()(Params... params) const { |
64 | return callback(callable, std::forward<Params>(params)...); |
65 | } |
66 | |
67 | operator bool() const { |
68 | return callback; |
69 | } |
70 | }; |
71 | |
72 | } // namespace c10 |
73 | |