#include "stdafx.h"
#include "Arena.h"
#include "Output.h"
#include "Listing.h"
#include "Remove64.h"
#include "RemoveInvalid.h"
#include "LayoutVars.h"
#include "Asm.h"
#include "AsmOut.h"
#include "Code/WindowsEh/Seh.h"
#include "Code/FnState.h"

namespace code {
	namespace x86 {

		static bool has64(Listing *in) {
			for (nat i = 0; i < in->count(); i++) {
				if (in->at(i)->size() == Size::sLong)
					return true;
			}
			return false;
		}

		Arena::Arena() {}

		Arena::TransformInfo Arena::transformInfo(Listing *l) const {
#if defined(WINDOWS) && defined(X86)
			code::eh::activateWindowsInfo(engine());
#endif

			if (has64(l)) {
				// Replace any 64-bit operations with 32-bit corresponding operations.
				l = code::transform(l, this, new (this) Remove64());
			}

			// Transform any unsupported op-codes into sequences of other op-codes. Eg. referencing
			// memory twice or similar.
			l = code::transform(l, this, new (this) RemoveInvalid());

			// Expand variables and function calls as well as function prolog and epilog. We need to
			// know all used registers for this to work, so it has to be run after the previous
			// transforms.
			LayoutVars *layout = new (this) LayoutVars();
			l = code::transform(l, this, layout);

			return TransformInfo(l, layout->layout);
		}

		void Arena::output(Listing *src, Output *to) const {
			code::x86::output(src, to);
			to->finish();
		}

		LabelOutput *Arena::labelOutput() const {
			return new (this) LabelOutput(4);
		}

		CodeOutput *Arena::codeOutput(Binary *owner, LabelOutput *size) const {
			return new (this) CodeOut(owner, size->offsets, size->size, size->refs);
		}

		void Arena::removeFnRegs(RegSet *from) const {
			code::Arena::removeFnRegs(from);
			from->remove(ptrD);
			// esi, edi (and actually ebx as well) are preserved.
		}

		RegSet *Arena::fnResultRegs() const {
			RegSet *result = new (this) RegSet();
			result->put(eax);
			result->put(edx);
			// Note: We use xmm0 as a placeholder for FP(0).
			result->put(xmm0);
			return result;
		}

		Instr *Arena::saveFnResultReg(Reg reg, Operand to) const {
			if (reg == xmm0) {
				return fstp(engine(), to);
			} else {
				return code::Arena::saveFnResultReg(reg, to);
			}
		}

		Instr *Arena::restoreFnResultReg(Reg reg, Operand from) const {
			if (reg == xmm0) {
				return fld(engine(), from);
			} else {
				return code::Arena::restoreFnResultReg(reg, from);
			}
		}


		Listing *Arena::redirect(Bool member, TypeDesc *result, Array<TypeDesc *> *params, Ref fn, Operand param) {
			Listing *l = new (this) Listing(this, member, result);

			// Add parameters. We only want to free them if we get an exception.
			for (Nat i = 0; i < params->count(); i++)
				l->createParam(params->at(i), freeOnException | freePtr);

			// Output the function.
			*l << prolog();

			if (!param.empty())
				*l << fnParam(ptrDesc(engine()), param);
			// It does not matter if the called function is a member in this case.
			*l << fnCall(fn, false, ptrDesc(engine()), ptrA);

			*l << epilog(); // preserves ptrA
			*l << jmp(ptrA);

			return l;
		}

		Listing *Arena::engineRedirect(TypeDesc *result, Array<TypeDesc *> *params, Ref fn, Operand engine) {
			Listing *l = new (this) Listing(this);

			if (resultParam(result)) {
				// The result is returned using a hidden parameter. The first parameter is, and has
				// to be, a pointer to the returned object. Here, the old return pointer and the
				// return value pointer are stored in the 'returnData' member of EnginePtr.
				*l << mov(ptrA, ptrRel(ptrStack, Offset::sPtr)); // Read the return value ptr.
				*l << push(engine);
				*l << push(ptrA); // Store the return value ptr once more.
			} else {
				// The result is returned in a register. The old pointer and the constant 0 will fit
				// inside the 'returnData' member of EnginePtr.
				*l << push(ptrConst(Offset(0)));
				*l << push(engine);
			}

			*l << call(fn, Size());
			*l << add(ptrStack, ptrConst(Size::sPtr * 2));
			*l << ret(Size());

			return l;
		}

		Nat Arena::firstParamId(MAYBE(TypeDesc *) desc) {
			if (!desc)
				return 1;

			// No difference with regards to the return value.
			return 0;
		}

		Operand Arena::firstParamLoc(Nat id) {
			if (id != 0)
				return Operand();

			return ptrRel(ptrStack, Offset::sPtr);
		}

		Reg Arena::functionDispatchReg() {
			return ptrA;
		}

		struct ModRmInfo {
			byte reg;
			byte mem;
			Int disp;
		};

		static ModRmInfo skipModRm(const byte *&code) {
			byte modrm = *(code++);
			ModRmInfo info = { byte((modrm >> 3) & 0x7), byte(modrm & 0x7), 0 };

			if ((modrm & 0xC0) == 0xC0) {
				// No extra displacement, no SIB.
				return info;
			}

			if ((modrm & 0x07) == 0x04) {
				// SIB byte, skip it.
				code++;
			}

			if ((modrm & 0xC0) == 0x80) {
				// 32-bit displacement.
				info.disp = *(const Int *)code;
				code += 4;
			} else if ((modrm & 0xC0) == 0x40) {
				// 8-bit displacement.
				info.disp = *(const signed char *)code;
				code++;
			} else if ((modrm & 0xC7) == 0x05) {
				// Special case of no displacement, 32-bit.
				info.disp = *(const Int *)code;
				code += 4;
			} else {
				// No displacement.
			}
			return info;
		}

		Arena::Skeleton *Arena::compatibleFrameSkeleton(Binary *binary, Nat offset) {
			Arena::Skeleton *result = frameSkeletonHead(binary);

			// Figure out which registers were saved where:
			Array<Operand> *preservedRegs = result->savedRegs;
			Array<Operand> *preservedLocs = result->savedLocs;

			// Map from register number to our representation.
			const Reg regMap[8] = { ptrA, ptrC, ptrD, ptrB, ptrStack, ptrFrame, ptrSi, ptrDi };

			Bool usingEh = false;

			// We don't have any metadata here. However, the prolog is fairly static looking, so we
			// can just parse that:
			const byte *code = (const byte *)binary->address();
			const byte *codeEnd = code + binary->size();
			while (code < codeEnd) {
				byte op = *(code++);
				if ((op & 0xF8) == 0x50) {
					// Push reg, no extra.
				} else if ((op & 0xF8) == 0x58) {
					// Pop reg, no extra.
				} else if (op == 0x89) {
					// Mov r/m32, r32
					ModRmInfo info = skipModRm(code);
					if (info.mem == 0x05 /* rbp */ && info.disp != 0) {
						if (info.reg != 0) { // We never preserve ptrA. That is a part of the SEH setup.
							preservedRegs->push(regMap[info.reg]);
							preservedLocs->push(ptrRel(ptrFrame, Offset(info.disp)));
						}
					}
				} else if (op == 0x8B) {
					// Mov r32, r/m32. We are never interested in this one.
					skipModRm(code);
				} else if (op == 0xC7) {
					// Mov 32-bit immediate.
					skipModRm(code);
					code += 4;
				} else if (op == 0x8D) {
					// LEA
					skipModRm(code);
				} else if (op == 0x83) {
					// Misc aritmetic operations with small immediate.
					skipModRm(code);
					code++;
				} else if (op == 0x81) {
					// Misc aritmetic operations with large immediate.
					skipModRm(code);
					code += 4;
				} else if (op == 0x64) {
					// FS override prefix
					usingEh = true;
				} else {
					// Stop as soon as we find something we don't understand. That means we have
					// passed the prolog (e.g. xor, call, ...).
					break;
				}
			}

			Nat extraWords = preservedRegs->count();
			if (usingEh) {
				extraWords += EH_WORDS;

				// Add extra offsets to save for the current block ptr...
				result->extraMetadata->push(Offset::sPtr * -1);
				// ...and the self pointer.
				result->extraMetadata->push(Offset::sPtr * -2);
			}

			// Find the current block and active piece:
			Nat active = findFunctionState(binary->address(), offset);
			decodeFnState(active, result->currentBlock, result->currentActivation);

			// Finish generating the skeleton.
			frameSkeletonTail(binary, result, extraWords, 1, false);

			return result;
		}

		void Arena::updateEhInfo(const void *function, size_t offset, void *framePointer) {
			// On X86, we need to update the SEH entry on the stack if there is one.
			Binary *binary = codeBinary(function);
			if (!binary->exceptionAware())
				return;

			Nat fnState = findFunctionState(function, offset);

			size_t *fp = (size_t *)framePointer;
			// Self pointer, used to find the correct EH information.
			fp[-3] = size_t(function);
			// Active block + activation.
			fp[-2] = fnState;
		}

		void Arena::resizeStackFrame(Listing *out, Reg tmpReg, Binary *newSz) {
			Int offset = Int(newSz->stackOffset());
			tmpReg = asSize(tmpReg, Size::sPtr);
			*out << mov(tmpReg, ptrFrame);
			if (offset > 0)
				*out << add(tmpReg, ptrConst(Nat(offset)));
			else if (offset < 0)
				*out << sub(tmpReg, ptrConst(Nat(-offset)));
			*out << mov(ptrStack, tmpReg);
		}

	}
}