#include "stdafx.h"
#include "Binary.h"
#include "Exception.h"
#include "Core/StrBuf.h"

namespace code {

	Binary::Binary() {}

	Binary::Binary(Arena *arena, Listing *listing) {
		compileI(arena, listing, false);
	}

	Binary::Info Binary::compile(Arena *arena, Listing *listing) {
		Binary *b = new (arena) Binary();
		return b->compileI(arena, listing, false);
	}

	Binary::Info Binary::compileTransformed(Arena *arena, Listing *listing) {
		Binary *b = new (arena) Binary();
		return b->compileI(arena, listing, true);
	}

	Binary::Info Binary::compileI(Arena *arena, Listing *listing, Bool skipTransform) {
		flags = 0;
		if (listing->exceptionCleanup())
			flags |= ehClean;
		if (listing->exceptionCaught())
			flags |= ehCatch;
		if (listing->member)
			flags |= isMemberFn;

		Listing *tfm = listing;
		Array<Offset> *layout;
		if (skipTransform) {
			layout = new (this) Array<Offset>();
		} else {
			Arena::TransformInfo info = arena->transformInfo(listing);
			tfm = info.listing;
			layout = info.varLayout;
		}

		fillResultAndParams(tfm);
		fillBlocks(tfm);

		LabelOutput *labels = arena->labelOutput();
		arena->output(tfm, labels);

		fillTryBlocks(tfm, labels);

		if (tfm->meta().id < labels->offsets->count()) {
			metaOffset = labels->offsets->at(tfm->meta().id);
		} else {
			metaOffset = 0;
			WARNING(L"No metadata seems to have been generated by the backend.");
			WARNING(L"Exception cleanup will not work!");
		}

		CodeOutput *output = arena->codeOutput(this, labels);
		arena->output(tfm, output);

		runtime::codeUpdatePtrs(output->codePtr());
		set(output->codePtr(), output->tell());

#ifdef DEBUG
		// All backends should output the pointer to the binary like this.
		assert(codeBinary(output->codePtr()) == this);
#endif

		// Also, we should have an array in element 2. Replace it wit a plain GcArray for slightly
		// more compact representation, and to allow us to use the 'filled' member as we please.
		{
			GcCode *refs = runtime::codeRefs(output->codePtr());
#ifdef DEBUG
			assert(refs->refs[1].kind == GcCodeRef::ptrStorage);
			assert(as<Array<TObject *>>((RootObject *)refs->refs[1].pointer));
#endif
			Array<TObject *> *updaters = (Array<TObject *> *)refs->refs[1].pointer;
			GcArray<TObject *> *copy = runtime::allocArray<TObject *>(engine(), &pointerArrayType, updaters->count());
			for (Nat i = 0; i < copy->count; i++)
				copy->v[i] = updaters->at(i);
			// Mark it as unsorted.
			copy->filled = copy->count + 1;
			// Store it back.
			refs->refs[1].pointer = copy;
		}

		return Info(this, tfm, labels->offsets, layout);
	}

	void Binary::toS(StrBuf *to) const {
		*to << S("Binary object:");

		const nat columns = 16;
		const byte *code = (const byte *)address();
		if (!code) {
			*to << S(" <null>");
			return;
		}

		size_t size = runtime::codeSize(code);
		for (size_t i = 0; i < size; i++) {
			if (i % columns == 0) {
				*to << S("\n") << hex(i) << S(" -");
			}

			*to << S(" ") << hex(code[i]);
		}
	}

	void Binary::fillResultAndParams(Listing *src) {
		Array<code::Var> *params = src->allParams();

		resultAndParams = runtime::allocArray<TypeDesc *>(engine(), &pointerArrayType, params->count() + 1);
		resultAndParams->v[0] = src->result;

		for (Nat i = 0; i < params->count(); i++)
			resultAndParams->v[i + 1] = src->paramDesc(params->at(i));
	}

	TypeDesc *Binary::result() const {
		if (resultAndParams)
			return resultAndParams->v[0];
		return null;
	}

	Array<TypeDesc *> *Binary::params() const {
		Array<TypeDesc *> *result = new (this) Array<TypeDesc *>();
		if (resultAndParams) {
			result->reserve(Nat(resultAndParams->count - 1));
			for (size_t i = 1; i < resultAndParams->count; i++)
				result->push(resultAndParams->v[i]);
		}
		return result;
	}

	const GcType Binary::blockType = {
		GcType::tArray,
		null,
		null,
		sizeof(Variable),
		1,
		{ OFFSET_OF(Variable, varInfo) },
	};

	const GcType Binary::tryInfoArrayType = {
		GcType::tArray,
		null,
		null,
		sizeof(TryInfo),
		1,
		{ OFFSET_OF(TryInfo, type) },
	};

	void Binary::fillBlocks(Listing *src) {
		Array<code::Block> *srcBlocks = src->allBlocks();

		blocks = runtime::allocArray<Block *>(engine(), &pointerArrayType, srcBlocks->count());

		for (Nat i = 0; i < srcBlocks->count(); i++) {
			code::Block block = srcBlocks->at(i);
			Array<Var> *vars = src->allVars(block);

			// Count variables that need finalization, and variables with VarInfo.
			// We only need variables that need finalization for unwinding, we need
			// variables with VarInfo for hot reloading.
			size_t count = 0;
			for (Nat j = 0; j < vars->count(); j++) {
				const Var &v = vars->at(j);
				if (src->varInfo(v) || (src->freeOpt(v) & freeOnException))
					count++;
			}

			Block *b = (Block *)runtime::allocArray(engine(), &blockType, count);
			blocks->v[i] = b;
			b->parent = src->parent(block).key();

			size_t at = 0;
			for (Nat j = 0; j < vars->count(); j++) {
				const Var &v = vars->at(j);
				Nat flags = src->freeOpt(v);
				Listing::VarInfo *varInfo = src->varInfo(v);

				// We only include the ones we actually need.
				if (varInfo == null && (flags & freeOnException) == 0)
					continue;

				if (flags & freePtr) {
					// No additional flags needed, but we set sPtr for good measure.
					flags |= Variable::sPtr;
				} else if (v.size() == Size::sPtr) {
					flags |= Variable::sPtr;
				} else if (v.size() == Size::sByte) {
					flags |= Variable::sByte;
				} else if (v.size() == Size::sInt) {
					flags |= Variable::sInt;
				} else if (v.size() == Size::sLong) {
					flags |= Variable::sLong;
				} else {
					flags |= Variable::sUnknown;
					if (flags & freeOnException) {
						throw new (this) InvalidValue(S("Can only use bytes, integers, longs and pointers for ")
													S("variable cleanup. Specify 'freePtr' to get a pointer to ")
													S("the value instead!"));
					}
				}

				if (src->isParam(v))
					flags |= (1 + src->paramIndex(v)) << Variable::sParamShift;

				b->vars[at].id = v.key();
				b->vars[at].flags = flags;
				b->vars[at].varInfo = varInfo;
				at++;
			}
		}
	}

	void Binary::fillTryBlocks(Listing *src, LabelOutput *labels) {
		Nat count = 0;

		Array<code::Block> *blocks = src->allBlocks();
		for (Nat i = 0; i < blocks->count(); i++) {
			if (Array<Listing::CatchInfo> *info = src->catchInfo(blocks->at(i)))
				count += info->count();
		}

		if (count == 0) {
			tryBlocks = null;
			return;
		}

		tryBlocks = runtime::allocArray<TryInfo>(engine(), &tryInfoArrayType, count);
		Nat at = 0;
		for (Nat i = 0; i < blocks->count(); i++) {
			code::Block b = blocks->at(i);
			if (Array<Listing::CatchInfo> *info = src->catchInfo(b)) {
				for (Nat j = 0; j < info->count(); j++) {
					tryBlocks->v[at].blockId = code::Block(b).key();
					tryBlocks->v[at].resumeOffset = labels->offsets->at(info->at(j).resume.id);
					tryBlocks->v[at].type = info->at(j).type;
					at++;
				}
			}
		}
	}

	void Binary::cleanup(StackFrame &frame) {
		for (size_t i = frame.block; i != code::Block().key(); i = blocks->v[i]->parent) {
			Block *b = blocks->v[i];

			// Reverse order is common.
			for (size_t j = b->count; j > 0; j--) {
				cleanup(frame, b->vars[j - 1]);
			}
		}
	}

	Nat Binary::cleanup(StackFrame &frame, Nat until) {
		for (size_t i = frame.block; i != code::Block().key(); i = blocks->v[i]->parent) {
			Block *b = blocks->v[i];

			// Reverse order is common.
			for (size_t j = b->count; j > 0; j--) {
				cleanup(frame, b->vars[j - 1]);
			}

			// Done?
			if (i == until)
				return Nat(blocks->v[i]->parent);
		}

		// Outside of all blocks.
		return code::Block().key();
	}

	VarCleanup *Binary::cleanupInfo() {
		// Element #0 is the total size, then VarCleanup-instances start.
		byte *data = (byte *)address() + metaOffset + sizeof(void *);
		return (VarCleanup *)data;
	}

	MAYBE(Reference *) Binary::findReferenceByOffset(Nat offset) {
		GcCode *refs = runtime::codeRefs((void *)address());

		size_t refSlot = 0;
		for (size_t i = 0; i < refs->refCount; i++) {
			const GcCodeRef &ref = refs->refs[i];
			if (ref.offset == offset) {
				refSlot = i;
				break;
			}
		}

		// Note: Id #0 is reserved for a pointer to us.
		if (refSlot == 0)
			return null;

		return findReferenceBySlot(Nat(refSlot));
	}

	class CompareUpdater {
	public:
		bool operator() (const TObject *lhs, const TObject *rhs) const {
			CodeUpdater *l = (CodeUpdater *)lhs;
			CodeUpdater *r = (CodeUpdater *)rhs;
			return l->getSlot() < r->getSlot();
		}
		bool operator() (const TObject *lhs, Nat rhs) const {
			CodeUpdater *l = (CodeUpdater *)lhs;
			return l->getSlot() < rhs;
		}
	};

	static void sortUpdaters(GcArray<TObject *> *updaters) {
		// First, partition according to whether it is a CodeUpdater or something else. Let 'filled'
		// keep track of how many elements are used for this.
		updaters->filled = 0;
		for (size_t i = 0; i < updaters->count; i++) {
			if (as<CodeUpdater>(updaters->v[i])) {
				// Cheap enough to swap the same pointer as a no-op, don't bother checking.
				std::swap(updaters->v[updaters->filled++], updaters->v[i]);
			}
		}

		// Now, we can sort the elements:
		std::sort(updaters->v, updaters->v + updaters->filled, CompareUpdater());
	}

	MAYBE(Reference *) Binary::findReferenceBySlot(Nat refSlot) {
		GcCode *refs = runtime::codeRefs((void *)address());

		// Element #1 is a reference to an array of TObjects that contains the data we are after:
		GcArray<TObject *> *updaters = (GcArray<TObject *> *)refs->refs[1].pointer;

		// If 'filled' is larger than 'count', then we need to sort the array first:
		if (updaters->filled > updaters->count)
			sortUpdaters(updaters);

		// Now, we can sort everything!
		TObject **begin = updaters->v;
		TObject **end = begin + updaters->filled;
		TObject **found = std::lower_bound(begin, end, refSlot, CompareUpdater());
		if (found == end)
			return null;

		CodeUpdater *elem = (CodeUpdater *)*found;
		if (elem->getSlot() == refSlot)
			return elem;
		return null;
	}

	void Binary::cleanup(StackFrame &frame, Variable &v) {
		if (v.flags & freeOnException) {
			VarCleanup *vars = cleanupInfo();

			VarCleanup &now = vars[v.id];
			void *freeFn = now.function;
			int offset = now.offset;
			nat activeAfter = now.activeAfter;

			// If not active, we don't destroy it.
			if (frame.activation < activeAfter)
				return;

			void *ptr = frame.toPtr(offset);

			if (v.flags & freeIndirection)
				ptr = *(void **)ptr;

			typedef void (*FPtr)(void *v);
			typedef void (*FByte)(Byte v);
			typedef void (*FInt)(Int v);
			typedef void (*FLong)(Long v);

			if (v.flags & freePtr) {
				FPtr p = (FPtr)freeFn;
				(*p)(ptr);
			} else {
				switch (v.flags & Variable::sMask) {
				case Variable::sUnknown:
					break;
				case Variable::sPtr: {
					FPtr p = (FPtr)freeFn;
					(*p)(*(void **)ptr);
					break;
				}
				case Variable::sByte: {
					FByte p = (FByte)freeFn;
					(*p)(*(Byte *)ptr);
					break;
				}
				case Variable::sInt: {
					FInt p = (FInt)freeFn;
					(*p)(*(Int *)ptr);
					break;
				}
				case Variable::sLong: {
					FLong p = (FLong)freeFn;
					(*p)(*(Long *)ptr);
					break;
				}
				}
			}
		}
	}

	bool Binary::hasCatch(Nat active, RootObject *exception, Resume &resume) {
		struct Compare {
			inline bool operator() (const TryInfo &l, Nat r) const {
				return l.blockId < r;
			}
		};

		if (!tryBlocks)
			return false;

		for (size_t block = active; block != code::Block().key(); block = blocks->v[block]->parent) {
			TryInfo *end = tryBlocks->v + tryBlocks->count;
			TryInfo *found = std::lower_bound(tryBlocks->v, end, Nat(block), Compare());

			// Check all possible matches.
			for (; found != end && found->blockId == block; found++) {
				if (runtime::isA(exception, found->type)) {
					// Find where to resume.
					byte *data = (byte *)address();
					resume.ip = data + found->resumeOffset;
					resume.stackOffset = stackOffset();

					// Remember how far to clean.
					resume.cleanUntil = block;

					return true;
				}
			}
		}

		return false;
	}

	ptrdiff_t Binary::rawStackOffset() const {
		byte *data = (byte *)address();
		// First entry in the metadata table is the stack offset.
		ptrdiff_t *table = (ptrdiff_t *)(data + metaOffset);
		return table[0];
	}

	ptrdiff_t Binary::stackOffset() const {
		ptrdiff_t raw = rawStackOffset();
		// If LSB is set, then this is just a size.
		if (raw & 0x1)
			return 0;
		else
			return raw;
	}

	size_t Binary::stackSize() const {
		ptrdiff_t raw = rawStackOffset();
		raw &= ~ptrdiff_t(0x1);
		return abs(raw);
	}

}