#include "builtins/builtins.h"
#include "precedence.h"
#include "helpers.h"
#include "tuple.h"

using namespace goose;
using namespace goose::eir;
using namespace goose::cir;
using namespace goose::parse;

namespace goose::builtins
{
    void SetupLogicOps( Env& e )
    {
        auto orOp = GetOrCreateOverloadSet( e, "operator_or"_sid );
        auto andOp = GetOrCreateOverloadSet( e, "operator_and"_sid );

        BuildParseRule( e, "!"_sid,
            PrefixOp( "operator_logical_not"_sid, precedence::UnaryOps,
                BuildGenericTupleOperator(),

                ForType< bool >( []( auto&& c, auto&& operand ) -> Value
                {
                    return BuildComputedValue( GetValueType< bool >(),
                        operand, cir::Not( c.locationId() ) );
                } )
            )
        );

        BuildParseRule( e, "~"_sid,
            PrefixOp( "operator_bitwise_not"_sid, precedence::UnaryOps,
                BuildGenericTupleOperator(),

                ForType< CustomPattern< IntegerType, IntegerType::Pattern > >( []( auto&& c, auto&& operand ) -> Value
                {
                    auto opTypeVal = *EIRToValue( operand.type() );
                    auto opType = *FromValue< IntegerType >( opTypeVal );
                    return BuildComputedValue( operand.type(),
                        operand, Value( operand.type(), APSInt::getMaxValue( opType.m_numBits, true ) ),
                        Xor( c.locationId() )
                    );
                } )
            )
        );

        BuildParseRule( e, "^"_sid,
            LeftAssInfixOp( "operator_xor"_sid, precedence::OrOp,
                BuildGenericTupleOperator(),

                // Logical xor
                ForType< bool, Xor >(),

                // ct_int xor
                ForType< BigInt >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    if( !lhs.isConstant() || !rhs.isConstant() )
                    {
                        DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( c.locationId(), "bitwise operations between ct_int values are only allowed on constants." );
                        return PoisonValue();
                    }

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, Xor( c.locationId() ) );
                } ),

                // runtime integer xor, defined to work for any two integers of same
                // bit size and signedness.
                ForType< CustomPattern< IntegerType, IntegerType::Pattern > >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, Xor( c.locationId() ) );
                } )
            )
        );

        BuildParseRule( e, "|"_sid,
            LeftAssInfixOp( "operator_or"_sid, precedence::OrOp,
                BuildGenericTupleOperator(),

                // ct_int or
                ForType< BigInt >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    if( !lhs.isConstant() || !rhs.isConstant() )
                    {
                        auto loc = Location::CreateSpanningLocation( lhs.locationId(), rhs.locationId() );
                        DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( loc, "bitwise operations between ct_int values are only allowed on constants." );
                        return PoisonValue();
                    }

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, Or( c.locationId() ) );
                } ),

                // runtime integer or, defined to work for any two integers of same
                // bit size and signedness.
                ForType< CustomPattern< IntegerType, IntegerType::Pattern > >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, Or( c.locationId() ) );
                } ),

                // bool or
                ForType< bool >( [orOp]< typename L, typename R >( auto&& c, L&& lhs, R&& rhs ) -> Value
                {
                    // Handle the case where lhs is constant, so that
                    // the result gets properly eagerly evaluated (in case
                    // we are using the expression as a compile time constant)
                    if( lhs.isConstant() )
                    {
                        if( *FromValue< bool >( lhs ) )
                            return forward< L >( lhs );

                        return forward< R >( rhs );
                    }

                    // This operator have different behaviors depending on the context:
                    // in normal code, we want to generate shortcut evaluation. But in propositions
                    // and ghost code, we want to simply generate a Or instruction with both operands
                    // always evaluated. So we delegate the work to another overload of operator_or that
                    // takes the builder as its first param and is overloaded according to it.
                    return InvokeOverloadSet( c,
                        orOp, MakeClosedTuple( c.builder(), forward< L >( lhs ), forward< R >( rhs ) ) );
                } )
            )
        );

        RegisterBuiltinFunc< Intrinsic< bool ( Value, bool, bool ) > >( e, orOp,
            []( auto&& c, auto&& b, auto&& lhs, auto&& rhs )
            {
                return BuildComputedValue( GetValueType< bool >(), lhs, rhs, Or( c.locationId() ) );
            } );

        RegisterBuiltinFunc< Intrinsic< bool ( TypeWrapper< ptr< CodeBuilder > >, bool, bool ) > >( e, orOp,
            []( auto&& c, auto&& b, auto&& lhs, auto&& rhs )
            {
                auto cb = *FromValue< TypeWrapper< ptr< CodeBuilder > > >( b );
                const auto& cfg = cb->cfg();

                // Build the control flow for shortcut evaluation.
                const auto& predBB = cfg->currentBB();

                auto pRhsBB = cfg->createBB();
                auto pSuccBB = cfg->createBB();

                // If the lhs is true, skip to the end directly.
                // Otherwise, jump to the BB that computes rhs.
                predBB->append( lhs );
                predBB->setTerminator( CondBranch( pSuccBB, pRhsBB ) );

                auto rhsIndex = util::GenerateNewUID();
                pRhsBB->append( rhs, CreateTemporary( rhsIndex, false, c.locationId() ) );
                pRhsBB->setTerminator( Branch( pSuccBB ) );

                auto resultIndex = util::GenerateNewUID();

                // Build the Phi instruction that will collect the final result.
                auto phi = Phi( GetValueType< bool >(), 2,
                    resultIndex, c.locationId() );

                // If coming directly from the lhs BB, we know the result is true.
                phi.setIncoming( predBB, ToValue( true ) );

                // Otherwise, the result is whatever was computed by the rhs block.
                phi.setIncoming( pRhsBB, BuildComputedValue( GetValueType< bool >(),
                    GetTemporary( GetValueType< bool >(), rhsIndex, c.locationId() ) ) );

                pSuccBB->append( move( phi ) );
                cfg->setCurrentBB( pSuccBB );

                // Build the result val which pulls the temporary created above.
                return BuildComputedValue( GetValueType< bool >(),
                    GetTemporary( GetValueType< bool >(), resultIndex, c.locationId() ) );
            } );

        BuildParseRule( e, "&"_sid,
            LeftAssInfixOp( "operator_and"_sid, precedence::AndOp,
                BuildGenericTupleOperator(),

                // ct_int and
                ForType< BigInt >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    if( !lhs.isConstant() || !rhs.isConstant() )
                    {
                        DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( c.locationId(), "bitwise operations between ct_int values are only allowed on constants." );
                        return PoisonValue();
                    }

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, And( c.locationId() ) );
                } ),

                // runtime integer and, defined to work for any two integers of same
                // bit size and signedness.
                ForType< CustomPattern< IntegerType, IntegerType::Pattern > >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, And( c.locationId() ) );
                } ),

                // bool and
                ForType< bool >( [andOp]< typename L, typename R >( auto&& c, L&& lhs, R&& rhs ) -> Value
                {
                    // Handle the case where lhs is constant, so that
                    // the result gets properly eagerly evaluated (in case
                    // we are using the expression as a compile time constant)
                    if( lhs.isConstant() )
                    {
                        if( *FromValue< bool >( lhs ) )
                            return forward< R >( rhs );

                        return forward< L >( lhs );
                    }

                    // This operator have different behaviors depending on the context:
                    // in normal code, we want to generate shortcut evaluation. But in propositions
                    // and ghost code, we want to simply generate a And instruction with both operands
                    // always evaluated. So we delegate the work to another overload of operator_and that
                    // takes the builder as its first param and is overloaded according to it.
                    return InvokeOverloadSet( c,
                        andOp, MakeClosedTuple( c.builder(), forward< L >( lhs ), forward< R >( rhs ) ) );
                } )
            )
        );

        RegisterBuiltinFunc< Intrinsic< bool ( Value, bool, bool ) > >( e, andOp,
            []( auto&& c, auto&& b, auto&& lhs, auto&& rhs )
            {
                return BuildComputedValue( GetValueType< bool >(), lhs, rhs, And( c.locationId() ) );
            } );

        RegisterBuiltinFunc< Intrinsic< bool ( TypeWrapper< ptr< CodeBuilder > >, bool, bool ) > >( e, andOp,
            []( auto&& c, auto&& b, auto&& lhs, auto&& rhs )
            {
                auto cb = *FromValue< TypeWrapper< ptr< CodeBuilder > > >( b );
                const auto& cfg = cb->cfg();

                // Build the control flow for shortcut evaluation.
                const auto& predBB = cfg->currentBB();

                auto pRhsBB = cfg->createBB();
                auto pSuccBB = cfg->createBB();

                // If the lhs is false, skip to the end directly.
                // Otherwise, jump to the BB that computes rhs.
                predBB->append( lhs );
                predBB->setTerminator( CondBranch( pRhsBB, pSuccBB ) );

                auto rhsIndex = util::GenerateNewUID();
                pRhsBB->append( rhs, CreateTemporary( rhsIndex, false, c.locationId() ) );
                pRhsBB->setTerminator( Branch( pSuccBB ) );

                auto resultIndex = util::GenerateNewUID();

                // Build the Phi instruction that will collect the final result.
                auto phi = Phi( GetValueType< bool >(), 2,
                    resultIndex, c.locationId() );

                // If coming directly from the lhs BB, we know the result is false.
                phi.setIncoming( predBB, ToValue( false ) );

                // Otherwise, the result is whatever was computed by the rhs block.
                phi.setIncoming( pRhsBB, BuildComputedValue( GetValueType< bool >(),
                    GetTemporary( GetValueType< bool >(), rhsIndex, c.locationId() ) ) );

                pSuccBB->append( move( phi ) );
                cfg->setCurrentBB( pSuccBB );

                // Build the result val which pulls the temporary created above.
                return BuildComputedValue( GetValueType< bool >(),
                    GetTemporary( GetValueType< bool >(), resultIndex, c.locationId() ) );
            } );

        BuildParseRule( e, "<<"_sid,
            LeftAssInfixOp( "operator_shift_left"_sid, precedence::BitShiftOp,
                BuildGenericTupleOperator(),

                // ct_int left shift
                ForTypes< BigInt, uint32_t >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    if( !lhs.isConstant() || !rhs.isConstant() )
                    {
                        auto loc = Location::CreateSpanningLocation( lhs.locationId(), rhs.locationId() );
                        DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( loc, "bitwise operations between ct_int values are only allowed on constants." );
                        return PoisonValue();
                    }

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, Shl( c.locationId() ) );
                } ),

                // runtime integer left shift.
                ForTypes< CustomPattern< IntegerType, IntegerType::Pattern >,
                    CustomPattern< IntegerType, IntegerType::PatternUnsigned > >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    using namespace goose::builtins::exprhelpers;

                    auto cfg = GetCFG( c );
                    assert( cfg );

                    // Shifting for a number of bits equal or larger than the bitsize
                    // of lhs is an undefined behavior, so we require verification that
                    // it won't happen.
                    // Extract the integer type of lhs to retrieve its bit size.
                    auto lhsType = *FromValue< IntegerType >( *EIRToValue( lhs.type() ) );
                    auto bitSizeValue = Value( rhs.type(), APSInt::get( lhsType.m_numBits ) );

                    auto cond = ULT( rhs, bitSizeValue );

                    DiagnosticsManager::GetInstance().defineCustomDiagnostic(
                        cond.locationId(), "assert"_sid, "the shift amount may be equal or greater than the bitsize." );

                    cfg->currentBB()->append(
                        move( cond ), cir::Assert( rhs.locationId() )
                    );

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, Shl( c.locationId() ) );
                } )
            )
        );

        BuildParseRule( e, ">>"_sid,
            LeftAssInfixOp( "operator_shift_right"_sid, precedence::BitShiftOp,
                BuildGenericTupleOperator(),

                // ct_int right shift
                ForTypes< BigInt, uint32_t >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    if( !lhs.isConstant() || !rhs.isConstant() )
                    {
                        DiagnosticsManager::GetInstance().emitSyntaxErrorMessage( c.locationId(), "bitwise operations between ct_int values are only allowed on constants." );
                        return PoisonValue();
                    }

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, AShr( c.locationId() ) );
                } ),

                // runtime signed integer right shift, defined to work for any two integers of same
                // bit size.
                ForTypes< CustomPattern< IntegerType, IntegerType::PatternSigned >,
                    CustomPattern< IntegerType, IntegerType::PatternUnsigned > >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    using namespace goose::builtins::exprhelpers;

                    auto cfg = GetCFG( c );
                    assert( cfg );

                    // Shifting for a number of bits equal or larger than the bitsize
                    // of lhs is an undefined behavior, so we require verification that
                    // it won't happen.
                    // Extract the integer type of lhs to retreieve its bit size.
                    auto lhsType = *FromValue< IntegerType >( *EIRToValue( lhs.type() ) );
                    auto bitSizeValue = Value( rhs.type(), APSInt::get( lhsType.m_numBits) );

                    auto cond = ULT( rhs, bitSizeValue );

                    DiagnosticsManager::GetInstance().defineCustomDiagnostic(
                        cond.locationId(), "assert"_sid, "the shift amount may be equal or greater than the bitsize." );

                    cfg->currentBB()->append(
                        move( cond ), cir::Assert( rhs.locationId() )
                    );

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, AShr( c.locationId() ) );
                } ),

                // runtime unsigned integer right shift, defined to work for any two integers of same
                // bit size.
                ForType< CustomPattern< IntegerType, IntegerType::PatternUnsigned > >(
                []( auto&& c, auto&& lhs, auto&& rhs ) -> Value
                {
                    using namespace goose::builtins::exprhelpers;

                    auto cfg = GetCFG( c );
                    assert( cfg );

                    // Shifting for a number of bits equal or larger than the bitsize
                    // of lhs is an undefined behavior, so we require verification that
                    // it won't happen.
                    // Extract the integer type of lhs to retreieve its bit size.
                    auto lhsType = *FromValue< IntegerType >( *EIRToValue( lhs.type() ) );
                    auto bitSizeValue = Value( rhs.type(), APSInt::get( lhsType.m_numBits ) );

                    auto cond = ULT( rhs, bitSizeValue );

                    DiagnosticsManager::GetInstance().defineCustomDiagnostic(
                        cond.locationId(), "assert"_sid, "the shift amount may be equal or greater than the bitsize." );

                    cfg->currentBB()->append(
                        move( cond ), cir::Assert( rhs.locationId() )
                    );

                    return BuildComputedValue( lhs.type(),
                        lhs, rhs, LShr( c.locationId() ) );
                } )
            )
        );
    }
}
