1 files changed, 83 insertions, 20 deletions

diff --git a/cps.cpp b/cps.cpp
index c1bddf2..e1a419f 100644
--- a/cps.cpp
+++ b/cps.cpp
@@ -22,55 +22,118 @@
  * CPS Conversion                                                          *
  ***************************************************************************/
 
-/** (cps x) => (cont x) */
+/** (cps x cont) => (cont x) */
 AST*
 AST::cps(TEnv& tenv, AST* cont)
 {
 	return tup<ACall>(loc, cont, this, 0);
 }
 
-AST*
-ATuple::cps(TEnv& tenv, AST* cont)
-{
-	ATuple* copy = tup<ATuple>(loc, NULL);
-	FOREACH(const_iterator, p, *this)
-		copy->push_back((*p)->cps(tenv, cont));
-	return copy;
-}
-
-/** (cps (fn (a ...) body ...)) => */
+/** (cps (fn (a ...) body) cont) => (cont (fn (a ... k) (cps body k))*/
 AST*
 AFn::cps(TEnv& tenv, AST* cont)
 {
-	AFn* copy = tup<AFn>(loc, tenv.penv.sym("fn"), prot(), 0);
+	ATuple*  copyProt = new ATuple(prot()->loc, *prot());
+	ASymbol* contArg  = tenv.penv.gensym("_k");
+	copyProt->push_back(contArg);
+	AFn* copy = tup<AFn>(loc, tenv.penv.sym("fn"), copyProt, 0);
 	const_iterator p = begin();
 	++(++p);
 	for (; p != end(); ++p)
-		copy->push_back((*p)->cps(tenv, cont));
-	return copy;
+		copy->push_back((*p)->cps(tenv, contArg));
+	return tup<ACall>(loc, cont, copy, 0);
+}
+
+AST*
+APrimitive::cps(TEnv& tenv, AST* cont)
+{
+	return value() ? tup<ACall>(loc, cont, this, 0) : ACall::cps(tenv, cont);
 }
 
 /** (cps (f a b ...)) => (a (fn (x) (b (fn (y) ... (cont (f x y ...)) */
 AST*
 ACall::cps(TEnv& tenv, AST* cont)
 {
-	return tup<ACall>(loc, cont, this, 0);
+	std::vector< std::pair<AFn*, AST*> > funcs;
+	AFn*     fn  = NULL;
+	ASymbol* arg = NULL;
+
+	// Make a continuation for each element (operator and arguments)
+	ssize_t firstFn = -1;
+	ssize_t lastFn = -1;
+	for (size_t i = 0; i < size(); ++i) {
+		if (!at(i)->to<ATuple*>()) {
+			funcs.push_back(make_pair((AFn*)NULL, at(i)));
+		} else {
+			arg = tenv.penv.gensym("a");
+
+			if (firstFn == -1)
+				firstFn = i;
+
+			AFn* thisFn = tup<AFn>(loc, tenv.penv.sym("fn"),
+					tup<ATuple>(at(i)->loc, arg, tenv.penv.gensym("_k"), 0),
+					0);
+
+			if (lastFn != -1)
+				fn->push_back(at(lastFn)->cps(tenv, thisFn));
+
+			funcs.push_back(make_pair(thisFn, arg));
+			fn = thisFn;
+			lastFn = i;
+		}
+	}
+
+	if (firstFn != -1) {
+		// Call our callee in the last argument's evaluation function
+		ACall* call = tup<ACall>(loc, 0);
+		assert(funcs.size() == size());
+		for (size_t i = 0; i < funcs.size(); ++i)
+			call->push_back(funcs[i].second);
+		if (!to<APrimitive*>())
+			call->push_back(cont);
+		else
+			call = tup<ACall>(loc, cont, call, 0);
+
+		assert(fn);
+		fn->push_back(call);
+		return at(firstFn)->cps(tenv, funcs[firstFn].first);
+	} else {
+		assert(at(0)->value());
+		ACall* ret = tup<ACall>(loc, 0);
+		for (size_t i = 0; i < size(); ++i)
+			ret->push_back(at(i));
+		if (!to<APrimitive*>())
+			ret->push_back(cont);
+		return ret;
+	}
 }
 
 /** (cps (def x y)) => (y (fn (x) (cont))) */
 AST*
 ADef::cps(TEnv& tenv, AST* cont)
 {
-	return tup<ADef>(loc, tenv.penv.sym("def"), sym(), at(2)->cps(tenv, cont), 0);
+	AST*   val     = at(2)->cps(tenv, cont);
+	ACall* valCall = val->to<ACall*>();
+	assert(valCall);
+	return tup<ADef>(loc, tenv.penv.sym("def"), sym(), valCall->at(1), 0);
 }
 
 /** (cps (if c t ... e)) => */
 AST*
 AIf::cps(TEnv& tenv, AST* cont)
 {
-	AFn* contFn = tup<AFn>(loc, tenv.penv.sym("if-fn"),
-			new ATuple(at(1)->loc, cont, 0), 0);
-	ACall* condCall = tup<ACall>(loc, contFn, 0);
-	return condCall;
+	ASymbol* argSym = tenv.penv.gensym("c");
+	if (at(1)->value()) {
+		return tup<AIf>(loc, tenv.penv.sym("if"), at(1),
+			at(2)->cps(tenv, cont),
+			at(3)->cps(tenv, cont), 0);
+	} else {
+		AFn* contFn = tup<AFn>(loc, tenv.penv.sym("fn"),
+				tup<ATuple>(at(1)->loc, argSym, tenv.penv.gensym("_k"), 0),
+				tup<AIf>(loc, tenv.penv.sym("if"), argSym,
+					at(2)->cps(tenv, cont),
+					at(3)->cps(tenv, cont), 0));
+		return at(1)->cps(tenv, contFn);
+	}
 }