1 files changed, 159 insertions, 29 deletions

diff --git a/src/lift.cpp b/src/lift.cpp
index 94f70e0..6a53165 100644
--- a/src/lift.cpp
+++ b/src/lift.cpp
@@ -17,64 +17,194 @@
 
 /** @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;
 
-void
-AFn::lift(CEnv& cenv) throw()
+AST*
+ASymbol::lift(CEnv& cenv, Code& code) throw()
 {
+	if (!cenv.liftStack.empty() && cppstr == cenv.liftStack.top().fn->name) {
+		return cenv.penv.sym("me"); // Reference to innermost function
+	} else if (!cenv.penv.handler(true, cppstr)
+			&& !cenv.penv.handler(false, cppstr)
+			&& !cenv.code.innermost(this)) {
+
+		const int32_t index = cenv.liftStack.top().index(this);
+
+		// Replace symbol with code to access free variable from closure
+		return tup<ADot>(loc, cenv.penv.sym("."),
+				cenv.penv.sym("me"),
+				new ALiteral<int32_t>(index, Cursor()),
+				NULL);
+	} else {
+		return this;
+	}
+}
+
+AST*
+ATuple::lift(CEnv& cenv, Code& code) throw()
+{
+	ATuple*  ret = new ATuple(*this);
+	iterator ri  = ret->begin();
+	FOREACHP(const_iterator, t, this)
+		*ri++ = (*t)->lift(cenv, code);
+	cenv.setTypeSameAs(ret, this);
+	return ret;
+}
+
+AST*
+AFn::lift(CEnv& cenv, Code& code) throw()
+{
+	AFn* impl = new AFn(this);
+	const string nameBase = cenv.penv.gensymstr(((name != "") ? name : "fn").c_str());
+	impl->name = "_" + nameBase;
+
+	cenv.liftStack.push(CEnv::FreeVars(this, impl->name));
+
 	// Create a new stub environment frame for parameters
 	cenv.push();
+	AType::const_iterator tp = cenv.type(this)->prot()->begin();
 	for (const_iterator p = prot()->begin(); p != prot()->end(); ++p)
-		cenv.def((*p)->as<ASymbol*>(), *p, NULL, NULL);
+		cenv.def((*p)->as<ASymbol*>(), *p, (*tp++)->as<AType*>(), NULL);
+
+	/* Add 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->prot()->push_front(cenv.penv.sym("_"));
+
 
 	// Lift body
-	for (iterator i = begin() + 2; i != end(); ++i)
-		(*i)->lift(cenv);
+	const AType* implRetT = NULL;
+	iterator ci = impl->begin() + 2;
+	for (const_iterator i = begin() + 2; i != end(); ++i, ++ci) {
+		*ci = (*i)->lift(cenv, code);
+		implRetT = cenv.type(*ci);
+	}
 
 	cenv.pop();
 
-	AType* type = cenv.type(this);
-	if (impls.find(type) || !type->concrete())
-		return;
+	// Set name of closure parameter to actual name of this function
+	*impl->prot()->begin() = cenv.penv.sym("me");
 
-	AType* protT = type->prot()->as<AType*>();
-	cenv.engine()->compileFunction(cenv, this, *protT);
-}
+	// Create definition for implementation fn
+	ASymbol* implName = cenv.penv.sym(impl->name);
+	ADef*    def       = tup<ADef>(loc, cenv.penv.sym("def"), implName, impl, NULL);
+	code.push_back(def);
 
-void
-ACall::lift(CEnv& cenv) throw()
-{
-	AFn* c = cenv.resolve(head())->to<AFn*>();
-	AType argsT(loc);
+	AType* implT = new AType(*cenv.type(this)); // Type of the implementation function
+	AType* tupT  = tup<AType>(loc, cenv.tenv.Tup, cenv.tenv.var(), NULL);
+	AType* consT = tup<AType>(loc, cenv.tenv.Tup, implT, NULL);
+	ACons* cons  = tup<ACons>(loc, cenv.penv.sym("cons"), implName, NULL); // Closure
 
-	// Lift arguments and build arguments type
-	for (iterator i = begin() + 1; i != end(); ++i) {
-		(*i)->lift(cenv);
-		argsT.push_back(cenv.type(*i));
+	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();
 
-	// Lift callee (if it's not a primitive)
-	if (c)
-		cenv.engine()->compileFunction(cenv, c, argsT);
+	implT->prot()->push_front(tupT);
+	*(implT->begin() + 2) = const_cast<AType*>(implRetT);
+
+	cenv.setType(impl, implT);
+	cenv.setType(cons, consT);
+
+	cenv.def(implName, impl, implT, NULL);
+	if (name != "")
+		cenv.def(cenv.penv.sym(name), this, consT, NULL);
+
+	return cons;
 }
 
-void
-ADot::lift(CEnv& cenv) throw()
+AST*
+ACall::lift(CEnv& cenv, Code& code) throw()
 {
-	(*(begin() + 1))->lift(cenv);
+	ACall*           copy = new ACall(this);
+	ATuple::iterator ri   = copy->begin();
+
+	// Lift all children (callee and arguments, recursively)
+	for (const_iterator i = begin(); i != end(); ++i)
+		*ri++ = (*i)->lift(cenv, code);
+
+	ASymbol* sym = head()->to<ASymbol*>();
+	if (sym && !cenv.liftStack.empty() && sym->cppstr == cenv.liftStack.top().fn->name) {
+		/* 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 (head()->to<AFn*>()) {
+		/* Special case: ((fn ...) ...)
+		 * Lifting (fn ...) yields: (cons _impl ...).
+		 * We don't want ((cons _impl ...) (cons _impl ...) ...),
+		 * so call the implementation function (_impl) directly:
+		 * (_impl (cons _impl ...) ...)
+		 */
+		ACons*       closure = (*copy->begin())->as<ACons*>();
+		ASymbol*     implSym = (*(closure->begin() + 1))->as<ASymbol*>();
+		const AType* implT   = cenv.type(cenv.resolve(implSym));
+		copy->push_front(implSym);
+		cenv.setType(copy, (*(implT->begin() + 2))->as<const AType*>());
+	} else {
+		// Call to a closure, prepend code to access implementation function
+		ADot* getFn = tup<ADot>(loc, cenv.penv.sym("."),
+				copy->head(),
+				new ALiteral<int32_t>(0, Cursor()), NULL);
+		const AType* calleeT = cenv.type(copy->head());
+		assert(**calleeT->begin() == *cenv.tenv.Tup);
+		const AType* implT = (*(calleeT->begin() + 1))->as<const AType*>();
+		copy->push_front(getFn);
+		cenv.setType(getFn, implT);
+		cenv.setType(copy, (*(implT->begin() + 2))->as<const AType*>());
+	}
+
+	return copy;
 }
 
-void
-ADef::lift(CEnv& cenv) throw()
+AST*
+ADef::lift(CEnv& cenv, Code& code) throw()
 {
 	// Define stub first for recursion
 	cenv.def(sym(), body(), cenv.type(body()), NULL);
 	AFn* c = body()->to<AFn*>();
 	if (c)
 		c->name = sym()->str();
-	body()->lift(cenv);
+
+	ADef* copy = new ADef(ATuple::lift(cenv, code)->as<ATuple*>());
+
+	if (copy->sym() == copy->body())
+		return NULL; // Definition created by AFn::lift when body was lifted
+
+	cenv.def(copy->sym(), copy->body(), cenv.type(copy->body()), NULL);
+	cenv.setTypeSameAs(copy, this);
+	return copy;
+}
+
+template<typename T>
+AST*
+lift_builtin_call(CEnv& cenv, T* call, Code& code) throw()
+{
+	ATuple*          copy = new T(call);
+	ATuple::iterator ri   = copy->begin() + 1;
+
+	// Lift all arguments
+	for (typename T::const_iterator i = call->begin() + 1; i != call->end(); ++i)
+		*ri++ = (*i)->lift(cenv, code);
+
+	cenv.setTypeSameAs(copy, call);
+	return copy;
 }
+
+AST* AIf::lift(CEnv& cenv, Code& code)        throw() { return lift_builtin_call(cenv, this, code); }
+AST* ACons::lift(CEnv& cenv, Code& code)      throw() { return lift_builtin_call(cenv, this, code); }
+AST* ADot::lift(CEnv& cenv, Code& code)       throw() { return lift_builtin_call(cenv, this, code); }
+AST* APrimitive::lift(CEnv& cenv, Code& code) throw() { return lift_builtin_call(cenv, this, code); }