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
25namespace 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.
33template <typename Fn>
34class function_ref;
35
36template <typename Ret, typename... Params>
37class 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