1 files changed, 57 insertions, 47 deletions

diff --git a/src/pprint.cpp b/src/pprint.cpp
index 7a463cc..003d447 100644
--- a/src/pprint.cpp
+++ b/src/pprint.cpp
@@ -32,7 +32,7 @@ newline(ostream& out, unsigned indent)
 		out << " ";
 }
 
-void
+ostream&
 print_tuple(ostream& out, const ATuple* tup, ATuple::const_iterator i,
             unsigned indent, bool newlines, CEnv* cenv, bool types, bool elem_types)
 {
@@ -53,6 +53,8 @@ print_tuple(ostream& out, const ATuple* tup, ATuple::const_iterator i,
 		
 		i = next;
 	}
+
+	return (out << ")");
 }
 	
 ostream&
@@ -97,62 +99,70 @@ print_to(ostream& out, const AST* ast, unsigned indent, CEnv* cenv, bool types)
 	if (tup) {
 		out << "(";
 		ATuple::const_iterator i = tup->begin();
-		const ASymbol* head = (i == tup->end()) ? NULL : (*i)->to<const ASymbol*>();
-		if (head) {
-			if (head == cenv->penv.sym("def")) {
-				out << (*i++) << " ";
-
-				// Print symbol (possibly with type annotation)
-				const AST* sym = *i++;
-				out << sym;
-				if (types)
-					out << " :" << cenv->tsubst.apply(cenv->tenv.var(sym));
-
-				// Print value on following lines, indented
-				newline(out, indent + 2);
-				print_tuple(out, tup, i, indent, true, cenv, types, false);
-				return out << ")";
+
+		std::string form = "";
+		if (i != tup->end()) {
+			const ASymbol* sym = (*i)->to<const ASymbol*>();
+			if (sym) {
+				form = sym->cppstr;
+			} else {
+				const ALexeme* lexeme = (*i)->to<const ALexeme*>();
+				if (lexeme)
+					form = *lexeme;
+			}
+		}
+		
+		if (form == "def") {
+			out << (*i++) << " ";
+
+			// Print symbol (possibly with type annotation)
+			const AST* sym = *i++;
+			out << sym;
+			if (types)
+				out << " :" << cenv->tsubst.apply(cenv->tenv.var(sym));
+
+			// Print value on following lines, indented
+			newline(out, indent + 2);
+			return print_tuple(out, tup, i, indent, true, cenv, types, false);
 				
-			} else if (head == cenv->penv.sym("fn")) {
-				out << (*i++) << " ";
-				const ATuple* pat = (*i++)->as<const ATuple*>();
+		} else if (form == "fn") {
+			out << (*i++) << " ";
+			const ATuple* pat = (*i++)->as<const ATuple*>();
 
-				// Print prototype (possibly with parameter type annotations)
-				out << "(";
-				print_tuple(out, pat, pat->begin(), indent, false, cenv, types, types);
-				out << ")";
+			// Print prototype (possibly with parameter type annotations)
+			out << "(";
+			print_tuple(out, pat, pat->begin(), indent, false, cenv, types, types);
 
-				// Print body expression(s) on following lines, indented
-				newline(out, indent + 2);
-				print_tuple(out, tup, i, indent + 2, true, cenv, types, false);
-				return out << ")";
+			// Print body expression(s) on following lines, indented
+			newline(out, indent + 2);
+			return print_tuple(out, tup, i, indent + 2, true, cenv, types, false);
 
-			} else if (head == cenv->penv.sym("if")) {
-				out << (*i++) << " ";
+		} else if (form == "if") {
+			out << (*i++) << " ";
 
-				// Print each condition and consequent pair separated by blank lines
-				for (; i != tup->end(); ) {
-					ATuple::const_iterator next = i;
-					++next;
+			// Print each condition and consequent pair separated by blank lines
+			for (; i != tup->end(); ) {
+				ATuple::const_iterator next = i;
+				++next;
 		
-					print_to(out, *i, indent + 2, cenv, types);
-					if (next != tup->end()) {
-						newline(out, indent + 2);
-						print_to(out, *next++, indent + 2, cenv, types);
-						newline(out, 0);
-						newline(out, indent + 2);
-					}
-
-					i = next;
+				print_to(out, *i, indent + 2, cenv, types);
+				if (next != tup->end()) {
+					newline(out, indent + 2);
+					print_to(out, *next++, indent + 2, cenv, types);
+					newline(out, 0);
+					newline(out, indent + 2);
 				}
+
+				i = next;
 			}
-		}
 
-		// Print plain tuple
-		print_tuple(out, tup, i, indent + 1, false, cenv, types, false);
-		return out << ")";
+			return out;
+			
+		} else {
+			return print_tuple(out, tup, i, indent + 1, false, cenv, types, false);
+		}
 	}
-
+	
 	return out << "?";
 }