// 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 IPartitionPass.h
//!
//! @brief Defines @ref omni::graph::exec::unstable::IPartitionPass.
#pragma once

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

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

// forward declarations needed by interface declaration
class IGraphBuilder;
class INode;
class ITopology;
class IPartitionPass;
class IPartitionPass_abi;

//! Base class for graph partitioning passes.
//!
//! Partition passes are typically run just after population passes and only on newly modified
//! @ref omni::graph::exec::unstable::INodeGraphDef objects. The job of a partition pass is to recognize patterns in the
//! newly populated graph and replace them with a new definition or augment existing one.
//!
//! Partition passes can only mutate the graph from the @ref omni::graph::exec::unstable::IPartitionPass::commit method
//! using provided @ref omni::graph::exec::unstable::IGraphBuilder. This will guarantee that the rest of the pipeline
//! is aware of changes made to the graph and avoid potential threading issues.
//!
//! See @ref  groupOmniGraphExecPasses for more pass related functionality.
class IPartitionPass_abi : public omni::core::Inherits<IPass, OMNI_TYPE_ID("omni.graph.exec.unstable.IPartitionPass")>
{
protected:
    //! Call from pass pipeline to initialize the pass for @p topology.
    //!
    //! This interface method implementation can't mutate given @p topology. Multiple passes can run concurrently on it.
    //!
    //! Returns True if initialization was successful and pipeline should issue calls to run and commit.
    //! Otherwise this pass will be destroyed and won't participate in partitioning @p topology.
    virtual bool initialize_abi(OMNI_ATTR("not_null, throw_if_null") ITopology* topology) noexcept = 0;

    //! Call from pass pipeline to discover nodes requiring partitioning.
    //!
    //! No topology changes are permitted at this point. Multiple passes will get a chance to receive this
    //! notification.
    //!
    //! Call to this method comes from graph traversal that may run multiple passes concurrently.
    virtual void run_abi(OMNI_ATTR("not_null, throw_if_null") INode* node) noexcept = 0;

    //! Call to verify generated partitions and commit new definition/s replacing discovered partitions.
    //!
    //! Commit of partitions is done serially and in the priority order of the pass. Passes with higher order will get
    //! the chance first. This is the only partition pass method that can mutate the graph.
    virtual void commit_abi(OMNI_ATTR("not_null, throw_if_null") IGraphBuilder* builder) noexcept = 0;
};

//! Smart pointer managing an instance of @ref IPartitionPass.
using PartitionPassPtr = omni::core::ObjectPtr<IPartitionPass>;

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

// generated API declaration
#define OMNI_BIND_INCLUDE_INTERFACE_DECL
#include <omni/graph/exec/unstable/IPartitionPass.gen.h>

//! @copydoc omni::graph::exec::unstable::IPartitionPass_abi
//!
//! @ingroup  groupOmniGraphExecPasses groupOmniGraphExecInterfaces
class omni::graph::exec::unstable::IPartitionPass
    : public omni::core::Generated<omni::graph::exec::unstable::IPartitionPass_abi>
{
};

// additional headers needed for API implementation
#include <omni/graph/exec/unstable/IGraphBuilder.h>
#include <omni/graph/exec/unstable/INode.h>
#include <omni/graph/exec/unstable/ITopology.h>

// generated API implementation
#define OMNI_BIND_INCLUDE_INTERFACE_IMPL
#include <omni/graph/exec/unstable/IPartitionPass.gen.h>
