MegEngine/imperative/src/impl/function_hook.h

/**
 * \file imperative/src/include/megbrain/imperative/function_hook.h
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#pragma once

#include "megbrain/utils/thin/function.h"

namespace mgb {
namespace imperative {

template <typename TFunction>
class FunctionHook;

template <template <typename> class TFunction, typename TRet, typename... TArgs>
class FunctionHook<TFunction<TRet(TArgs...)>> {
public:
    using FunctionType = TFunction<TRet(TArgs...)>;
    explicit FunctionHook(FunctionType* fptr) : m_fptr{fptr} { m_backup = *fptr; }

public:
    template <
            typename THook,
            typename = std::enable_if_t<
                    std::is_invocable_r_v<TRet, THook, FunctionType, TArgs...>, void>>
    FunctionHook& apply_hook(THook&& hook) {
        // Replace with hooked version
        *m_fptr = [func = *m_fptr,
                   hook = std::forward<THook>(hook)](TArgs... args) -> TRet {
            return hook(func, std::forward<TArgs>(args)...);
        };
        // Convinent for chain call
        return *this;
    }

private:
    FunctionType* m_fptr;
    FunctionType m_backup;

public:
    ~FunctionHook() { *m_fptr = std::move(m_backup); }
};

template <typename TFunction>
auto make_shared_hook(TFunction* fptr) {
    return std::make_shared<FunctionHook<TFunction>>(fptr);
}

}  // namespace imperative
}  // namespace mgb