// Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//

//! @file NodeDefLambda.h
//!
//! @brief Defines @ref omni::graph::exec::unstable::NodeDefLambda.
#pragma once

#include <omni/graph/exec/unstable/NodeDef.h>

#include <functional>

namespace omni
{
namespace graph
{
namespace exec
{
namespace unstable
{

//! Wrapper of a lambda function into opaque node definition
//!
//! This class is great for quick prototyping and nodes that won't be statically analyzed. It is recommended to convert
//! such nodes into real classes to avoid paying the price of std::function call from already virtual method and to
//! provide type information.
// ef-docs node-def-lambda-begin
class NodeDefLambda : public NodeDef
{
public:
    //! Templated constructor for wrapper class
    //!
    //! @tparam Fn  Function type taking execution task and returning status. Compiler should detect it for you.
    //! @param  definitionName  Definition name is considered as a token that transformation passes can register against
    //! @param  fn    Execute function body.  Signature should be `Status(ExecutionTask&)`.
    //! @param  schedInfo  Fixed at runtime scheduling constraint
    //!
    //! May throw.
    template <typename Fn>
    static omni::core::ObjectPtr<NodeDefLambda> create(const char* definitionName, Fn&& fn, SchedulingInfo schedInfo)
    {
        OMNI_THROW_IF_ARG_NULL(definitionName);
        return omni::core::steal(new NodeDefLambda(definitionName, std::forward<Fn>(fn), schedInfo));
    }

protected:
    //! Templated and protected constructor for wrapper class.
    //!
    //! Use factory method to construct objects of this class.
    template <typename Fn>
    NodeDefLambda(const char* definitionName, Fn&& fn, SchedulingInfo schedInfo)
        : NodeDef(definitionName), m_fn(std::move(fn)), m_schedulingInfo(schedInfo)
    {
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IDef::execute_abi for @ref NodeDefLambda
    //!
    //! Calling captured lambda
    Status execute_abi(ExecutionTask* info) noexcept override
    {
        return m_fn(*info);
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IDef::getSchedulingInfo_abi for @ref NodeDefLambda
    //!
    //! Returns scheduling information provided to the constructor
    SchedulingInfo getSchedulingInfo_abi(const ExecutionTask* info) noexcept override
    {
        return m_schedulingInfo;
    }

private:
    std::function<Status(ExecutionTask&)> m_fn; //!< Execute function body
    SchedulingInfo m_schedulingInfo; //!< Scheduling constraint
};
// ef-docs node-def-lambda-end

} // namespace unstable
} // namespace exec
} // namespace graph
} // namespace omni
