// Copyright 2015 Dolphin Emulator Project
// Licensed under GPLv2+
// Refer to the license.txt file included.

#include "Common/Arm64Emitter.h"
#include "Common/Common.h"

#include "Core/HW/MMIO.h"

#include "Core/PowerPC/JitArm64/Jit.h"
#include "Core/PowerPC/JitArm64/Jit_Util.h"
template <typename T>
class MMIOWriteCodeGenerator : public MMIO::WriteHandlingMethodVisitor<T>
{
public:
	MMIOWriteCodeGenerator(ARM64XEmitter* emit, BitSet32 gprs_in_use, BitSet32 fprs_in_use,
	                      ARM64Reg src_reg, u32 address)
		: m_emit(emit), m_gprs_in_use(gprs_in_use), m_fprs_in_use(fprs_in_use),
		  m_src_reg(src_reg), m_address(address)
	{
	}

	virtual void VisitNop()
	{
		// Do nothing
	}
	virtual void VisitDirect(T* addr, u32 mask)
	{
		WriteRegToAddr(8 * sizeof (T), addr, mask);
	}
	virtual void VisitComplex(const std::function<void(u32, T)>* lambda)
	{
		CallLambda(8 * sizeof (T), lambda);
	}

private:

	void StoreFromRegister(int sbits, ARM64Reg reg)
	{
		switch (sbits)
		{
		case 8:
			m_emit->STRB(INDEX_UNSIGNED, reg, X0, 0);
		break;
		case 16:
			m_emit->STRH(INDEX_UNSIGNED, reg, X0, 0);
		break;
		case 32:
			m_emit->STR(INDEX_UNSIGNED, reg, X0, 0);
		break;
		default:
			_assert_msg_(DYNA_REC, false, "Unknown size %d passed to MMIOWriteCodeGenerator!", sbits);
		break;
		}
	}

	void WriteRegToAddr(int sbits, const void* ptr, u32 mask)
	{
		m_emit->MOVI2R(X0, (u64)ptr);

		// If we do not need to mask, we can do the sign extend while loading
		// from memory. If masking is required, we have to first zero extend,
		// then mask, then sign extend if needed (1 instr vs. ~4).
		u32 all_ones = (1ULL << sbits) - 1;
		if ((all_ones & mask) == all_ones)
		{
			StoreFromRegister(sbits, m_src_reg);
		}
		else
		{
			m_emit->MOVI2R(W1, mask);
			m_emit->AND(W1, m_src_reg, W1, ArithOption(W1, ST_LSL, 0));
			StoreFromRegister(sbits, W1);
		}
	}

	void CallLambda(int sbits, const std::function<void(u32, T)>* lambda)
	{
		ARM64FloatEmitter float_emit(m_emit);

		m_emit->ABI_PushRegisters(m_gprs_in_use);
		float_emit.ABI_PushRegisters(m_fprs_in_use, X1);
			m_emit->MOVI2R(W1, m_address);
			m_emit->MOV(W2, m_src_reg);
			m_emit->BLR(m_emit->ABI_SetupLambda(lambda));
		float_emit.ABI_PopRegisters(m_fprs_in_use, X1);
		m_emit->ABI_PopRegisters(m_gprs_in_use);
	}

	ARM64XEmitter* m_emit;
	BitSet32 m_gprs_in_use;
	BitSet32 m_fprs_in_use;
	ARM64Reg m_src_reg;
	u32 m_address;
};
// Visitor that generates code to read a MMIO value.
template <typename T>
class MMIOReadCodeGenerator : public MMIO::ReadHandlingMethodVisitor<T>
{
public:
	MMIOReadCodeGenerator(ARM64XEmitter* emit, BitSet32 gprs_in_use, BitSet32 fprs_in_use,
	                      ARM64Reg dst_reg, u32 address, bool sign_extend)
		: m_emit(emit), m_gprs_in_use(gprs_in_use), m_fprs_in_use(fprs_in_use),
		  m_dst_reg(dst_reg), m_address(address), m_sign_extend(sign_extend)
	{
	}

	virtual void VisitConstant(T value)
	{
		LoadConstantToReg(8 * sizeof (T), value);
	}
	virtual void VisitDirect(const T* addr, u32 mask)
	{
		LoadAddrMaskToReg(8 * sizeof (T), addr, mask);
	}
	virtual void VisitComplex(const std::function<T(u32)>* lambda)
	{
		CallLambda(8 * sizeof (T), lambda);
	}

private:

	void LoadConstantToReg(int sbits, u32 value)
	{
		m_emit->MOVI2R(m_dst_reg, value);
		if (m_sign_extend)
			m_emit->SBFM(m_dst_reg, m_dst_reg, 0, sbits - 1);
	}

	void LoadToRegister(int sbits, bool dont_extend)
	{
		switch (sbits)
		{
		case 8:
			if (m_sign_extend && !dont_extend)
				m_emit->LDRSB(INDEX_UNSIGNED, m_dst_reg, X0, 0);
			else
				m_emit->LDRB(INDEX_UNSIGNED, m_dst_reg, X0, 0);
		break;
		case 16:
			if (m_sign_extend && !dont_extend)
				m_emit->LDRSH(INDEX_UNSIGNED, m_dst_reg, X0, 0);
			else
				m_emit->LDRH(INDEX_UNSIGNED, m_dst_reg, X0, 0);
		break;
		case 32:
			m_emit->LDR(INDEX_UNSIGNED, m_dst_reg, X0, 0);
		break;
		default:
			_assert_msg_(DYNA_REC, false, "Unknown size %d passed to MMIOReadCodeGenerator!", sbits);
		break;
		}
	}

	void LoadAddrMaskToReg(int sbits, const void* ptr, u32 mask)
	{
		m_emit->MOVI2R(X0, (u64)ptr);

		// If we do not need to mask, we can do the sign extend while loading
		// from memory. If masking is required, we have to first zero extend,
		// then mask, then sign extend if needed (1 instr vs. ~4).
		u32 all_ones = (1ULL << sbits) - 1;
		if ((all_ones & mask) == all_ones)
		{
			LoadToRegister(sbits, false);
		}
		else
		{
			LoadToRegister(sbits, true);
			m_emit->MOVI2R(W0, mask);
			m_emit->AND(m_dst_reg, m_dst_reg, W0, ArithOption(W0, ST_LSL, 0));
			if (m_sign_extend)
				m_emit->SBFM(m_dst_reg, m_dst_reg, 0, sbits - 1);
		}
	}

	void CallLambda(int sbits, const std::function<T(u32)>* lambda)
	{
		ARM64FloatEmitter float_emit(m_emit);

		m_emit->ABI_PushRegisters(m_gprs_in_use);
		float_emit.ABI_PushRegisters(m_fprs_in_use, X1);
			m_emit->MOVI2R(W1, m_address);
			m_emit->BLR(m_emit->ABI_SetupLambda(lambda));
		float_emit.ABI_PopRegisters(m_fprs_in_use, X1);
		m_emit->ABI_PopRegisters(m_gprs_in_use);

		if (m_sign_extend)
			m_emit->SBFM(m_dst_reg, W0, 0, sbits - 1);
		else
			m_emit->UBFM(m_dst_reg, W0, 0, sbits - 1);
	}

	ARM64XEmitter* m_emit;
	BitSet32 m_gprs_in_use;
	BitSet32 m_fprs_in_use;
	ARM64Reg m_dst_reg;
	u32 m_address;
	bool m_sign_extend;
};

void MMIOLoadToReg(MMIO::Mapping* mmio, Arm64Gen::ARM64XEmitter* emit,
                   BitSet32 gprs_in_use, BitSet32 fprs_in_use,
                   ARM64Reg dst_reg, u32 address, u32 flags)
{
	if (flags & BackPatchInfo::FLAG_SIZE_8)
	{
		MMIOReadCodeGenerator<u8> gen(emit, gprs_in_use, fprs_in_use, dst_reg,
							address, flags & BackPatchInfo::FLAG_EXTEND);
		mmio->GetHandlerForRead<u8>(address).Visit(gen);
	}
	else if (flags & BackPatchInfo::FLAG_SIZE_16)
	{
		MMIOReadCodeGenerator<u16> gen(emit, gprs_in_use, fprs_in_use, dst_reg,
							 address, flags & BackPatchInfo::FLAG_EXTEND);
		mmio->GetHandlerForRead<u16>(address).Visit(gen);
	}
	else if (flags & BackPatchInfo::FLAG_SIZE_32)
	{
		MMIOReadCodeGenerator<u32> gen(emit, gprs_in_use, fprs_in_use, dst_reg,
							 address, flags & BackPatchInfo::FLAG_EXTEND);
		mmio->GetHandlerForRead<u32>(address).Visit(gen);
	}
}

void MMIOWriteRegToAddr(MMIO::Mapping* mmio, Arm64Gen::ARM64XEmitter* emit,
                        BitSet32 gprs_in_use, BitSet32 fprs_in_use,
                        ARM64Reg src_reg, u32 address, u32 flags)
{
	if (flags & BackPatchInfo::FLAG_SIZE_8)
	{
		MMIOWriteCodeGenerator<u8> gen(emit, gprs_in_use, fprs_in_use, src_reg,
							address);
		mmio->GetHandlerForWrite<u8>(address).Visit(gen);
	}
	else if (flags & BackPatchInfo::FLAG_SIZE_16)
	{
		MMIOWriteCodeGenerator<u16> gen(emit, gprs_in_use, fprs_in_use, src_reg,
							 address);
		mmio->GetHandlerForWrite<u16>(address).Visit(gen);
	}
	else if (flags & BackPatchInfo::FLAG_SIZE_32)
	{
		MMIOWriteCodeGenerator<u32> gen(emit, gprs_in_use, fprs_in_use, src_reg,
							 address);
		mmio->GetHandlerForWrite<u32>(address).Visit(gen);
	}
}
