1 files changed, 93 insertions, 105 deletions

diff --git a/src/pprint.cpp b/src/pprint.cpp
index 28df1cc..7a463cc 100644
--- a/src/pprint.cpp
+++ b/src/pprint.cpp
@@ -22,7 +22,41 @@
 #include "resp.hpp"
 
 ostream&
-operator<<(ostream& out, const AST* ast)
+print_to(ostream& out, const AST* ast, unsigned indent, CEnv* cenv, bool types);
+
+static void
+newline(ostream& out, unsigned indent)
+{
+	out << endl;
+	for (unsigned i = 0; i < indent; ++i)
+		out << " ";
+}
+
+void
+print_tuple(ostream& out, const ATuple* tup, ATuple::const_iterator i,
+            unsigned indent, bool newlines, CEnv* cenv, bool types, bool elem_types)
+{
+	for (; i != tup->end(); ) {
+		ATuple::const_iterator next = i;
+		++next;
+		
+		print_to(out, *i, newlines ? indent + 2 : indent, cenv, types);
+		if (elem_types)
+			out << " :" << cenv->tsubst.apply(cenv->tenv.var(*i));
+
+		if (next != tup->end()) {
+			if (newlines)
+				newline(out, indent + 2);
+			else
+				out << " ";
+		}
+		
+		i = next;
+	}
+}
+	
+ostream&
+print_to(ostream& out, const AST* ast, unsigned indent, CEnv* cenv, bool types)
 {
 	const ALexeme* lexeme = ast->to<const ALexeme*>();
 	if (lexeme)
@@ -62,122 +96,76 @@ operator<<(ostream& out, const AST* ast)
 	const ATuple* tup = ast->to<const ATuple*>();
 	if (tup) {
 		out << "(";
-		for (ATuple::const_iterator i = tup->begin(); i != tup->end(); ) {
-			ATuple::const_iterator next = i;
-			++next;
-			out << (*i) << ((next != tup->end()) ? " " : "");
-			i = next;
+		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 << ")";
+				
+			} else if (head == cenv->penv.sym("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 body expression(s) on following lines, indented
+				newline(out, indent + 2);
+				print_tuple(out, tup, i, indent + 2, true, cenv, types, false);
+				return out << ")";
+
+			} else if (head == cenv->penv.sym("if")) {
+				out << (*i++) << " ";
+
+				// 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 plain tuple
+		print_tuple(out, tup, i, indent + 1, false, cenv, types, false);
 		return out << ")";
 	}
 
 	return out << "?";
 }
 
-enum TupleStyle {
-	STYLE_ONE_LINE, // Entire tuple printed on one line
-	STYLE_ALIGN, // Align all elements aligned after opening '('
-	STYLE_CALL // Call with all arguments aligned after head
-};
-
-static bool
-can_one_line(const AST* ast, const unsigned max_width, unsigned& width, unsigned& depth)
-{
-	static const unsigned MAX_DEPTH = 3;
-
-	ostringstream ss;
-	ss << ast;
-	if ((width += ss.str().length()) >= max_width)
-		return false;
-
-	const ATuple* tup = ast->to<const ATuple*>();
-	if (tup) {
-		if (++depth >= MAX_DEPTH)
-			return false;
-
-		for (ATuple::const_iterator i = tup->begin(); i != tup->end(); ++i) {
-			ostringstream ss;
-			ss << *i;
-			width += ss.str().length();
-			if (!can_one_line(*i, max_width, width, depth))
-				return false;
-		}
-	}
-
-	return true;
-}
-	
-static TupleStyle
-tuple_style(const ATuple* tup, const unsigned max_width, unsigned& indent)
-{
-	unsigned depth = 0;
-	unsigned width = indent;
-	if (can_one_line(tup, max_width, width, depth))
-		return STYLE_ONE_LINE;
-
-	if (tup->to<const ADef*>() || tup->to<const AFn*>()) {
-		indent += 2;
-		return STYLE_CALL;
-	}
-	
-	if (tup->to<const ACall*>()) {
-		indent += tup->head()->str().length() + 1;
-		return STYLE_CALL;
-	}
-
-	return STYLE_ALIGN;
-}
-
-static void
-newline(ostream& out, unsigned indent)
-{
-	out << endl;
-	for (unsigned i = 0; i < indent; ++i)
-		out << " ";
-}
-	
-void
-pprint_internal(ostream& out, const AST* ast, unsigned indent)
+ostream&
+operator<<(ostream& out, const AST* ast)
 {
-	const ATuple* tup = ast->to<const ATuple*>();
-	if (tup) {
-		out << "(";
-		unsigned               child_indent = indent + 1;
-		ATuple::const_iterator i            = tup->begin();
-		if (i != tup->end()) {
-			const TupleStyle style = tuple_style(tup, 80, child_indent);
-			pprint_internal(out, *i++, indent + 1); // Print head
-			switch (style) {
-			case STYLE_ONE_LINE:
-				while (i != tup->end())
-					out << " " << *i++;
-				break;
-			case STYLE_ALIGN:
-				while (i != tup->end()) {
-					newline(out, indent + 1);
-					pprint_internal(out, *i++, child_indent);
-				}
-				break;
-			case STYLE_CALL:
-				out << " ";
-				pprint_internal(out, *i, child_indent); // Print first argument
-				while (++i != tup->end()) {
-					newline(out, child_indent);
-					pprint_internal(out, *i, child_indent);
-				}
-				break;
-			}
-		}
-		out << ")";
-	} else {
-		out << ast;
-	}
+	return print_to(out, ast, 0, NULL, false);
 }
 
 /// Pretty-print @a ast to @a out
 void
-pprint(ostream& out, const AST* ast)
+pprint(ostream& out, const AST* ast, CEnv* cenv, bool types)
 {
-	pprint_internal(out, ast, 0);
+	print_to(out, ast, 0, cenv, types);
 	out << endl;
 }