c10/util/FunctionRef.h

//===- llvm/ADT/STLExtras.h - Useful STL related functions ------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file contains some templates that are useful if you are working with the
// STL at all.
//
// No library is required when using these functions.
//
//===----------------------------------------------------------------------===//

// c10: modified from llvm::function_ref
// c10: added more SFINAE to enable use in overloaded functions

#pragma once

namespace c10 {

/// An efficient, type-erasing, non-owning reference to a callable. This is
/// intended for use as the type of a function parameter that is not used
/// after the function in question returns.
///
/// This class does not own the callable, so it is not in general safe to store
/// a function_ref.
template<typename Fn> class function_ref;

template<typename Ret, typename ...Params>
class function_ref<Ret(Params...)> {
Ret (*callback)(intptr_t callable, Params ...params) = nullptr;
intptr_t callable;

template<typename Callable>
static Ret callback_fn(intptr_t callable, Params ...params) {
    return (*reinterpret_cast<Callable*>(callable))(
        std::forward<Params>(params)...);
}

public:
function_ref() = default;
function_ref(std::nullptr_t) {}

template <typename Callable>
function_ref(Callable &&callable,
            typename std::enable_if<
                !std::is_same<typename std::remove_reference<Callable>::type,
                                function_ref>::value>::type * = nullptr,
            typename std::enable_if<
                 std::is_convertible<
                   typename std::result_of<Callable&&(Params&&...)>::type,
                   Ret>::value>::type * = nullptr)
    : callback(callback_fn<typename std::remove_reference<Callable>::type>),
        callable(reinterpret_cast<intptr_t>(&callable)) {}

Ret operator()(Params ...params) const {
    return callback(callable, std::forward<Params>(params)...);
}

operator bool() const { return callback; }
};
 
}