1 files changed, 27 insertions, 20 deletions

diff --git a/src/cps.cpp b/src/cps.cpp
index eb44fd9..6711556 100644
--- a/src/cps.cpp
+++ b/src/cps.cpp
@@ -62,21 +62,24 @@ ACall::cps(TEnv& tenv, AST* cont)
 	// 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
-	ssize_t firstFn = -1;
-	ssize_t index   = 0;
-	FOREACH(iterator, i, *this) {
+	iterator firstFnIter = end();
+	AFn*     firstFn     = NULL;
+	ssize_t  index       = 0;
+	FOREACHP(iterator, i, this) {
 		if (!(*i)->to<ATuple*>()) {
 			funcs.push_back(make_pair((AFn*)NULL, (*i)));
 		} else {
 			arg = tenv.penv.gensym("a");
 
-			if (firstFn == -1)
-				firstFn = index;
-
 			AFn* thisFn = tup<AFn>(loc, tenv.penv.sym("fn"),
 					tup<ATuple>((*i)->loc, arg, 0),
 					0);
 
+			if (firstFnIter == end()) {
+				firstFnIter = i;
+				firstFn = thisFn;
+			}
+
 			if (fn)
 				fn->push_back((*i)->cps(tenv, thisFn));
 
@@ -86,7 +89,7 @@ ACall::cps(TEnv& tenv, AST* cont)
 		++index;
 	}
 
-	if (firstFn != -1) {
+	if (firstFnIter != end()) {
 		// Call this call's callee in the last argument evaluator
 		ACall* call = tup<ACall>(loc, 0);
 		assert(funcs.size() == size());
@@ -95,11 +98,11 @@ ACall::cps(TEnv& tenv, AST* cont)
 
 		assert(fn);
 		fn->push_back(call->cps(tenv, cont));
-		return at(firstFn)->cps(tenv, funcs[firstFn].first);
+		return (*firstFnIter)->cps(tenv, firstFn);
 	} else {
 		assert(head()->value());
 		ACall* ret = tup<ACall>(loc, 0);
-		FOREACH(iterator, i, *this)
+		FOREACHP(iterator, i, this)
 			ret->push_back((*i));
 		if (!to<APrimitive*>())
 			ret->push_back(cont);
@@ -111,10 +114,10 @@ ACall::cps(TEnv& tenv, AST* cont)
 AST*
 ADef::cps(TEnv& tenv, AST* cont)
 {
-	AST*   val     = at(2)->cps(tenv, cont);
+	AST*   val     = body()->cps(tenv, cont);
 	ACall* valCall = val->to<ACall*>();
-	assert(valCall);
-	return tup<ADef>(loc, tenv.penv.sym("def"), sym(), valCall->at(1), 0);
+	ACall::iterator i = valCall->begin();
+	return tup<ADef>(loc, tenv.penv.sym("def"), sym(), *++i, 0);
 }
 
 /** (cps (if c t ... e)) => */
@@ -122,16 +125,20 @@ AST*
 AIf::cps(TEnv& tenv, AST* cont)
 {
 	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);
+	const_iterator i = begin();
+	AST* cond = *++i;
+	AST* exp  = *++i;
+	AST* next = *++i;
+	if (cond->value()) {
+		return tup<AIf>(loc, tenv.penv.sym("if"), cond,
+			exp->cps(tenv, cont),
+			next->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<ATuple>(cond->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);
+					exp->cps(tenv, cont),
+					next->cps(tenv, cont), 0));
+		return cond->cps(tenv, contFn);
 	}
 }