#include "stdafx.h"
#include "ReplaceContext.h"
#include "ExactPart.h"
#include "Engine.h"
#include "Package.h"
#include "Code.h"

namespace storm {

	ReplaceContext::ReplaceContext() {
		typeEq = new (this) Map<Type *, Type *>();
	}

	Bool ReplaceContext::same(Type *a, Type *b) {
		return typeEq->get(a, a) == typeEq->get(b, b);
	}

	Type *ReplaceContext::normalize(Type *t) {
		if (t)
			return typeEq->get(t, t);
		else
			return null;
	}

	Value ReplaceContext::normalize(Value v) {
		v.type = typeEq->get(v.type, v.type);
		return v;
	}

	Set<Type *> *ReplaceContext::allNewTypes() {
		Set<Type *> *types = new (this) Set<Type *>();
		for (Map<Type *, Type *>::Iter i = typeEq->begin(), end = typeEq->end(); i != end; ++i)
			types->put(i.k());
		return types;
	}

	void ReplaceContext::buildTypeEquivalence(NameSet *oldRoot, NameSet *newRoot) {
		currentOldRoot = oldRoot;
		try {
			buildTypeEqRec(oldRoot, newRoot);
			currentOldRoot = null;
		} catch (...) {
			currentOldRoot = null;
			throw;
		}
	}

	void ReplaceContext::addEquivalence(Type *oldType, Type *newType) {
		typeEq->put(newType, oldType);
		if (currentOldRoot) {
			buildTypeEqRec(oldType, newType);
		}
	}

	void ReplaceContext::buildTypeEqRec(NameSet *currOld, NameSet *currNew) {
		currNew->forceLoad();
		for (NameSet::Iter i = currNew->begin(), end = currNew->end(); i != end; ++i) {
			Named *newEntity = i.v();

			// Try to find a matching old entity:
			ExactPart *p = new (this) ExactPart(newEntity->name, resolveParams(currentOldRoot, newEntity->params));
			Named *oldEntity = currOld->find(p, Scope());
			if (!oldEntity)
				continue;

			// Ask the entity to find mark all equivalent types.
			newEntity->findEquivalentTypes(oldEntity, this);
		}
	}

	Array<Value> *ReplaceContext::resolveParams(NameSet *root, Array<Value> *params) {
		if (params->empty())
			return params;

		params = new (this) Array<Value>(*params);
		for (Nat i = 0; i < params->count(); i++)
			params->at(i) = resolveParam(root, params->at(i));

		return params;
	}

	Value ReplaceContext::resolveParam(NameSet *root, Value param) {
		if (param == Value())
			return param;

		// Do we know of this parameter already?
		if (Type *t = typeEq->get(param.type, null)) {
			param.type = t;
			return t;
		}

		// No, we need to look it up ourselves. It has probably not been processed yet.
		Array<Named *> *prev = new (this) Array<Named *>();
		Named *current = param.type;
		while (current && current == root) {
			prev->push(current);
			current = as<Named>(current->parent());
		}

		NameSet *back = as<NameSet>(current);

		// No luck.
		if (!back)
			return param;

		// Now, try to match each name in 'prev' inside 'root'.
		while (prev->any() && back) {
			Named *n = prev->last();
			ExactPart *p = new (this) ExactPart(n->name, resolveParams(root, n->params));
			back = as<NameSet>(back->find(p, Scope(back)));

			prev->pop();
		}

		if (Type *t = as<Type>(back))
			param.type = t;
		return param;
	}

}
