1 files changed, 119 insertions, 86 deletions

diff --git a/src/cps.cpp b/src/cps.cpp
index 88ab425..c55f728 100644
--- a/src/cps.cpp
+++ b/src/cps.cpp
@@ -25,123 +25,156 @@
 
 #include "resp.hpp"
 
-/** (cps x cont) => (cont x) */
-static const AST*
-cps_value(TEnv& tenv, AST* cont) const
+static bool
+is_value(CEnv& cenv, const AST* exp)
 {
-	return tup(loc, cont, this, 0);
+	const ATuple* const call = exp->to_tuple();
+	if (!call)
+		return true; // Atom
+
+	if (!is_primitive(cenv.penv, exp))
+		return false; // Non-primitive fn call
+	
+	for (ATuple::const_iterator i = call->iter_at(1); i != call->end(); ++i)
+		if (!is_value(cenv, *i))
+			return false; // Primitive with non-value argument
+
+	return true; // Primitive with all value arguments
 }
 
-/** (cps (fn (a ...) body) cont) => (cont (fn (a ... k) (cps body k)) */
+/** [v]k => (k v) */
 static const AST*
-cps_fn(TEnv& tenv, AST* cont) const
+cps_value(CEnv& cenv, const AST* v, const AST* k)
 {
-	ATuple*  copyProt = new ATuple(*prot());
-	ASymbol* contArg  = tenv.penv.gensym("_k");
-	copyProt->push_back(contArg);
-	AFn* copy = tup(loc, tenv.penv.sym("fn"), copyProt, 0);
-	const_iterator p = begin();
-	++(++p);
-	for (; p != end(); ++p)
-		copy->push_back((*p)->(tenv, contArg));
-	return tup(loc, cont, copy, 0);
+	return tup(v->loc, k, v, 0);
 }
 
+/** [(fn (a ...) r)]k => (k (fn (a ... k2) [r]k2)) */
 static const AST*
-cps_primitive(TEnv& tenv, AST* cont) const
+cps_fn(CEnv& cenv, const ATuple* fn, const AST* cont)
 {
-	return value() ? tup(loc, cont, this, 0) : ATuple::(tenv, cont);
-}
+	const ASymbol* k2 = cenv.penv.gensym("__k");
 
-/** (cps (f a b ...)) => (a (fn (x) (b (fn (y) ... (cont (f x y ...)) */
-static const AST*
-cps_tuple(TEnv& tenv, AST* cont) const
-{
-	std::vector< std::pair<AFn*, AST*> > funcs;
-	AFn*     fn  = NULL;
-	ASymbol* arg = NULL;
-
-	// Make a continuation for each element (operator and arguments)
-	// Argument evaluation continuations are not themselves in CPS.
-	// Each makes a tail call to the next, and the last makes a tail
-	// call to the continuation of this call
-	const_iterator firstFnIter = end();
-	AFn*           firstFn     = NULL;
-	ssize_t        index       = 0;
-	FOREACHP(const_iterator, i, this) {
-		if (!(*i)->to_tuple()) {
-			funcs.push_back(make_pair((AFn*)NULL, (*i)));
-		} else {
-			arg = tenv.penv.gensym("a");
+	List copyProt;
+	FOREACHP(ATuple::const_iterator, i, fn->prot())
+		copyProt.push_back(*i);
+	copyProt.push_back(k2);
 
-			AFn* thisFn = tup(loc, tenv.penv.sym("fn"),
-					tup((*i)->loc, arg, 0),
-					0);
+	assert(fn->fst());
+	assert(copyProt.head);
+	List copy;
+	copy.push_back(cenv.penv.sym("fn"));
+	copy.push_back(copyProt);
 
-			if (firstFnIter == end()) {
-				firstFnIter = i;
-				firstFn = thisFn;
-			}
+	for (ATuple::const_iterator i = fn->iter_at(2); i != fn->end(); ++i)
+		copy.push_back(resp_cps(cenv, *i, k2));
 
-			if (fn)
-				fn->push_back((*i)->(tenv, thisFn));
+	return copy;
+}
 
-			funcs.push_back(make_pair(thisFn, arg));
-			fn = thisFn;
+/** [(f a b ...)]k => [a](fn (__a) [b](fn (__b) ... (f __a __b ... k))) */
+static const AST*
+cps_call(CEnv& cenv, const ATuple* call, const AST* k)
+{
+	// Build innermost application first
+	List body;
+	typedef std::vector<const AST*> ExpVec;
+	ExpVec exprs;
+	ExpVec args;
+	FOREACHP(ATuple::const_iterator, i, call) {
+		exprs.push_back(*i);
+		if (is_value(cenv, *i)) {
+			body.push_back(*i);
+			args.push_back(*i);
+		} else {
+			const ASymbol* sym = cenv.penv.gensym("__a");
+			body.push_back(sym);
+			args.push_back(sym);
 		}
-		++index;
 	}
 
-	if (firstFnIter != end()) {
-		// Call this call's callee in the last argument evaluator
-		ATuple* call = tup(loc, 0);
-		assert(funcs.size() == size());
-		for (size_t i = 0; i < funcs.size(); ++i)
-			call->push_back(funcs[i].second);
-
-		assert(fn);
-		fn->push_back(call->(tenv, cont));
-		return (*firstFnIter)->(tenv, firstFn);
+	const AST* cont;
+	if (cenv.penv.primitives.find(call->fst()->str()) != cenv.penv.primitives.end()) {
+		cont = tup(Cursor(), k, body.head, 0);
 	} else {
-		assert(fst()->value());
-		ATuple* ret = tup(loc, 0);
-		FOREACHP(const_iterator, i, this)
-			ret->push_back((*i));
-		if (!is_primitive(this))
-			ret->push_back(cont);
-		return ret;
+		body.push_back(k);
+		cont = body;
 	}
+
+	// Wrap application in fns to evaluate parameters (from right to left)
+	std::vector<const AST*>::const_reverse_iterator a = args.rbegin();
+	for (ExpVec::const_reverse_iterator e = exprs.rbegin(); e != exprs.rend(); ++e, ++a) {
+		if (!is_value(cenv, *e)) {
+			cont = resp_cps(cenv, *e, tup(Cursor(), cenv.penv.sym("fn"),
+			                              tup(Cursor(), *a, 0),
+			                              cont,
+			                              0));
+		}
+	}
+
+	return cont;
 }
 
-/** (cps (def x y)) => (y (fn (x) (cont))) */
+/** [(def x y)]k => (def x [y]k) */
 static const AST*
-cps_def(TEnv& tenv, AST* cont) const
+cps_def(CEnv& cenv, const ATuple* def, const AST* k)
 {
+	List copy(def->loc, def->fst(), def->frst(), 0);
+	copy.push_back(resp_cps(cenv, def->list_ref(2), k));
+	return copy;
+	/*
 	AST*    val     = body()->(tenv, cont);
 	ATuple* valCall = val->to_tuple();
 	ATuple::iterator i = valCall->begin();
 	return tup(loc, tenv.penv.sym("def"), sym(), *++i, 0);
+	*/
 }
 
-/** (cps (if c t ... e)) => */
+/** [(if c t e)]k => [c](fn (__c) (if c [t]k [e]k)) */
 static const AST*
-cps_iff(TEnv& tenv, AST* cont) const
+cps_if(CEnv& cenv, const ATuple* aif, const AST* k)
 {
-	ASymbol* argSym = tenv.penv.gensym("c");
-	const_iterator i = begin();
-	AST* cond = *++i;
-	AST* exp  = *++i;
-	AST* next = *++i;
-	if (cond->value()) {
-		return tup(loc, tenv.penv.sym("if"), cond,
-			exp->(tenv, cont),
-			next->(tenv, cont), 0);
+	ATuple::const_iterator i = aif->begin();
+	const AST* const c = *++i;
+	const AST* const t = *++i;
+	const AST* const e = *++i;
+	if (is_value(cenv, c)) {
+		return tup(aif->loc, cenv.penv.sym("if"), c,
+		           resp_cps(cenv, t, k),
+		           resp_cps(cenv, e, k), 0);
 	} else {
-		AFn* contFn = tup(loc, tenv.penv.sym("fn"),
-				tup(cond->loc, argSym, tenv.penv.gensym("_k"), 0),
-				tup(loc, tenv.penv.sym("if"), argSym,
-					exp->(tenv, cont),
-					next->(tenv, cont), 0));
-		return cond->(tenv, contFn);
+		/*
+		  const ASymbol* const condSym = cenv.penv.gensym("c");
+		  const ATuple* contFn = tup(loc, tenv.penv.sym("fn"),
+		  tup(cond->loc, argSym, tenv.penv.gensym("_k"), 0),
+		  tup(loc, tenv.penv.sym("if"), argSym,
+		  exp->(tenv, cont),
+		  next->(tenv, cont), 0));
+		  return cond->(tenv, contFn);
+		*/
+		return aif;
 	}
 }
+
+const AST*
+resp_cps(CEnv& cenv, const AST* ast, const AST* k) throw()
+{
+	if (is_value(cenv, ast))
+		return cps_value(cenv, ast, k);
+
+	const ATuple* const call = ast->to_tuple();
+	if (call) {
+		const ASymbol* const sym  = call->fst()->to_symbol();
+		const std::string    form = sym ? sym->sym() : "";
+		if (form == "def")
+			return cps_def(cenv, call, k);
+		else if (form == "fn")
+			return cps_fn(cenv, call, k);
+		else if (form == "if")
+			return cps_if(cenv, call, k);
+		else
+			return cps_call(cenv, call, k);
+	}
+
+	return cps_value(cenv, ast, k);
+}