/* Resp: A programming language
 * Copyright (C) 2008-2009 David Robillard <dave@drobilla.net>
 *
 * Resp is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 *
 * Resp is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General
 * Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with Resp.  If not, see <http://www.gnu.org/licenses/>.
 */

/** @file
 * @brief Lift functions (compilation pass 1)
 * After this pass:
 *  - All function definitions are top-level
 *  - All references to functions are replaced with references to
 *    a closure (a tuple with the function and necessary context)
 */

#include "resp.hpp"

using namespace std;

static AST*
lift_symbol(CEnv& cenv, Code& code, ASymbol* sym) throw()
{
	const std::string& cppstr = sym->cppstr;
	if (!cenv.liftStack.empty() && cppstr == cenv.name(cenv.liftStack.top().fn)) {
		return cenv.penv.sym("_me"); // Reference to innermost function
	} else if (!cenv.penv.handler(true, cppstr)
			&& !cenv.penv.handler(false, cppstr)
			&& !cenv.code.innermost(sym)) {

		const int32_t index = cenv.liftStack.top().index(sym);

		// Replace symbol with code to access free variable from closure
		return tup<ATuple>(sym->loc, cenv.penv.sym("."),
				cenv.penv.sym("_me"),
				new ALiteral<int32_t>(index, Cursor()),
				NULL);
	} else {
		return sym;
	}
}

static AST*
lift_fn(CEnv& cenv, Code& code, ATuple* fn) throw()
{
	ATuple* impl = new ATuple(*fn);
	const string fnName      = cenv.name(fn);
	const string nameBase    = cenv.penv.gensymstr(((fnName != "") ? fnName : "fn").c_str());
	const string implNameStr = string("_") + nameBase;
	cenv.setName(impl, implNameStr);

	cenv.liftStack.push(CEnv::FreeVars(fn, implNameStr));

	// Create a new stub environment frame for parameters
	cenv.push();
	const AType*          type      = cenv.type(fn);
	AType::const_iterator tp        = type->prot()->begin();
	AType*                implProtT = new AType(*type->prot()->as<const AType*>());
	ATuple::iterator      ip        = implProtT->begin();
	for (ATuple::const_iterator p = fn->prot()->begin(); p != fn->prot()->end(); ++p) {
		const AType* paramType = (*tp++)->as<const AType*>();
		if (paramType->kind == AType::EXPR && *paramType->head() == *cenv.tenv.Fn) {
			AType* fnType = new AType(*paramType);
			fnType->set_prot(new AType(const_cast<AType*>(cenv.tenv.var()), fnType->prot()->as<AType*>(), Cursor()));
			paramType = tup<const AType>((*p)->loc, cenv.tenv.Tup, fnType, NULL);
		}
		cenv.def((*p)->as<const ASymbol*>(), *p, paramType, NULL);
		*ip++ = new AType(*paramType);
	}

	/* Prepend closure parameter with dummy name (undefined symbol).
	 * The name of this parameter will be changed to the name of this
	 * function after lifting the body (so recursive references correctly
	 * refer to this function by the closure parameter).
	 */
	impl->set_prot(new ATuple(cenv.penv.sym("_"), impl->prot()));

	// Lift body
	const AType* implRetT = NULL;
	ATuple::iterator ci = impl->iter_at(2);
	for (ATuple::iterator i = fn->iter_at(2); i != fn->end(); ++i, ++ci) {
		*ci = resp_lift(cenv, code, *i);
		implRetT = cenv.type(*ci);
	}

	cenv.pop();

	// Set name of closure parameter to "me"
	*impl->prot()->begin() = cenv.penv.sym("_me");

	// Create definition for implementation fn
	ASymbol* implName = cenv.penv.sym(implNameStr);
	ATuple*   def      = tup<ATuple>(fn->loc, cenv.penv.sym("def"), implName, impl, NULL);
	code.push_back(def);

	AType* implT = new AType(*type); // Type of the implementation function
	TList  tupT(fn->loc, cenv.tenv.Tup, cenv.tenv.var(), NULL);
	TList  consT(fn->loc, cenv.tenv.Tup, implT, NULL);
	List<ATuple, AST> cons(fn->loc, cenv.penv.sym("Closure"), implName, NULL);

	implT->list_ref(1) = implProtT;

	const CEnv::FreeVars& freeVars = cenv.liftStack.top();
	for (CEnv::FreeVars::const_iterator i = freeVars.begin(); i != freeVars.end(); ++i) {
		cons.push_back(*i);
		tupT.push_back(const_cast<AType*>(cenv.type(*i)));
		consT.push_back(const_cast<AType*>(cenv.type(*i)));
	}
	cenv.liftStack.pop();

	implT->set_prot(new AType(tupT, implT->prot(), Cursor()));
	implT->list_ref(2) = const_cast<AType*>(implRetT);

	cenv.setType(impl, implT);
	cenv.setType(cons, consT);

	cenv.def(implName, impl, implT, NULL);
	if (cenv.name(fn) != "")
		cenv.def(cenv.penv.sym(cenv.name(fn)), fn, consT, NULL);

	return cons;
}

static AST*
lift_call(CEnv& cenv, Code& code, ATuple* call) throw()
{
	List<ATuple, AST> copy;

	// Lift all children (callee and arguments, recursively)
	for (ATuple::iterator i = call->begin(); i != call->end(); ++i)
		copy.push_back(resp_lift(cenv, code, *i));

	copy.head->loc = call->loc;

	const AType* copyT = NULL;
	
	ASymbol* sym = call->head()->to<ASymbol*>();
	if (sym && !cenv.liftStack.empty() && sym->cppstr == cenv.name(cenv.liftStack.top().fn)) {
		/* Recursive call to innermost function, call implementation directly,
		 * reusing the current "_me" closure parameter (no cons or .).
		 */
		copy.push_front(cenv.penv.sym(cenv.liftStack.top().implName));
	} else if (is_form(call, "fn")) {
		/* Special case: ((fn ...) ...)
		 * Lifting (fn ...) yields: (Fn _impl ...).
		 * We don't want ((Fn _impl ...) (Fn _impl ...) ...),
		 * so call the implementation function (_impl) directly and pass the
		 * closure as the first parameter:
		 * (_impl (Fn _impl ...) ...)
		 */
		ATuple*      closure = copy.head->list_ref(0)->as<ATuple*>();
		ASymbol*     implSym = closure->list_ref(1)->as<ASymbol*>();
		const AType* implT   = cenv.type(cenv.resolve(implSym));
		copy.push_front(implSym);
		copyT = implT->list_ref(2)->as<const AType*>();
	} else {
		// Call to a closure, prepend code to access implementation function
		ATuple* getFn = tup<ATuple>(call->loc, cenv.penv.sym("."),
		                        copy.head->head(),
		                        new ALiteral<int32_t>(0, Cursor()), NULL);
		const AType* calleeT = cenv.type(copy.head->head());
		assert(**calleeT->begin() == *cenv.tenv.Tup);
		const AType* implT = calleeT->list_ref(1)->as<const AType*>();
		copy.push_front(getFn);
		cenv.setType(getFn, implT);
		copyT = implT->list_ref(2)->as<const AType*>();
	}

	cenv.setType(copy, copyT);
	return copy;
}

static AST*
lift_def(CEnv& cenv, Code& code, ATuple* def) throw()
{
	// Define stub first for recursion
	const ASymbol* const sym  = def->list_ref(1)->as<const ASymbol*>();
	AST*           const body = def->list_ref(2);
	cenv.def(sym, body, cenv.type(body), NULL);
	if (is_form(body, "fn"))
		cenv.setName(body->as<const ATuple*>(), sym->str());

	assert(def->list_ref(1)->to<const ASymbol*>());
	List<ATuple, AST> copy;
	copy.push_back(def->head());
	copy.push_back(resp_lift(cenv, code, def->list_ref(1)));
	for (ATuple::iterator t = def->iter_at(2); t != def->end(); ++t)
		copy.push_back(resp_lift(cenv, code, *t));
	
	cenv.setTypeSameAs(copy, def);

	if (copy.head->list_ref(1) == copy.head->list_ref(2))
		return NULL; // Definition created by lift_fn when body was lifted

	cenv.def(copy.head->list_ref(1)->as<const ASymbol*>(),
	         copy.head->list_ref(2),
	         cenv.type(copy.head->list_ref(2)),
	         NULL);
	return copy;
}

static AST*
lift_builtin_call(CEnv& cenv, Code& code, ATuple* call) throw()
{
	List<ATuple, AST> copy;
	copy.push_back(call->head());

	// Lift all arguments
	for (ATuple::iterator i = call->iter_at(1); i != call->end(); ++i)
		copy.push_back(resp_lift(cenv, code, *i));

	cenv.setTypeSameAs(copy, call);
	return copy;
}

AST*
resp_lift(CEnv& cenv, Code& code, AST* ast) throw()
{
	ASymbol* const sym = ast->to<ASymbol*>();
	if (sym)
		return lift_symbol(cenv, code, sym);

	ATuple* const call = ast->to<ATuple*>();
	if (call) {
		const ASymbol* const sym  = call->head()->to<const ASymbol*>();
		const std::string    form = sym ? sym->cppstr : "";
		if (is_primitive(cenv.penv, call))
			return lift_builtin_call(cenv, code, call);
		else if (form == "fn")
			return lift_fn(cenv, code, call);
		else if (form == "def")
			return lift_def(cenv, code, call);
		else if (form == "if")
			return lift_builtin_call(cenv, code, call);
		else if (form == "cons" || isupper(form[0]))
			return lift_builtin_call(cenv, code, call);
		else if (form == ".")
			return lift_builtin_call(cenv, code, call);
		else if (form == "quote")
			return lift_builtin_call(cenv, code, call);
		else if (form == "match" || form == "def-type")
			return call; // FIXME
		else
			return lift_call(cenv, code, call);
	}

	return ast;
}