#include "stdafx.h"
#include "Arena.h"
#include "Reg.h"
#include "X86/Arena.h"
#include "X64/Arena.h"
#include "Arm64/Arena.h"
#include "Core/Str.h"

#include "Listing.h"
#include "Binary.h"

namespace code {

	Arena::Arena() {}

	Ref Arena::external(const wchar *name, const void *ptr) const {
		return Ref(externalSource(name, ptr));
	}

	RefSource *Arena::externalSource(const wchar *name, const void *ptr) const {
		RefSource *src = new (this) StrRefSource(name);
		src->setPtr(ptr);
		return src;
	}

	Listing *Arena::transform(Listing *l) const {
		return transformInfo(l).listing;
	}

	void Arena::removeFnRegs(RegSet *from) const {
		from->remove(ptrA);
		from->remove(ptrB);
		from->remove(ptrC);
	}

	Instr *Arena::saveFnResultReg(Reg reg, Operand to) const {
		return mov(engine(), to, reg);
	}

	Instr *Arena::restoreFnResultReg(Reg reg, Operand from) const {
		return mov(engine(), reg, from);
	}


#if defined(X86) && defined(WINDOWS)
	Arena *arena(EnginePtr e) {
		return new (e.v) x86::Arena();
	}
#elif defined(X64) && defined(WINDOWS)
	Arena *arena(EnginePtr e) {
		return new (e.v) x64::WindowsArena();
	}
#elif defined(X64) && defined(POSIX)
	Arena *arena(EnginePtr e) {
		return new (e.v) x64::PosixArena();
	}
#elif defined(ARM64) && defined(POSIX)
	Arena *arena(EnginePtr e) {
		return new (e.v) arm64::Arena();
	}
#else
#error "Please note which is the default arena for your platform."
#endif

	Binary *codeBinaryImpl(GcCode *refs) {
		return (Binary *)refs->refs[0].pointer;
	}

	Binary *codeBinary(const void *fn) {
		// All backends do this.
		return codeBinaryImpl(runtime::codeRefs((void *)fn));
	}

	void Arena::updateEhInfo(const void *function, size_t offset, void *framePointer) {
		// Nothing to do in the default implementation.
	}

	Arena::Skeleton::Skeleton(Listing *listing)
		: listing(listing), currentBlock(0), currentActivation(0), accessMode(-1) {
		savedRegs = new (this) Array<Operand>();
		savedLocs = new (this) Array<Operand>();
		varOffsets = new (this) Array<Operand>();
		extraMetadata = new (this) Array<Offset>();
	}

	static Size makeSize(Nat size, Int currentOffset) {
		Size r(size);
		if ((size & 0x3) || (currentOffset & 0x3)) {
			return r.alignedAs(Size::sByte);
		} else if ((size & 0x7) || (currentOffset & 0x7)) {
			return r.alignedAs(Size::sInt);
		} else {
			return r.alignedAs(Size::sPtr);
		}
	}

	// State for keeping track of added variables.
	class VarState {
	public:
		// Create, initialize with the maximum offset.
		VarState(Listing *to, Int startOffset, Nat minAlign)
			: to(to), currentOffset(startOffset), minAlign(minAlign) {}

		// Add a new variable.
		Var add(FreeOpt freeOpt, Nat size, Listing::VarInfo *info, Operand free, Int offset) {
			if (size == 0) {
				// Emit this variable as a part of the previous one:
				size = Nat(currentOffset - offset);
				currentOffset = offset;

				Var v = to->createVar(to->root(), makeSize(size), free, freeOpt);
				to->varInfo(v, info);
				return v;
			} else {
				// Output any necessary padding:
				Nat alignedSz = roundUp(size, minAlign);
				output(offset + alignedSz, false);

				// Emit a variable now.
				currentOffset -= alignedSz;
				Var v = to->createVar(to->root(), makeSize(size), free, freeOpt);
				to->varInfo(v, info);
				return v;
			}
		}

		// Finish any remaining variables. Emit empty space until the specified offset.
		void done(Int targetOffset) {
			output(targetOffset, true);
		}

	private:
		Listing *to;
		Int currentOffset;
		Nat minAlign;

		void output(Int offset, Bool last) {
			Int size = currentOffset - offset;
			currentOffset = offset;

			if (size > 0) {
				Block parent = to->root();
				// For the last one: create a separate block to make the space usable for other things.
				if (last)
					parent = to->createBlock(to->root());
				to->createVar(parent, makeSize(Nat(size)));
			}
		}

		Size makeSize(Nat size) {
			return code::makeSize(size, currentOffset);
		}
	};

	Arena::Skeleton *Arena::frameSkeletonHead(Binary *binary) {
			Listing *l = new (this) Listing(binary->isMember(), binary->result());
			Array<TypeDesc *> *parameters = binary->params();
			for (Nat i = 0; i < parameters->count(); i++)
				l->createParam(parameters->at(i));

			return new (this) Skeleton(l);
	}

	static Nat toNat(Size s, Bool is64) {
		if (is64)
			return s.size64();
		else
			return s.size32();
	}

	static void frameSkeletonTailImpl(Binary *binary, Arena::Skeleton *result,
									Int maxOffset, Int minOffset,
									Nat minAlign, Bool is64) {
		Listing *l = result->listing;
		Array<Var> *paramVars = l->allParams();

		// Find variables that are active.
		GcArray<Binary::Block *> *blocks = binary->blockInfo();
		VarCleanup *cleanup = binary->cleanupInfo();

		// Collect all variables we need to process. Note that we traverse from the current node
		// towards the root, which is why we add nodes in reverse order.
		vector<std::pair<Nat, Nat>> ids;
		for (size_t i = result->currentBlock; i < blocks->count; i = blocks->v[i]->parent) {
			Binary::Block *current = blocks->v[i];

			for (Nat varId = Nat(current->count); varId > 0; varId--) {
				const Binary::Variable &v = current->vars[varId - 1];

				// Skip parameters entirely in this step. They are already handled.
				if (v.flags & Binary::Variable::sParamMask) {
					// However, we can add take its metadata and add that.
					Nat paramId = (v.flags & Binary::Variable::sParamMask) >> Binary::Variable::sParamShift;
					paramId--;

					if (v.varInfo)
						l->varInfo(paramVars->at(paramId), v.varInfo);

					continue;
				}

				FreeOpt freeOpts = FreeOpt(v.flags & Binary::Variable::sFreeOptMask);
				if ((freeOpts & freeOnException) || v.varInfo)
					ids.push_back(std::make_pair(Nat(i), varId - 1));
			}
		}

		// Traverse the offsets and add them:
		VarState varState(l, maxOffset, minAlign);

		for (size_t i = ids.size(); i > 0; i--) {
			Binary::Block *current = blocks->v[ids[i - 1].first];
			const Binary::Variable &v = current->vars[ids[i - 1].second];
			const VarCleanup &c = cleanup[v.id];

			// Store the variable.
			Nat size = 0;
			if ((v.flags & freePtr) == 0) {
				if (v.flags & Binary::Variable::sPtr)
					size = toNat(Size::sPtr, is64);
				else if (v.flags & Binary::Variable::sByte)
					size = toNat(Size::sByte, is64);
				else if (v.flags & Binary::Variable::sInt)
					size = toNat(Size::sInt, is64);
				else if (v.flags & Binary::Variable::sLong)
					size = toNat(Size::sLong, is64);
			}

			FreeOpt freeOpts = FreeOpt(v.flags & Binary::Variable::sFreeOptMask);
			Operand freeOp;
			if (c.function) {
				// Find the function that was stored here:
				size_t offset = size_t(&c.function) - size_t(binary->address());
				Reference *found = binary->findReferenceByOffset(Nat(offset));
				if (found)
					freeOp = Operand(found);
			}

			if (c.activeAfter >= result->currentActivation) {
				// If it is active, clear its 'freeInactive' flag to make it immediately active.
				freeOpts &= ~freeInactive;
			}

			Var added = varState.add(freeOpts, size, v.varInfo, freeOp, c.offset);
			while (added.key() >= result->varOffsets->count())
				result->varOffsets->push(Operand());
			result->varOffsets->at(added.key()) = xRel(added.size(), ptrFrame, Offset(c.offset));
		}

		varState.done(minOffset);
	}

	void Arena::frameSkeletonTail(Binary *binary, Skeleton *result, Nat extraWords, Nat minAlign, Bool is64) {
		frameSkeletonTailImpl(binary, result,
							-Int(extraWords * toNat(Size::sPtr, is64)),
							Int(binary->stackOffset()),
							minAlign, is64);
	}

	void Arena::frameSkeletonTailBelow(Binary *binary, Skeleton *result, Nat extraBelow, Nat extraAbove,
									Nat minAlign, Bool is64) {
		frameSkeletonTailImpl(binary, result,
							Int(binary->stackSize() - extraAbove * toNat(Size::sPtr, is64)),
							Int(extraBelow * toNat(Size::sPtr, is64)),
							minAlign, is64);
	}

}