1 files changed, 31 insertions, 18 deletions

diff --git a/src/repl.cpp b/src/repl.cpp
index 83324f2..25ae163 100644
--- a/src/repl.cpp
+++ b/src/repl.cpp
@@ -122,6 +122,7 @@ eval(CEnv& cenv, Cursor& cursor, istream& is, bool execute)
 		// Flatten expressions
 		Code flattened;  // Type and function definitions
 		Code exprs;  // Other top-level expressions (main code)
+		const AST* retT = NULL;
 		for (Code::const_iterator i = lifted.begin(); i != lifted.end(); ++i) {
 			const ATuple* call = (*i)->to_tuple();
 			if (call && (is_form(*i, "def-type")
@@ -131,31 +132,35 @@ eval(CEnv& cenv, Cursor& cursor, istream& is, bool execute)
 				const ATuple* tup = (*i)->to_tuple();
 				if (!tup || !tup->empty()) {
 					exprs.push_back(*i);
+					retT = cenv.type(*i);
 				}
 			}
 		}
-		if (cenv.args.find("-F") != cenv.args.end())
-			return dump(cenv, flattened);
-
-		// Compile type and function definitions
-		for (Code::const_iterator i = flattened.begin(); i != flattened.end(); ++i) {
-			resp_compile(cenv, *i);
-		}
 
-		// Compile other top-level expressions into "main" function
+		// Flatten main code into flattened
 		if (!exprs.empty()) {
-			const AST*    type = cenv.type(exprs.back());
-			const ATuple* fnT  = tup(cursor, cenv.tenv.Fn, new ATuple(cursor), type, 0);
-
-			CFunc f   = cenv.engine()->startFn(cenv, "main", new ATuple(cursor), fnT);
-			CVal  val = NULL;
+			const ASymbol* main = cenv.penv.sym("main");
+			List mainT(Cursor(), cenv.penv.sym("Fn"), new ATuple(Cursor()), retT, NULL);
+			cenv.def(main, NULL, mainT, NULL);
+			
+			flattened.push_back(
+				tup(Cursor(), cenv.penv.sym("fn-start"), main, NULL));
+
+			const AST* mainRet = NULL;
 			for (Code::const_iterator i = exprs.begin(); i != exprs.end(); ++i)
-				val = resp_compile(cenv, *i);
-			cenv.engine()->finishFn(cenv, val, type);
+				mainRet = resp_flatten(cenv, flattened, *i);
+			
+			flattened.push_back(
+				tup(Cursor(), cenv.penv.sym("fn-end"), main, mainRet, NULL));
+		}
 
-			// Call main and print result
-			if (cenv.args.find("-S") == cenv.args.end())
-				callPrintCollect(cenv, f, ast, type, execute);
+		if (cenv.args.find("-F") != cenv.args.end()) {
+			return dump(cenv, flattened);
+		}
+
+		// Compile flattened code
+		for (Code::const_iterator i = flattened.begin(); i != flattened.end(); ++i) {
+			resp_compile(cenv, *i);
 		}
 
 		if (cenv.args.find("-S") != cenv.args.end()) {
@@ -163,6 +168,13 @@ eval(CEnv& cenv, Cursor& cursor, istream& is, bool execute)
 			return 0;
 		}
 
+		// Call main and print result
+		if (!exprs.empty()) {
+			CFunc f = cenv.engine()->getFn(cenv, "main");
+			if (cenv.args.find("-S") == cenv.args.end())
+				callPrintCollect(cenv, f, ast, retT, execute);
+		}
+
 	} catch (Error& e) {
 		cenv.err << e.what() << endl;
 		return 1;
@@ -190,6 +202,7 @@ repl(CEnv& cenv)
 
 			Code lifted;
 			ast = resp_lift(cenv, lifted, ast);
+
 			const AST*    type = cenv.type(ast);
 			const ATuple* fnT  = tup(cursor, cenv.tenv.Fn, new ATuple(cursor), type, 0);
 			CFunc f = NULL;