#include "verify.h"
#include "builtins/builtins.h"
using namespace goose::diagnostics;

// This is an implementation of the "Software Verification Using k-Induction" algorithm.
// http://www.cprover.org/kinduction/

namespace goose::verify
{
    bool Func::checkLoop( const cir::BasicBlock& header, queue< const BasicBlock* >& parentWorkQueue )
    {
        ProfileZoneScoped;

        if( ms_TraceMode )
            cout << "  == Checking loop " << header.index() << endl;

        auto origAssHandler = m_builder.assertionHandler();
        AssertionHandlerGuard ahg( m_builder );

        m_viz.beginLoopVerification();

        m_remapper.beginLoopUnrolling( header.index() );

        bool completed = false;
        uint32_t k = 0;

        // Max K is hardcoded for now. Maybe we'll allow it to be changed per function.
        // Or maybe we won't limit K but rather the number of expressions added to the solver
        // (so that we have a limit that works consistently regardless of loop nesting depth)
        uint32_t maxK = 16;

        while( !completed && k < maxK )
        {
            ++k;

            if( ms_TraceMode )
                cout << "  == Unrolling base case " << k << " for loop " << header.index() << endl;

            {
                DiagnosticsContext dc( header.locationId(), "In this loop." );

                // Unroll the base case just once for each increasing K: we save the state of the z3 solver
                // before generating the induction case, so we can incrementally add unrolling of the base
                // case one at a time to save some work.
                m_viz.nextLoopIteration();
                m_remapper.nextLoopIteration();
                if( !buildZ3Expressions( header, &parentWorkQueue ) )
                {
                    // If a counter example was found while emitting the base case,
                    // we know the verification failed and we can stop here.
                    completed = true;
                    break;
                }
            }

            // Push the solver state, builder state and remapper before unrolling the induction case,
            // so that we can rollback all of this for the next attempt, if any.
            m_solver.push();

            unordered_set< uint32_t > convertedBBIndices = m_convertedBBIndices;
            auto savedBuilder = m_builder;
            auto savedRemapper = m_remapper;

            // Set a different assertion handler for the induction case: we want to just emit them as assumptions.
            m_builder.setAssertionHandler( [&]( const z3::expr& expr, const z3::expr& exprToCheck, LocationId locationId )
            {
                m_builder.assume( expr );
                return true;
            } );

            VerifyViz::Guard vvg( m_viz );
            m_viz.beginLoopInduction();

            // Havoc all the addresses modified during the current loop and emit
            // the first unrolling for the induction case.
            m_remapper.nextLoopIteration();
            uint32_t firstInductionIterationHeaderIndex = m_remapper.remapBBId( header );
            m_builder.setCurrentBB( firstInductionIterationHeaderIndex );

            if( ms_TraceMode )
                cout << "  == Havocing modified addresses for loop " << header.index() << endl;

            m_cfg->forEachAddressModifiedByLoop( header.index(), [&]( auto&& type, auto&& addr )
            {
                HavocAddress( m_builder, firstInductionIterationHeaderIndex, type, addr );
            } );

            if( ms_TraceMode )
                cout << "  == Unrolling first induction case for loop " << header.index() << endl;

            buildZ3Expressions( header, &parentWorkQueue );

            // Emit the rest of the induction case.
            uint32_t i = 1;
            while( i < k )
            {
                if( ms_TraceMode )
                    cout << "  == Unrolling induction case " << i << " for loop " << header.index() << endl;

                m_viz.nextLoopIteration();
                m_remapper.nextLoopIteration();
                buildZ3Expressions( header, &parentWorkQueue );
                ++i;
            }

            // For the final loop instance, if we aren't at the last value of k before giving up,
            // we want to resolve the assertions without emitting diagnostics:
            // if they fail it means we need to increase k and try again.
            // Otherwise, we want to use the regular assertion handler that will emit error
            // messages for all failed assertions.
            if( k == maxK )
                m_builder.setAssertionHandler( origAssHandler );
            else
            {
                m_builder.setAssertionHandler( [&]( const z3::expr& expr, const z3::expr& exprToCheck, LocationId locationId )
                {
                    if( ms_TraceMode )
                        return true;

                    z3::expr_vector exprVec( GetZ3Context() );
                    exprVec.push_back( exprToCheck );
                    bool result = m_solver.check( exprVec ) == z3::check_result::unsat;

                    if( ( !result && ms_DumpSolverOnFailure ) || ( result && ms_DumpSolverOnSuccess ) )
                        cout << "Solver dump:\n" << m_solver << "check_unsat_ind: " << exprToCheck << endl << endl;

                    return result;
                } );
            }

            m_viz.nextLoopIteration();
            m_remapper.nextLoopIteration();
            if( buildZ3Expressions( header, &parentWorkQueue ) && !m_builder.hasCheckFailed() )
            {
                completed = true;
                break;
            }

            m_convertedBBIndices = move( convertedBBIndices );
            m_solver.pop();

            m_remapper = move( savedRemapper );
            m_builder = move( savedBuilder );
        }

        m_remapper.endLoopUnrolling();

        m_viz.endLoopVerification();

        if( ms_TraceMode )
            cout << "  == Finished checking loop " << header.index() << endl;
        return completed && !m_builder.hasCheckFailed();
    }
}
