//
// Created by ssunny on 7/1/16.
//

#include <Register.h>
#include "SymEvalSemantics.h"

using namespace rose::BinaryAnalysis::InstructionSemantics2;

///////////////////////////////////////////////////////
//                                           StateARM64
///////////////////////////////////////////////////////

BaseSemantics::SValuePtr SymEvalSemantics::StateARM64::readRegister(const RegisterDescriptor &desc,
                                                                    const BaseSemantics::SValuePtr &/*dflt*/, BaseSemantics::RiscOperators */*ops*/) {
    ASSERT_require(desc.is_valid());
    SymEvalSemantics::RegisterStateARM64Ptr registers = SymEvalSemantics::RegisterStateARM64::promote(registerState());
    return registers->readRegister(desc, addr);
}

void SymEvalSemantics::StateARM64::writeRegister(const RegisterDescriptor &reg, const BaseSemantics::SValuePtr &value,
                                                 BaseSemantics::RiscOperators */*ops*/) {
    ASSERT_require(reg.is_valid());
    ASSERT_not_null(value);
    SymEvalSemantics::RegisterStateARM64Ptr registers = SymEvalSemantics::RegisterStateARM64::promote(registerState());
    registers->writeRegister(reg, value, res, aaMap);
}

BaseSemantics::SValuePtr SymEvalSemantics::StateARM64::readMemory(const BaseSemantics::SValuePtr &address,
                                                                  const BaseSemantics::SValuePtr &/*dflt*/,
                                                                  BaseSemantics::RiscOperators */*addrOps*/,
                                                                  BaseSemantics::RiscOperators */*valOps*/,
                                                                  size_t readSize) {
    ASSERT_not_null(address);
    SymEvalSemantics::MemoryStateARM64Ptr memory = SymEvalSemantics::MemoryStateARM64::promote(memoryState());
    return memory->readMemory(address, readSize);
}

void SymEvalSemantics::StateARM64::writeMemory(const BaseSemantics::SValuePtr &addr,
                                               const BaseSemantics::SValuePtr &value,
                                               BaseSemantics::RiscOperators */*addrOps*/,
                                               BaseSemantics::RiscOperators */*valOps*/,
                                               size_t writeSize) {
    ASSERT_not_null(addr);
    ASSERT_not_null(value);
    SymEvalSemantics::MemoryStateARM64Ptr memory = SymEvalSemantics::MemoryStateARM64::promote(memoryState());
    memory->writeMemory(addr, value, res, aaMap, writeSize);
}

///////////////////////////////////////////////////////
//                                   RegisterStateARM64
///////////////////////////////////////////////////////

BaseSemantics::SValuePtr SymEvalSemantics::RegisterStateARM64::readRegister(const RegisterDescriptor &reg,
                                                                            Dyninst::Address addr) {
    return SymEvalSemantics::SValue::instance(wrap(convert(reg), addr));
}

BaseSemantics::SValuePtr SymEvalSemantics::RegisterStateARM64::readRegister(const RegisterDescriptor &reg,
                                                                            const BaseSemantics::SValuePtr &/*dflt*/,
                                                                            BaseSemantics::RiscOperators */*ops*/) {
    ASSERT_always_forbid("overridden RegisterState::readRegister() should never be called for ARM64, always use the non-virtual readRegister that takes Dyninst::Address as an argument.");
}

void SymEvalSemantics::RegisterStateARM64::writeRegister(const RegisterDescriptor &reg,
                                                         const BaseSemantics::SValuePtr &value,
                                                         Dyninst::DataflowAPI::Result_t &res,
                                                         std::map<Dyninst::Absloc, Dyninst::Assignment::Ptr> &aaMap) {
    std::map<Dyninst::Absloc, Dyninst::Assignment::Ptr>::iterator i = aaMap.find(convert(reg));
    if (i != aaMap.end()) {
        SymEvalSemantics::SValuePtr value_ = SymEvalSemantics::SValue::promote(value);
        res[i->second] = value_->get_expression();
    }
}

void SymEvalSemantics::RegisterStateARM64::writeRegister(const RegisterDescriptor &/*reg*/,
                                                         const BaseSemantics::SValuePtr &/*value*/,
                                                         BaseSemantics::RiscOperators */*ops*/) {
    ASSERT_always_forbid("overridden RegisterState::writeRegister() should never be called for ARM64, always use the non-virtual writeRegister that also takes additional parameters.");
}

Dyninst::Absloc SymEvalSemantics::RegisterStateARM64::convert(const RegisterDescriptor &reg) {
    Dyninst::MachRegister mreg;

    unsigned int major = reg.get_major();
    unsigned int size = reg.get_nbits();

    switch (major) {
        case armv8_regclass_gpr: {
            unsigned int minor = reg.get_minor();

            if (minor == armv8_gpr_zr) {
                mreg = Dyninst::MachRegister((size == 32) ? Dyninst::aarch64::wzr : Dyninst::aarch64::zr);
            } else {
                Dyninst::MachRegister base = (size == 32) ? Dyninst::aarch64::w0 : Dyninst::aarch64::x0;
                mreg = Dyninst::MachRegister(base.val() + (minor - armv8_gpr_r0));
            }
        }
            break;
        case armv8_regclass_simd_fpr: {
            Dyninst::MachRegister base;
            unsigned int minor = reg.get_minor();

            switch(size) {
                case 8: base = Dyninst::aarch64::b0;
                    break;
                case 16: base = Dyninst::aarch64::h0;
                    break;
                case 32: base = Dyninst::aarch64::s0;
                    break;
                case 64:
                    if (reg.get_offset() == 64) {
                        base = Dyninst::aarch64::hq0;
                    } else {
                        base = Dyninst::aarch64::d0;
                    }
                    break;
                case 128: base = Dyninst::aarch64::q0;
                    break;
                default:assert(!"invalid size of RegisterDescriptor!");
                    break;
            }
            mreg = Dyninst::MachRegister(base.val() + (minor - armv8_simdfpr_v0));
        }
            break;
        case armv8_regclass_pc:
            mreg = Dyninst::MachRegister(Dyninst::aarch64::pc);
            break;
        case armv8_regclass_sp:
            mreg = Dyninst::MachRegister((size == 32) ? Dyninst::aarch64::wsp : Dyninst::aarch64::sp);
            break;
        case armv8_regclass_pstate: {
            unsigned int offset = reg.get_offset();
            switch (offset) {
                case armv8_pstatefield_n:
                    mreg = Dyninst::MachRegister(Dyninst::aarch64::n);
                    break;
                case armv8_pstatefield_z:
                    mreg = Dyninst::MachRegister(Dyninst::aarch64::z);
                    break;
                case armv8_pstatefield_c:
                    mreg = Dyninst::MachRegister(Dyninst::aarch64::c);
                    break;
                case armv8_pstatefield_v:
                    mreg = Dyninst::MachRegister(Dyninst::aarch64::v);
                    break;
                default:
                    ASSERT_always_forbid("No part of the PSTATE register other than NZCV should be used.");
                    break;
            }
        }
            break;
        default:
            ASSERT_always_forbid("Unexpected register major type.");
    }

    return Dyninst::Absloc(mreg);
}

///////////////////////////////////////////////////////
//                                     MemoryStateARM64
///////////////////////////////////////////////////////

//TODO: what is Len in this case?

BaseSemantics::SValuePtr SymEvalSemantics::MemoryStateARM64::readMemory(const BaseSemantics::SValuePtr &address,
                                                                        size_t readSize) {
    SymEvalSemantics::SValuePtr addr = SymEvalSemantics::SValue::promote(address);
    return SymEvalSemantics::SValue::instance(Dyninst::DataflowAPI::RoseAST::create(Dyninst::DataflowAPI::ROSEOperation(Dyninst::DataflowAPI::ROSEOperation::derefOp, readSize),
                                                                                    addr->get_expression(),
                                                                                    Dyninst::DataflowAPI::ConstantAST::create(Dyninst::DataflowAPI::Constant(1, 1))));
}

BaseSemantics::SValuePtr SymEvalSemantics::MemoryStateARM64::readMemory(const BaseSemantics::SValuePtr &address,
                                                                        const BaseSemantics::SValuePtr &/*dflt*/,
                                                                        BaseSemantics::RiscOperators */*addrOps*/,
                                                                        BaseSemantics::RiscOperators */*valOps*/) {
    ASSERT_always_forbid("overridden MemoryState::readMemory() should never be called for ARM64, always use the non-virtual readMemory that also takes additional parameters.");
}

void SymEvalSemantics::MemoryStateARM64::writeMemory(const BaseSemantics::SValuePtr &address,
                                                     const BaseSemantics::SValuePtr &value,
                                                     Dyninst::DataflowAPI::Result_t &res,
                                                     std::map<Dyninst::Absloc, Dyninst::Assignment::Ptr> &aaMap,
                                                     size_t writeSize) {
    std::map<Dyninst::Absloc, Dyninst::Assignment::Ptr>::iterator i = aaMap.find(Dyninst::Absloc(0));
    SymEvalSemantics::SValuePtr addr = SymEvalSemantics::SValue::promote(address);

    if (i != aaMap.end()) {
        i->second->out().setGenerator(addr->get_expression());
        //i->second->out().setSize(Len);

        SymEvalSemantics::SValuePtr data = SymEvalSemantics::SValue::promote(value);
        res[i->second] = data->get_expression();
    }
}

void SymEvalSemantics::MemoryStateARM64::writeMemory(const BaseSemantics::SValuePtr &/*addr*/,
                                                     const BaseSemantics::SValuePtr &/*value*/,
                                                     BaseSemantics::RiscOperators */*addrOps*/,
                                                     BaseSemantics::RiscOperators */*valOps*/) {
    ASSERT_always_forbid("overridden MemoryState::writeMemory() should never be called for ARM64, always use the non-virtual writeMemory that also takes additional parameters.");
}

///////////////////////////////////////////////////////
//                                        RiscOperators
///////////////////////////////////////////////////////

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::and_(const BaseSemantics::SValuePtr &a_,
                                                               const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::andOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::or_(const BaseSemantics::SValuePtr &a_,
                                                              const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::orOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::xor_(const BaseSemantics::SValuePtr &a_,
                                                               const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::xorOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::invert(const BaseSemantics::SValuePtr &a_) {
    return createUnaryAST(Dyninst::DataflowAPI::ROSEOperation::invertOp, a_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::extract(const BaseSemantics::SValuePtr &a_, uint64_t begin,
                                                                  uint64_t end) {
    BaseSemantics::SValuePtr begin_ = SymEvalSemantics::SValue::instance(64, begin);
    BaseSemantics::SValuePtr end_ = SymEvalSemantics::SValue::instance(64, end);

    return createTernaryAST(Dyninst::DataflowAPI::ROSEOperation::extractOp, a_, begin_, end_, end - begin);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::ite(const BaseSemantics::SValuePtr &sel_,
                                                              const BaseSemantics::SValuePtr &a_,
                                                              const BaseSemantics::SValuePtr &b_) {
    return createTernaryAST(Dyninst::DataflowAPI::ROSEOperation::ifOp, sel_, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::concat(const BaseSemantics::SValuePtr &a_,
                                                                 const BaseSemantics::SValuePtr &b_) {
    //TODO: should be able to specify number of bits to concat for each expression
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::concatOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::leastSignificantSetBit(const BaseSemantics::SValuePtr &a_) {
    return createUnaryAST(Dyninst::DataflowAPI::ROSEOperation::LSBSetOp, a_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::mostSignificantSetBit(const BaseSemantics::SValuePtr &a_) {
    return createUnaryAST(Dyninst::DataflowAPI::ROSEOperation::MSBSetOp, a_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::rotateLeft(const BaseSemantics::SValuePtr &a_,
                                                                     const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::rotateLOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::rotateRight(const BaseSemantics::SValuePtr &a_,
                                                                      const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::rotateROp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::shiftLeft(const BaseSemantics::SValuePtr &a_,
                                                                    const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::shiftLOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::shiftRight(const BaseSemantics::SValuePtr &a_,
                                                                     const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::shiftROp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::shiftRightArithmetic(const BaseSemantics::SValuePtr &a_,
                                                                               const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::shiftRArithOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::equalToZero(const BaseSemantics::SValuePtr &a_) {
    return createUnaryAST(Dyninst::DataflowAPI::ROSEOperation::equalToZeroOp, a_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::signExtend(const BaseSemantics::SValuePtr &a_,
                                                                     uint64_t newwidth) {
    BaseSemantics::SValuePtr width_ = SymEvalSemantics::SValue::instance(64, newwidth);

    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::signExtendOp, a_, width_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::add(const BaseSemantics::SValuePtr &a_,
                                                              const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::addOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::addWithCarries(const BaseSemantics::SValuePtr &a_,
                                                                         const BaseSemantics::SValuePtr &b_,
                                                                         const BaseSemantics::SValuePtr &c_,
                                                                         BaseSemantics::SValuePtr &carry_out) {
    BaseSemantics::SValuePtr aa_ = unsignedExtend(a_, a_->get_width() + 1);
    BaseSemantics::SValuePtr bb_ = unsignedExtend(b_, b_->get_width() + 1);
    BaseSemantics::SValuePtr sum_ = createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::addOp, aa_,
                                                    createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::addOp, bb_, c_));

    BaseSemantics::SValuePtr cc_ = createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::xorOp, aa_,
                                                   createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::xorOp, bb_, sum_));
    carry_out = extract(cc_, 1, a_->get_width() + 1);

    return extract(sum_, 0, a_->get_width());
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::negate(const BaseSemantics::SValuePtr &a_) {
    return createUnaryAST(Dyninst::DataflowAPI::ROSEOperation::negateOp, a_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::signedDivide(const BaseSemantics::SValuePtr &a_,
                                                                       const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::sDivOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::signedModulo(const BaseSemantics::SValuePtr &a_,
                                                                       const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::sModOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::signedMultiply(const BaseSemantics::SValuePtr &a_,
                                                                         const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::sMultOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::unsignedDivide(const BaseSemantics::SValuePtr &a_,
                                                                       const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::uDivOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::unsignedModulo(const BaseSemantics::SValuePtr &a_,
                                                                       const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::uModOp, a_, b_);
}

BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::unsignedMultiply(const BaseSemantics::SValuePtr &a_,
                                                                         const BaseSemantics::SValuePtr &b_) {
    return createBinaryAST(Dyninst::DataflowAPI::ROSEOperation::uMultOp, a_, b_);
}

//TODO: do we deal with byte ordering?
BaseSemantics::SValuePtr SymEvalSemantics::RiscOperatorsARM64::readMemory(const RegisterDescriptor &/*segreg*/,
                                                                     const BaseSemantics::SValuePtr &addr,
                                                                     const BaseSemantics::SValuePtr &dflt,
                                                                     const BaseSemantics::SValuePtr &/*cond*/) {
    return currentState()->readMemory(addr, dflt, this, this);
}

void SymEvalSemantics::RiscOperatorsARM64::writeMemory(const RegisterDescriptor &/*segreg*/,
                                                  const BaseSemantics::SValuePtr &addr,
                                                  const BaseSemantics::SValuePtr &data,
                                                  const BaseSemantics::SValuePtr &/*cond*/) {
    currentState()->writeMemory(addr, data, this, this);
}
