#include "stdafx.h"
#include "ControlFlow.h"
#include "Code/Listing.h"
#include "Code/Arena.h"
#include "Gc/CodeTable.h"
#include "Gc/DwarfTable.h"
#include "Type.h"
#include "BuiltInSource.h"

namespace storm {

	static void flattenHelper(const ControlFlowItem &current, Array<ControlFlowItem> *out) {
		if (current.isCall()) {
			out->push(current);
		} else if (current.isLoop()) {
			Array<ControlFlowItem> *body = current.loop();

			out->push(ControlFlowItem(current.offset(), current.endOffset(), body));

			for (Nat i = 0; i < body->count(); i++) {
				flattenHelper(body->at(i), out);
			}
		}
	}

	Array<ControlFlowItem> *ControlFlowItem::flatten() const {
		Array<ControlFlowItem> *out = new (data) Array<ControlFlowItem>();
		flattenHelper(*this, out);
		return out;
	}

	Array<ControlFlowItem> *flatten(Array<ControlFlowItem> *items) {
		Array<ControlFlowItem> *out = new (items) Array<ControlFlowItem>();
		for (Nat i = 0; i < items->count(); i++)
			flattenHelper(items->at(i), out);
		return out;
	}

	void ControlFlowItem::toS(StrBuf *to) const {
		*to << S("@") << offset() << S(" ");

		if (hasFunction())
			*to << function()->identifier();
		else if (hasBuiltIn())
			*to << builtIn()->title();
		else if (isLoop())
			*to << S("-> ") << endOffset() << S(" ") << loop();

		statusSuffix(to, status());
	}

	void ControlFlowItem::statusSuffix(StrBuf *to, Status status) {
		switch (status) {
		case none:
			break;
		case removed:
			*to << S(" (removed)");
			break;
		}
	}

	void ControlFlowItem::replaceOffsets(Map<Nat, Nat> *replace) {
		if (replace->has(offset())) {
			offset(replace->get(offset()));
		}

		if (isLoop()) {
			Array<ControlFlowItem> *item = loop();
			for (Nat i = 0; i < item->count(); i++)
				item->at(i).replaceOffsets(replace);
		}
	}

	void replaceOffsets(Array<ControlFlowItem> *items, Map<Nat, Nat> *replace) {
		for (Nat i = 0; i < items->count(); i++)
			items->at(i).replaceOffsets(replace);
	}

	static ControlFlowItem filterFunctions(Nat offset, Function *f) {
		if (!f)
			return ControlFlowItem();

		// Don't include destructors - they are almost always not present explcitly in the source listings.
		if (*f->name == Type::DTOR)
			return ControlFlowItem();
		return ControlFlowItem(offset, f);
	}

	// Find a function object from a RefSource:
	static ControlFlowItem findFunction(Nat offset, MAYBE(code::RefSource *) source) {
		if (NamedSource *n = as<NamedSource>(source)) {
			return filterFunctions(offset, as<Function>(n->named()));
		} else if (BuiltInSource *b = as<BuiltInSource>(source)) {
			return ControlFlowItem(offset, b);
		} else {
			return ControlFlowItem();
		}
	}

	// Find a function object from a pointer to the code.
	static ControlFlowItem findFunctionFromCode(Nat offset, void *code, Nat refId) {
		code::Binary *binary = code::codeBinary(code);
		if (code::Reference *ref = binary->findReferenceBySlot(refId)) {
			return findFunction(offset, ref->source());
		} else {
			return ControlFlowItem();
		}
	}

	// Find a function object from a Ref. Should correspond roughly to 'findFunctionFromCode'.
	static ControlFlowItem findFunctionFromRef(Nat offset, code::Ref ref) {
		return findFunction(offset, ref.source());
	}


	class CodeRefSort {
	public:
		GcCodeRef *src;

		CodeRefSort(GcCodeRef *src) : src(src) {}

		bool operator() (size_t a, size_t b) const {
			return src[a].offset < src[b].offset;
		}
	};

	Array<ControlFlowItem> *controlFlowList(code::Binary *code) {
		return controlFlowListRaw((void *)code->address());
	}

	// Generate a control flow list from a pre-compiled function:
	Array<ControlFlowItem> *controlFlowListRaw(void *code) {
		code::Binary *b = code::codeBinary(code);
		Engine &e = b->engine();

		GcCode *refs = runtime::codeRefs(code);

		// First, sort the refs. They are not necessarily in increasing order.
		vector<size_t> order(refs->refCount, 0);
		for (size_t i = 0; i < order.size(); i++)
			order[i] = i;

		std::sort(order.begin(), order.end(), CodeRefSort(refs->refs));

		// Then, traverse them and create our result:
		Array<ControlFlowItem> *result = new (e) Array<ControlFlowItem>();

		for (size_t i = 0; i < refs->refCount; i++) {
			GcCodeRef &ref = refs->refs[i];

			if (ref.kind == GcCodeRef::jump) {
				// Create a function item if we can find the underlying function.
				ControlFlowItem call = findFunctionFromCode(Nat(ref.offset), code, Nat(i));
				if (call.isCall())
					result->push(call);
			} else if (ref.kind == GcCodeRef::backEdge) {
				// Look back in 'result' to find things to include in the loop.
				size_t target = size_t(ref.pointer);

				Nat includeUntil = result->count();
				while (includeUntil > 0) {
					if (result->at(includeUntil - 1).offset() <= target)
						break;

					includeUntil--;
				}

				Array<ControlFlowItem> *sub = new (e) Array<ControlFlowItem>();
				sub->reserve(result->count() - includeUntil);
				while (result->count() > includeUntil) {
					sub->push(result->last());
					result->pop();
				}
				sub->reverse();

				result->push(ControlFlowItem(Nat(target), Nat(ref.offset), sub));
			}
		}

		return result;
	}

	// Generate a control flow list from a Listing:
	Array<ControlFlowItem> *controlFlowList(code::Listing *code) {
		using namespace code;

		// Just go through the listing, keep track of labels we have found in order to find back-edges.
		Array<Nat> *labelTarget = new (code) Array<Nat>(code->labelCount(), code->count());

		Array<ControlFlowItem> *result = new (code) Array<ControlFlowItem>();

		for (Nat i = 0; i < code->count(); i++) {
			Instr *instr = code->at(i);

			// Mark labels:
			if (Array<Label> *labels = code->labels(i)) {
				for (Nat j = 0; j < labels->count(); j++) {
					labelTarget->at(labels->at(j).key()) = i;
				}
			}

			// Look at the instruction. We are only interested in 'fnCall' and 'jmp':
			switch (instr->op()) {
			case op::fnCall:
			case op::fnCallRef:
				if (instr->src().type() == opReference) {
					ControlFlowItem call = findFunctionFromRef(i, instr->src().ref());
					if (call.isCall())
						result->push(call);
				}
				break;
			case op::jmp:
			case op::jmpBlock:
				if (instr->dest().type() == opLabel) {
					Nat target = labelTarget->at(instr->dest().label().key());
					// Not a back-edge?
					if (target >= i)
						break;

					// Copy elements, and replace:
					Nat includeUntil = result->count();
					while (includeUntil > 0) {
						if (result->at(includeUntil - 1).offset() <= target)
							break;

						includeUntil--;
					}

					Array<ControlFlowItem> *sub = new (code) Array<ControlFlowItem>();
					sub->reserve(result->count() - includeUntil);
					while (result->count() > includeUntil) {
						sub->push(result->last());
						result->pop();
					}
					sub->reverse();

					result->push(ControlFlowItem(target, i, sub));
				}
				break;
			}
		}

		return result;
	}

}