1 files changed, 3 insertions, 3 deletions
diff --git a/lib/CodeGen/StatementGen.cc b/lib/CodeGen/StatementGen.cc
index 5b6a13c..1a9c285 100644
--- a/lib/CodeGen/StatementGen.cc
+++ b/lib/CodeGen/StatementGen.cc
@@ -99,9 +99,9 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
 
     if (auto *Ret = dynamic_cast<ReturnStmtAST*>(Stmt)) {
         Function *F = Builder->GetInsertBlock()->getParent();
-        const FuncSig *Sig = FuncGen->getSignature(F->getName().str());
-        // Pseudocode cannot name a function "main" (reserved runtime symbol),
-        // so no signature means we are emitting the top level.
+        const FuncSig *Sig = F->getName() == "main" ? nullptr
+            : FuncGen->getSignature(F->getName().str());
+        // The generated wrapper has no signature; user MAIN uses __cps_main.
         if (!Sig) {
             reportError("RETURN is only allowed inside a FUNCTION or PROCEDURE");
             return;