// -*- C++ -*- (c) 2019 Vladimír Štill <divin@vstill.eu>

DIVINE_RELAX_WARNINGS
#include <llvm/Pass.h>
#include <llvm/IR/Value.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Constants.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/IntrinsicInst.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/CallSite.h>
#include <llvm/Transforms/Utils/BasicBlockUtils.h>
#include <llvm/Transforms/Utils/Cloning.h>
#include <llvm/Analysis/CaptureTracking.h>
DIVINE_UNRELAX_WARNINGS
#include <lart/support/pass.h>
#include <lart/support/meta.h>
#include <lart/support/util.h>
#include <lart/lart.h>

#include <brick-llvm>
#include <brick-string>
#include <brick-assert>

namespace lart::nontermination {

struct Termsec
{
    static PassMeta meta()
    {
        return passMetaC( "termsec", "Instrument detection of terminating blocks.",
                []( PassVector &ps, std::string opt ) {
                    ps.emplace_back< Termsec >( opt );
                } );
    }

    explicit Termsec( std::string mode ) : _mode( read_nonterm( mode ) ) { }

    void run( llvm::Module &m )
    {
        auto hook_begin = m.getFunction( "__lart_termsec_begin" );
        auto hook_end = m.getFunction( "__lart_termsec_end" );
        auto impl_begin = m.getFunction( "__lart_termsec_begin_impl" );
        ASSERT( impl_begin );
        auto impl_end = m.getFunction( "__lart_termsec_end_impl" );
        ASSERT( impl_end );
        auto resched_cond_hook = m.getFunction( "__lart_termsec_reschedule_cond_acc_hook" );
        auto resched_hook = m.getFunction( "__lart_termsec_reschedule_acc_hook" );
        auto resched = m.getFunction( "__dios_reschedule" );
        auto suspend = m.getFunction( "__dios_suspend" );

        if ( _mode != NonterminationMode::Global )
        {
            using FP = std::pair< llvm::Function *, llvm::Function * >;
            for ( auto [ hook, impl ] : { FP{ hook_begin, impl_begin }, FP{ hook_end, impl_end } } )
            {
                std::vector< llvm::User * > users( hook->user_begin(), hook->user_end() );
                for ( auto usr : users )
                {
                    llvm::CallSite call( usr );
                    if ( call )
                    {
                        auto kind_c = llvm::dyn_cast< llvm::ConstantInt >( call.getArgument( 0 ) );
                        if ( !kind_c )
                            UNREACHABLE( "Unexpected non-constant termsec kind value in ", usr );
                        NonterminationMode kind = NonterminationMode( kind_c->getZExtValue() );
                        if ( _mode & kind )
                            call.setCalledFunction( impl );
                    }
                }
            }
        }

        for ( auto dios_end : { resched, suspend } )
        {
            if ( dios_end ) {
                llvm::IRBuilder<> irb( &*dios_end->begin(), dios_end->begin()->getFirstInsertionPt() );
                llvm::CallSite cs = irb.CreateCall( _mode == NonterminationMode::Global
                                                    ? resched_hook : resched_cond_hook, { } );
                llvm::InlineFunctionInfo ifi;
                llvm::InlineFunction( cs, ifi );
            }
        }
    }

    NonterminationMode _mode;
};

PassMeta termsecPass()
{
    return Termsec::meta();
}

} // namespace lart::nontermination