1 files changed, 53 insertions, 45 deletions

diff --git a/src/repl.cpp b/src/repl.cpp
index edaa003..b457e06 100644
--- a/src/repl.cpp
+++ b/src/repl.cpp
@@ -81,7 +81,7 @@ dump(CEnv& cenv, const Code& code)
 }
 
 const AST*
-compile(CEnv& cenv, const Code& parsed, Code& defs, Code& exprs)
+compile(CEnv& cenv, const Code& parsed, Code& defs, bool& hasMain, const char* mainName)
 {
 	const AST* exp = NULL;
 
@@ -111,6 +111,7 @@ compile(CEnv& cenv, const Code& parsed, Code& defs, Code& exprs)
 
 	// Flatten expressions
 	const AST* retT = NULL;
+	Code       exprs;
 	for (Code::const_iterator i = lifted.begin(); i != lifted.end(); ++i) {
 		const ATuple* call = (*i)->to_tuple();
 		if (call && (is_form(*i, "def-type")
@@ -125,6 +126,25 @@ compile(CEnv& cenv, const Code& parsed, Code& defs, Code& exprs)
 		}
 	}
 
+	// Flatten top-level code into main function
+	if (!exprs.empty()) {
+		const ASymbol* main = cenv.penv.sym(mainName);
+
+		List mainT(Cursor(), cenv.penv.sym("Fn"), new ATuple(Cursor()), retT, NULL);
+		cenv.def(main, NULL, mainT, NULL);
+			
+		defs.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)
+			mainRet = resp_flatten(cenv, defs, *i);
+			
+		defs.push_back(
+			tup(Cursor(), cenv.penv.sym("fn-end"), main, mainRet, NULL));
+	}
+
+	hasMain = !exprs.empty();
 	return retT;
 }
 
@@ -145,27 +165,9 @@ eval(CEnv& cenv, Cursor& cursor, istream& is, bool execute)
 			return 0;
 		}
 
-		Code       defs;
-		Code       exprs;
-		const AST* retT = compile(cenv, parsed, defs, exprs);
-
-		// Flatten main code into `defs'
-		if (!exprs.empty()) {
-			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);
-			
-			defs.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)
-				mainRet = resp_flatten(cenv, defs, *i);
-			
-			defs.push_back(
-				tup(Cursor(), cenv.penv.sym("fn-end"), main, mainRet, NULL));
-		}
-
+		Code defs;
+		bool hasMain;
+		const AST* retT = compile(cenv, parsed, defs, hasMain, "main");
 		if (cenv.args.find("-F") != cenv.args.end()) {
 			dump(cenv, defs);
 			return 0;
@@ -174,15 +176,15 @@ eval(CEnv& cenv, Cursor& cursor, istream& is, bool execute)
 		// Compile flattened code
 		for (Code::const_iterator i = defs.begin(); i != defs.end(); ++i) {
 			resp_compile(cenv, *i);
+			Object::pool.addRoot(*i);
 		}
-
 		if (cenv.args.find("-S") != cenv.args.end()) {
 			cenv.engine()->writeModule(cenv, cenv.out);
 			return 0;
 		}
 
 		// Call main and print result
-		if (!exprs.empty())
+		if (hasMain)
 			callPrintCollect(cenv, cenv.engine()->getFn(cenv, "main"),
 			                 ast, retT, execute);
 
@@ -199,9 +201,8 @@ repl(CEnv& cenv)
 {
 	cenv.repl = true;
 
-	const AST*   exp        = NULL;
-	const AST*   ast        = NULL;
-	const string replFnName = cenv.penv.gensymstr("_repl");
+	const AST* exp = NULL;
+	const AST* ast = NULL;
 	while (1) {
 		cenv.out << "() ";
 		cenv.out.flush();
@@ -211,24 +212,31 @@ repl(CEnv& cenv)
 			if (!readParseType(cenv, cursor, std::cin, exp, ast))
 				break;
 
-			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;
-			try {
-				// Create function for this repl loop
-				f = cenv.engine()->startFn(cenv, replFnName, new ATuple(cursor), fnT);
-				cenv.engine()->finishFn(cenv, resp_compile(cenv, ast), type);
-
-				if (cenv.args.find("-S") != cenv.args.end())
-					cenv.engine()->writeModule(cenv, cenv.out);
-				else
-					callPrintCollect(cenv, f, ast, type, true);
-			} catch (Error& e) {
-				cenv.out << e.msg << endl;
-				cenv.engine()->eraseFn(cenv, f);
+			Code parsed;
+			parsed.push_back(ast);
+				
+			Code defs;
+			bool hasMain;
+			const std::string replName = cenv.penv.gensymstr("_repl");
+			const AST* retT = compile(cenv, parsed, defs, hasMain, replName.c_str());
+			if (cenv.args.find("-F") != cenv.args.end()) {
+				dump(cenv, defs);
+				continue;
+			}
+
+			// Compile flattened code
+			for (Code::const_iterator i = defs.begin(); i != defs.end(); ++i) {
+				resp_compile(cenv, *i);
+			}
+			if (cenv.args.find("-S") != cenv.args.end()) {
+				cenv.engine()->writeModule(cenv, cenv.out);
+				continue;
+			}
+
+			// Call main and print result
+			if (hasMain) {
+				callPrintCollect(cenv, cenv.engine()->getFn(cenv, replName),
+				                 ast, retT, true);
 			}
 
 		} catch (Error& e) {