#include "pch.h"

#include "Assembly.h"
#include "CppUnitTest.h"

using namespace Microsoft::VisualStudio::CppUnitTestFramework;

int globalFunc(int a, int b)
{
    return a * b;
}

class A
{
public:
    virtual int getValue()
    {
        return 0;
    }
};

class B : public A
{
public:
    virtual int getValue() override
    {
        return 10;
    }
};

namespace FrameworkTests
{
    TEST_CLASS (AssemblyTests)
    {
    public:
        TEST_METHOD (Add8Test)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addMoveValue(Framework::Assembly::RAX, (__int64)0);
            codeBlock.addMoveValue(Framework::Assembly::RAX,
                Framework::Assembly::RDX,
                Framework::Assembly::LOWER8);
            codeBlock.addAddition(Framework::Assembly::RAX,
                Framework::Assembly::RCX,
                Framework::Assembly::LOWER8);
            int (*add)(int a, int b) = (int (*)(int, int))codeBlock.compile();
            int result = add(22, 15);
            int result2 = add(0x0FFFFF43, 0x0FFFFF05);
            Assert::AreEqual(37, result);
            Assert::AreEqual(0x48, result2);
        }

        TEST_METHOD (Add16Test)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addMoveValue(Framework::Assembly::RAX, (__int64)0);
            codeBlock.addMoveValue(Framework::Assembly::RAX,
                Framework::Assembly::RDX,
                Framework::Assembly::LOWER16);
            codeBlock.addAddition(Framework::Assembly::RAX,
                Framework::Assembly::RCX,
                Framework::Assembly::LOWER16);
            int (*add)(int a, int b) = (int (*)(int, int))codeBlock.compile();
            int result = add(22, 15);
            int result2 = add(0x11111111, 0x22222222);
            Assert::AreEqual(37, result);
            Assert::AreEqual(0x3333, result2);
        }

        TEST_METHOD (Add32Test)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addMoveValue(Framework::Assembly::RAX,
                Framework::Assembly::RDX,
                Framework::Assembly::LOWER32);
            codeBlock.addAddition(Framework::Assembly::RAX,
                Framework::Assembly::RCX,
                Framework::Assembly::LOWER32);
            __int64 (*add)(__int64 a, __int64 b)
                = (__int64 (*)(__int64, __int64))codeBlock.compile();
            __int64 result = add(22, 15);
            __int64 result2 = add(0x1111111111111111, 0x1111111111111111);
            Assert::AreEqual((__int64)37, result);
            Assert::AreEqual((__int64)0x22222222, result2);
        }

        TEST_METHOD (Add64Test)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addMoveValue(
                Framework::Assembly::RAX, Framework::Assembly::RDX);
            codeBlock.addAddition(
                Framework::Assembly::RAX, Framework::Assembly::RCX);
            __int64 (*add)(__int64 a, __int64 b)
                = (__int64 (*)(__int64, __int64))codeBlock.compile();
            __int64 result = add(22, 15);
            __int64 result2 = add(0x1111111111111111, 0x1111111111111111);
            Assert::AreEqual((__int64)37, result);
            Assert::AreEqual((__int64)0x2222222222222222, result2);
        }

        TEST_METHOD (AddFloatTest)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addAddition(Framework::Assembly::MM0,
                Framework::Assembly::MM1,
                Framework::Assembly::SINGLE_FLOAT,
                Framework::Assembly::X);
            float (*add)(float a, float b)
                = (float (*)(float, float))codeBlock.compile();
            float result = add(4.24f, 8.54f);
            float result2 = add(0.3f, 7.6f);
            Assert::AreEqual((float)(4.24f + 8.54f), result);
            Assert::AreEqual((float)(0.3f + 7.6f), result2);
        }

        TEST_METHOD (AddDoubleTest)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addAddition(Framework::Assembly::MM0,
                Framework::Assembly::MM1,
                Framework::Assembly::SINGLE_DOUBLE,
                Framework::Assembly::X);
            double (*add)(double a, double b)
                = (double (*)(double, double))codeBlock.compile();
            double result = add(4.24, 8.54);
            double result2 = add(0.3, 7.6);
            Assert::AreEqual(4.24 + 8.54, result);
            Assert::AreEqual(0.3 + 7.6, result2);
        }

        TEST_METHOD (returnRefTest)
        {
            char cRef = 1;
            short sRef = 2;
            int iRef = 3;
            __int64 lRef = 4;
            float fRef = 5.0f;
            double dRef = 6.0;
            Framework::Assembly::AssemblyBlock ccodeBlock;
            ccodeBlock.addLoadValue(&cRef, Framework::Assembly::RAX);
            char (*getc)() = (char (*)())ccodeBlock.compile();
            char cresult = getc();
            Assert::AreEqual((char)1, cresult);
            cRef = 11;
            cresult = getc();
            Assert::AreEqual((char)11, cresult);
            Framework::Assembly::AssemblyBlock scodeBlock;
            scodeBlock.addLoadValue(&sRef, Framework::Assembly::RAX);
            short (*gets)() = (short (*)())scodeBlock.compile();
            short sresult = gets();
            Assert::AreEqual((short)2, sresult);
            sRef = 22;
            sresult = gets();
            Assert::AreEqual((short)22, sresult);
            Framework::Assembly::AssemblyBlock icodeBlock;
            icodeBlock.addLoadValue(&iRef, Framework::Assembly::RAX);
            int (*geti)() = (int (*)())icodeBlock.compile();
            int iresult = geti();
            Assert::AreEqual((int)3, iresult);
            iRef = 33;
            iresult = geti();
            Assert::AreEqual((int)33, iresult);
            Framework::Assembly::AssemblyBlock lcodeBlock;
            lcodeBlock.addLoadValue(&lRef, Framework::Assembly::RAX);
            __int64 (*getl)() = (__int64 (*)())lcodeBlock.compile();
            __int64 lresult = getl();
            Assert::AreEqual((__int64)4, lresult);
            lRef = 44;
            lresult = getl();
            Assert::AreEqual((__int64)44, lresult);
            Framework::Assembly::AssemblyBlock fcodeBlock;
            fcodeBlock.addLoadValue(&fRef, Framework::Assembly::MM0);
            float (*getf)() = (float (*)())fcodeBlock.compile();
            float fresult = getf();
            Assert::AreEqual((float)5.0f, fresult);
            fRef = 55.0f;
            fresult = getf();
            Assert::AreEqual((float)55.0f, fresult);
            Framework::Assembly::AssemblyBlock dcodeBlock;
            dcodeBlock.addLoadValue(&dRef, Framework::Assembly::MM0);
            double (*getd)() = (double (*)())dcodeBlock.compile();
            double dresult = getd();
            Assert::AreEqual((double)6.0, dresult);
            dRef = 66.0;
            dresult = getd();
            Assert::AreEqual((double)66.0, dresult);
        }
        int c;

        int testMethod(int a, int b)
        {
            return a + b + c;
        }

        TEST_METHOD (testCall)
        {
            Framework::Assembly::AssemblyBlock gfcodeBlock;
            gfcodeBlock.addMoveValue(Framework::Assembly::RCX, 20);
            gfcodeBlock.addMoveValue(Framework::Assembly::RDX, 50);
            gfcodeBlock.addCall(globalFunc,
                Framework::Assembly::INT_VALUE,
                {Framework::Assembly::RCX, Framework::Assembly::RDX});
            int (*f)() = gfcodeBlock.compileToFunction<int (*)()>();
            int result = f();
            Assert::AreEqual(globalFunc(20, 50), result);
            Framework::Assembly::AssemblyBlock lfcodeBlock;
            lfcodeBlock.addLoadAddress(this, Framework::Assembly::RCX);
            lfcodeBlock.addMoveValue(Framework::Assembly::RDX, 20);
            lfcodeBlock.addMoveValue(Framework::Assembly::R8, 50);
            lfcodeBlock.addMemberCall(&AssemblyTests::testMethod,
                Framework::Assembly::INT_VALUE,
                {Framework::Assembly::RCX,
                    Framework::Assembly::RDX,
                    Framework::Assembly::R8});
            c = 1;
            f = lfcodeBlock.compileToFunction<int (*)()>();
            result = f();
            A* a = new B();
            Assert::AreEqual(testMethod(20, 50), result);
            Framework::Assembly::AssemblyBlock vfcodeBlock;
            vfcodeBlock.addLoadAddress(a, Framework::Assembly::RCX);
            vfcodeBlock.addMemberCall(
                &A::getValue, Framework::Assembly::INT_VALUE);
            c = 1;
            f = vfcodeBlock.compileToFunction<int (*)()>();
            result = f();
            Assert::AreEqual(a->getValue(), result);
            delete a;
        }

        TEST_METHOD (mulTest)
        {
            float fa = 3.0f;
            float fb = 4.0f;
            double da = 5.0;
            double db = 6.0;
            Framework::Assembly::AssemblyBlock fcodeBlock;
            fcodeBlock.addLoadValue(&fa, Framework::Assembly::MM0);
            fcodeBlock.addLoadValue(&fb, Framework::Assembly::MM1);
            fcodeBlock.addMultiplication(Framework::Assembly::MM0,
                Framework::Assembly::MM1,
                Framework::Assembly::SINGLE_FLOAT,
                Framework::Assembly::X);
            float (*fmul)() = fcodeBlock.compileToFunction<float (*)()>();
            float fresult = fmul();
            Assert::AreEqual(fa * fb, fresult);
            Framework::Assembly::AssemblyBlock dcodeBlock;
            dcodeBlock.addLoadValue(&da, Framework::Assembly::MM0);
            dcodeBlock.addLoadValue(&db, Framework::Assembly::MM1);
            dcodeBlock.addMultiplication(Framework::Assembly::MM0,
                Framework::Assembly::MM1,
                Framework::Assembly::SINGLE_DOUBLE,
                Framework::Assembly::X);
            double (*dmul)() = dcodeBlock.compileToFunction<double (*)()>();
            double dresult = dmul();
            Assert::AreEqual(da * db, dresult);
        }

        TEST_METHOD (jumpTest)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addMoveValue(Framework::Assembly::RAX, 0);
            codeBlock.addJump(Framework::Assembly::JMP, "end");
            codeBlock.addMoveValue(Framework::Assembly::RAX, 1);
            codeBlock.defineJumpTarget("start");
            codeBlock.addMoveValue(Framework::Assembly::RAX, 2);
            codeBlock.defineJumpTarget("end");
            codeBlock.addTest(
                Framework::Assembly::RAX, Framework::Assembly::RAX);
            codeBlock.addJump(Framework::Assembly::JZ, "start");
            int (*test)() = codeBlock.compileToFunction<int (*)()>();
            int result = test();
            Assert::AreEqual(2, result);
        }

        TEST_METHOD (stackTest)
        {
            Framework::Assembly::AssemblyBlock codeBlock;
            codeBlock.addMoveValue(Framework::Assembly::RAX, (char)10);
            codeBlock.addMoveValue(Framework::Assembly::RBX, (short)20);
            codeBlock.addMoveValue(Framework::Assembly::RCX, (int)30);
            codeBlock.addMoveValue(Framework::Assembly::RDX, (__int64)40);
            codeBlock.addMoveValue(
                Framework::Assembly::MM0, 5.f, Framework::Assembly::R8);
            codeBlock.addMoveValue(
                Framework::Assembly::MM1, 6.0, Framework::Assembly::R8);
            codeBlock.addPush(
                Framework::Assembly::RAX, Framework::Assembly::LOWER8);
            codeBlock.addPush(
                Framework::Assembly::RBX, Framework::Assembly::LOWER16);
            codeBlock.addPush(
                Framework::Assembly::RCX, Framework::Assembly::LOWER32);
            codeBlock.addPush(
                Framework::Assembly::RDX, Framework::Assembly::FULL64);
            codeBlock.addPush(Framework::Assembly::MM0,
                Framework::Assembly::SINGLE_FLOAT,
                Framework::Assembly::X);
            codeBlock.addPush(Framework::Assembly::MM1,
                Framework::Assembly::SINGLE_DOUBLE,
                Framework::Assembly::X);
            codeBlock.addPop(Framework::Assembly::MM1,
                Framework::Assembly::SINGLE_DOUBLE,
                Framework::Assembly::X);
            codeBlock.addMoveValue(Framework::Assembly::MM2, 6.0);
            codeBlock.addCompare(Framework::Assembly::MM1,
                Framework::Assembly::MM2,
                Framework::Assembly::SINGLE_DOUBLE);
            codeBlock.addJump(Framework::Assembly::JNE, "error");
            codeBlock.addPop(Framework::Assembly::MM0,
                Framework::Assembly::SINGLE_FLOAT,
                Framework::Assembly::X);
            codeBlock.addMoveValue(Framework::Assembly::MM2, 5.f);
            codeBlock.addCompare(Framework::Assembly::MM0,
                Framework::Assembly::MM2,
                Framework::Assembly::SINGLE_FLOAT);
            codeBlock.addJump(Framework::Assembly::JNE, "error");
            codeBlock.addPop(
                Framework::Assembly::RDX, Framework::Assembly::FULL64);
            codeBlock.addCompare(Framework::Assembly::RDX, 40);
            codeBlock.addJump(Framework::Assembly::JNE, "error");
            codeBlock.addPop(
                Framework::Assembly::RCX, Framework::Assembly::LOWER32);
            codeBlock.addCompare(Framework::Assembly::RCX, 30);
            codeBlock.addJump(Framework::Assembly::JNE, "error");
            codeBlock.addPop(
                Framework::Assembly::RBX, Framework::Assembly::LOWER16);
            codeBlock.addCompare(Framework::Assembly::RBX, (short)20);
            codeBlock.addJump(Framework::Assembly::JNE, "error");
            codeBlock.addPop(
                Framework::Assembly::RAX, Framework::Assembly::LOWER8);
            codeBlock.addCompare(Framework::Assembly::RAX, (char)10);
            codeBlock.addJump(Framework::Assembly::JNE, "error");
            codeBlock.addMoveValue(Framework::Assembly::RAX, 1);
            codeBlock.addJump(Framework::Assembly::JMP, "end");
            codeBlock.defineJumpTarget("error");
            codeBlock.addMoveValue(Framework::Assembly::RAX, 0);
            codeBlock.defineJumpTarget("end");
            int (*test)() = codeBlock.compileToFunction<int (*)()>();
            int result = test();
            Assert::AreEqual(1, result);
        }
    };
} // namespace FrameworkTests