feat(codegen): FUNCTION MAIN entry point

Invoke the internal entry function through the generated exception handler
and truncate its INTEGER result to the process exit status.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

4 files changed, 42 insertions, 11 deletions

diff --git a/include/cps/FunctionAST.h b/include/cps/FunctionAST.h
index 9b2d49b..4dfe908 100644
--- a/include/cps/FunctionAST.h
+++ b/include/cps/FunctionAST.h
@@ -1,5 +1,6 @@
 #pragma once
 #include "cps/AST.h"
+#include "cps/Identifier.h"
 #include <cstdint>
 #include <vector>
 #include <string>
@@ -46,7 +47,10 @@ public:
     bool isExternal() const { return IsExternal; }
     const std::string &getOwnerClass() const { return OwnerClass; }
     bool isPublic() const { return IsPublic; }
-    std::string getLinkName() const { return OwnerClass.empty() ? Name : OwnerClass + "." + Name; }
+    std::string getLinkName() const {
+        if (!OwnerClass.empty()) return OwnerClass + "." + Name;
+        return sameIdentifier(Name, "MAIN") ? "__cps_main" : Name;
+    }
 };
 
 class FunctionDefAST : public StmtAST {
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 6aa1558..eba7d77 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -189,6 +189,8 @@ void CodeGen::compile(const std::vector<SourceUnit> &Sources) {
 
     IdentifierMap<size_t> Definitions;
     std::optional<size_t> CodeUnit;
+    FunctionDefAST *Entry = nullptr;
+    size_t EntryUnit = 0;
     auto Reserve = [&](const std::string &Name) {
         auto [It, Added] = Definitions.emplace(Name, CurrentUnit);
         if (!Added && It->second != CurrentUnit)
@@ -202,6 +204,15 @@ void CodeGen::compile(const std::vector<SourceUnit> &Sources) {
             std::string TypeName;
             if (auto *F = dynamic_cast<FunctionDefAST*>(S.get())) {
                 Reserve(F->getProto()->getName());
+                if (sameIdentifier(F->getProto()->getName(), "MAIN")) {
+                    if (Entry) reportError("Multiple MAIN definitions in %s and %s",
+                                           Sources[EntryUnit].Path.c_str(), Sources[I].Path.c_str());
+                    Entry = F;
+                    EntryUnit = I;
+                    if (!F->getProto()->getParams().empty() ||
+                        !sameIdentifier(F->getProto()->getReturnType(), "INTEGER"))
+                        reportError("MAIN must be a FUNCTION with no parameters returning INTEGER");
+                }
             } else if (auto *E = dynamic_cast<EnumTypeDeclAST*>(S.get())) {
                 TypeName = E->getName();
                 for (const auto &V : E->getValueNames()) Reserve(V);
@@ -231,7 +242,14 @@ void CodeGen::compile(const std::vector<SourceUnit> &Sources) {
         }
     }
     if (HadError) return;
-    MainUnit = CodeUnit.value_or(0);
+    if (Entry && CodeUnit) {
+        selectUnit(EntryUnit);
+        SaveAndRestore<int> Location(DiagnosticLine, Entry->getLine());
+        reportError("FUNCTION MAIN in %s cannot coexist with top-level code in %s",
+                    Sources[EntryUnit].Path.c_str(), Sources[*CodeUnit].Path.c_str());
+        return;
+    }
+    MainUnit = Entry ? EntryUnit : CodeUnit.value_or(0);
     FunctionType *FT = FunctionType::get(Type::getInt32Ty(*TheContext), false);
     Function *F = Function::Create(FT, Function::ExternalLinkage, "main", TheModule.get());
     BasicBlock *BB = BasicBlock::Create(*TheContext, "entry", F);
@@ -268,7 +286,8 @@ void CodeGen::compile(const std::vector<SourceUnit> &Sources) {
             if (auto *Class = dynamic_cast<ClassTypeDeclAST*>(Stmt.get())) {
                 for (const auto &Method : Class->getMethods()) {
                     SaveAndRestore<int> MethodLocation(DiagnosticLine, Method->getLine());
-                    FuncGen->emitPrototype(Method->getProto());
+                    if (findMethod(Class->getName(), Method->getProto()->getName()) == Method->getProto())
+                        FuncGen->emitPrototype(Method->getProto());
                 }
             }
         }
@@ -300,13 +319,18 @@ void CodeGen::compile(const std::vector<SourceUnit> &Sources) {
         for (const auto &Stmt : Sources[I].Statements) emitStmt(Stmt.get());
     }
     selectUnit(MainUnit);
-    if (!CodeUnit && !HadError) {
+    if (!CodeUnit && !Entry && !HadError) {
         SaveAndRestore<int> Location(DiagnosticLine, Sources.size() > 1 ? 1 : 0);
         reportError("no main program");
     }
 
     if (!Builder->GetInsertBlock()->hasTerminator()) {
-        Exceptions->recordExit(Builder->CreateRet(ConstantInt::get(*TheContext, APInt(32, 0))));
+        Value *Status = ConstantInt::get(*TheContext, APInt(32, 0));
+        if (Entry) {
+            Value *Result = Exceptions->emitCall(TheModule->getFunction("__cps_main"), {}, "main_result");
+            Status = Builder->CreateTrunc(Result, Type::getInt32Ty(*TheContext));
+        }
+        Exceptions->recordExit(Builder->CreateRet(Status));
     }
 
     finishStackObjects(F);
diff --git a/lib/CodeGen/FunctionGen.cc b/lib/CodeGen/FunctionGen.cc
index 220a97a..47971cf 100644
--- a/lib/CodeGen/FunctionGen.cc
+++ b/lib/CodeGen/FunctionGen.cc
@@ -12,7 +12,7 @@ using namespace cps;
 // FUNCTION/PROCEDURE with one of these names would hijack the runtime.
 static bool isReservedRuntimeName(const std::string &Name) {
     static const char *Reserved[] = {
-        "main", "printf", "scanf", "malloc", "free", "strlen", "memcpy",
+        "__cps_main", "printf", "scanf", "malloc", "free", "strlen", "memcpy",
         "strcpy", "strcat", "toupper", "tolower", "sprintf", "snprintf", "strtol", "strtoll",
         "strtod", "memset", "exit",
         "fopen", "fclose", "fgetc", "ungetc", "fputs", "fputc",
@@ -105,7 +105,7 @@ void FunctionGen::createArgumentAllocas(Function *F, const std::vector<ParamDecl
 }
 
 const FuncSig *FunctionGen::getSignature(const std::string &Name) const {
-    auto It = Signatures.find(Name);
+    auto It = Signatures.find(sameIdentifier(Name, "MAIN") ? "__cps_main" : Name);
     return It == Signatures.end() || (Visible && !Visible(It->second.OwnerUnit))
                ? nullptr : &It->second;
 }
@@ -180,7 +180,10 @@ Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
         return nullptr;
     }
     if (!F) {
-        F = Function::Create(FT, Function::ExternalLinkage, Proto->getLinkName(), &Module);
+        F = Function::Create(FT,
+            Proto->getOwnerClass().empty() && sameIdentifier(Proto->getName(), "MAIN")
+                ? Function::InternalLinkage : Function::ExternalLinkage,
+            Proto->getLinkName(), &Module);
     }
     Exceptions.prepareFunction(F);
 
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;