FunctionRef.h 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  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. // c10: modified from llvm::function_ref
  16. // c10: added more SFINAE to enable use in overloaded functions
  17. #pragma once
  18. #include <cstdint>
  19. #include <type_traits>
  20. #include <utility>
  21. namespace c10 {
  22. /// An efficient, type-erasing, non-owning reference to a callable. This is
  23. /// intended for use as the type of a function parameter that is not used
  24. /// after the function in question returns.
  25. ///
  26. /// This class does not own the callable, so it is not in general safe to store
  27. /// a function_ref.
  28. template <typename Fn>
  29. class function_ref;
  30. template <typename Ret, typename... Params>
  31. class function_ref<Ret(Params...)> {
  32. Ret (*callback)(intptr_t callable, Params... params) = nullptr;
  33. intptr_t callable;
  34. template <typename Callable>
  35. static Ret callback_fn(intptr_t callable, Params... params) {
  36. return (*reinterpret_cast<Callable*>(callable))(std::forward<Params>(
  37. params)...);
  38. }
  39. public:
  40. function_ref() = default;
  41. function_ref(std::nullptr_t) {}
  42. template <typename Callable>
  43. function_ref(
  44. Callable&& callable,
  45. typename std::enable_if<!std::is_same<
  46. typename std::remove_reference<Callable>::type,
  47. function_ref>::value>::type* = nullptr,
  48. typename std::enable_if<std::is_convertible<
  49. typename c10::invoke_result_t<Callable, Params...>,
  50. Ret>::value>::type* = nullptr)
  51. : callback(callback_fn<typename std::remove_reference<Callable>::type>),
  52. callable(reinterpret_cast<intptr_t>(&callable)) {}
  53. Ret operator()(Params... params) const {
  54. return callback(callable, std::forward<Params>(params)...);
  55. }
  56. operator bool() const {
  57. return callback;
  58. }
  59. };
  60. } // namespace c10