feat(classes): support ordinary member methods and this references

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

7 files changed, 138 insertions, 29 deletions

diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index ffabf8d..01393a9 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -169,6 +169,13 @@ void CodeGen::compile(const std::vector<std::unique_ptr<StmtAST>> &Statements) {
         }
     }
 
+    for (const auto &[Class, Methods] : ClassMethods) {
+        for (const auto &[Name, Method] : Methods) {
+            SaveAndRestore<int> Location(DiagnosticLine, Method->getLine());
+            FuncGen->emitPrototype(Method->getProto());
+        }
+    }
+
     for (const auto &Stmt : Statements) {
         SaveAndRestore<int> Location(DiagnosticLine, Stmt->getLine());
         emitStmt(Stmt.get());
diff --git a/lib/CodeGen/ConstantGen.cc b/lib/CodeGen/ConstantGen.cc
index e14e22f..12bf4b1 100644
--- a/lib/CodeGen/ConstantGen.cc
+++ b/lib/CodeGen/ConstantGen.cc
@@ -16,6 +16,10 @@ void CodeGen::emitScopedBlock(const std::vector<std::unique_ptr<StmtAST>> &Body)
 }
 
 bool CodeGen::prepareSymbolDeclaration(const std::string &Name, bool IsConstant) {
+    if (!CurrentClass.empty() && sameIdentifier(Name, "this")) {
+        reportError("Cannot redeclare this inside a method");
+        return false;
+    }
     auto &Scope = ConstantScopes.back();
     auto It = Symbols.find(Name);
     if (It != Symbols.end()) {
diff --git a/lib/CodeGen/Designator.cc b/lib/CodeGen/Designator.cc
index fcc6919..adf08f1 100644
--- a/lib/CodeGen/Designator.cc
+++ b/lib/CodeGen/Designator.cc
@@ -40,6 +40,59 @@ void CodeGen::runTypePrePass(const std::vector<std::unique_ptr<StmtAST>> &Statem
         }
         if (!Ok) reportError("%s", Err.c_str());
     }
+
+    for (const auto &S : Statements) {
+        auto *Class = dynamic_cast<ClassTypeDeclAST*>(S.get());
+        if (!Class) continue;
+        auto &Methods = ClassMethods[Class->getName()];
+        IdentifierSet Names;
+        for (const auto &Field : Class->getFields()) {
+            Names.insert(Field.Name);
+            if (sameIdentifier(Field.Name, "this")) {
+                SaveAndRestore<int> Location(DiagnosticLine, Class->getLine());
+                reportError("A class field cannot be named this");
+            }
+        }
+        for (const auto &Method : Class->getMethods()) {
+            SaveAndRestore<int> Location(DiagnosticLine, Method->getLine());
+            const std::string &Name = Method->getProto()->getName();
+            if (sameIdentifier(Name, "NEW")) {
+                reportError("NEW constructors are not implemented");
+            } else if (sameIdentifier(Name, "this")) {
+                reportError("A method cannot be named this");
+            } else if (!Names.insert(Name).second) {
+                reportError("Duplicate member %s in class %s", Name.c_str(), Class->getName().c_str());
+            } else {
+                Methods.emplace(Name, Method.get());
+            }
+        }
+    }
+}
+
+const PrototypeAST *CodeGen::findMethod(const std::string &Class, const std::string &Name) const {
+    auto C = ClassMethods.find(Class);
+    if (C == ClassMethods.end()) return nullptr;
+    auto M = C->second.find(Name);
+    return M == C->second.end() ? nullptr : M->second->getProto();
+}
+
+Value *CodeGen::emitMemberCall(MemberCallExprAST *Call) {
+    SaveAndRestore<int> Location(DiagnosticLine, Call->getLine());
+    LValueInfo Receiver;
+    if (!emitLValue(Call->getReceiver(), Receiver)) return nullptr;
+    const PrototypeAST *Method = findMethod(Receiver.TypeName, Call->getMethod());
+    if (!Method) {
+        reportError("Type %s has no method %s", Receiver.TypeName.c_str(), Call->getMethod().c_str());
+        return nullptr;
+    }
+    if (!Method->isPublic() && !sameIdentifier(CurrentClass, Receiver.TypeName)) {
+        reportError("Cannot access private method %s of type %s",
+                    Call->getMethod().c_str(), Receiver.TypeName.c_str());
+        return nullptr;
+    }
+    std::vector<Value*> Args = {Receiver.Addr};
+    if (!marshalCallArgs(Method->getLinkName(), Call->getArgs(), Call->getLine(), Args)) return nullptr;
+    return FuncGen->emitCall(Method->getLinkName(), Args);
 }
 
 // Resolves a designator (or legacy variable/array-access node) to the address
@@ -126,7 +179,7 @@ bool CodeGen::emitLValue(ExprAST *E, LValueInfo &Out) {
                 reportError("Record type %s has no field %s", TypeName.c_str(), A.FieldName.c_str());
                 return false;
             }
-            if (!FI->IsPublic) {
+            if (!FI->IsPublic && !sameIdentifier(CurrentClass, TypeName)) {
                 reportError("Cannot access private field %s of type %s",
                             A.FieldName.c_str(), TypeName.c_str());
                 return false;
diff --git a/lib/CodeGen/ExpressionGen.cc b/lib/CodeGen/ExpressionGen.cc
index 9ea79d8..fde948c 100644
--- a/lib/CodeGen/ExpressionGen.cc
+++ b/lib/CodeGen/ExpressionGen.cc
@@ -127,6 +127,9 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
         return emitBinaryOp(Bin->getOp(), L, LT, R, RT, Bin->getLine());
     }
 
+    if (auto *Call = dynamic_cast<MemberCallExprAST*>(Expr))
+        return emitMemberCall(Call);
+
     if (auto *Call = dynamic_cast<CallExprAST*>(Expr))
         return emitCall(Call->getCallee(), Call->getArgs(), Call->getLine());
 
diff --git a/lib/CodeGen/FunctionGen.cc b/lib/CodeGen/FunctionGen.cc
index 10bd3fa..d5bc5d4 100644
--- a/lib/CodeGen/FunctionGen.cc
+++ b/lib/CodeGen/FunctionGen.cc
@@ -45,9 +45,12 @@ Type *FunctionGen::getLLVMType(const std::string &TypeName) {
     return Resolved;
 }
 
-void FunctionGen::createArgumentAllocas(Function *F, const std::vector<ParamDecl> &Params) {
+void FunctionGen::createArgumentAllocas(Function *F, const std::vector<ParamDecl> &Params, bool IsMethod) {
     Function::arg_iterator AI = F->arg_begin();
+    if (IsMethod) ++AI;
     for (const ParamDecl &P : Params) {
+        bool InvalidThis = IsMethod && sameIdentifier(P.Name, "this");
+        if (InvalidThis) reportError("A method parameter cannot be named this");
         if (Types.lookupEnumConstant(P.Name)) {
             reportError("parameter '%s' collides with an enum value", P.Name.c_str());
         }
@@ -72,10 +75,11 @@ void FunctionGen::createArgumentAllocas(Function *F, const std::vector<ParamDecl
                 LBs.push_back(LB);
                 UBs.push_back(UB);
             }
-            if (BindArrayParam) BindArrayParam(P, ArgVal, LBs, UBs);
+            if (!InvalidThis && BindArrayParam) BindArrayParam(P, ArgVal, LBs, UBs);
             continue;
         }
 
+        if (InvalidThis) continue;
         if (P.Mode == PassMode::ByRef) {
             Symbols[P.Name] = {ArgVal, P.TypeName, false};
             continue;
@@ -96,32 +100,35 @@ const FuncSig *FunctionGen::getSignature(const std::string &Name) const {
 }
 
 Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
-    if (isReservedRuntimeName(Proto->getName())) {
-        reportError("'%s' is a reserved runtime symbol and cannot be a FUNCTION/PROCEDURE name",
-                Proto->getName().c_str());
-        return nullptr;
-    }
-    if (CodeGen::isBuiltinName(Proto->getName())) {
-        reportError("'%s' is a built-in function and cannot be redefined",
-                Proto->getName().c_str());
-        return nullptr;
-    }
-    if (Types.lookupEnumConstant(Proto->getName())) {
-        reportError("'%s' is an enum value and cannot be a FUNCTION/PROCEDURE name",
-                Proto->getName().c_str());
-        return nullptr;
-    }
-    if (Types.resolve(Proto->getName())) {
-        reportError("'%s' is a type name and cannot be a FUNCTION/PROCEDURE name",
-                Proto->getName().c_str());
-        return nullptr;
+    if (Proto->getOwnerClass().empty()) {
+        if (isReservedRuntimeName(Proto->getName())) {
+            reportError("'%s' is a reserved runtime symbol and cannot be a FUNCTION/PROCEDURE name",
+                    Proto->getName().c_str());
+            return nullptr;
+        }
+        if (CodeGen::isBuiltinName(Proto->getName())) {
+            reportError("'%s' is a built-in function and cannot be redefined",
+                    Proto->getName().c_str());
+            return nullptr;
+        }
+        if (Types.lookupEnumConstant(Proto->getName())) {
+            reportError("'%s' is an enum value and cannot be a FUNCTION/PROCEDURE name",
+                    Proto->getName().c_str());
+            return nullptr;
+        }
+        if (Types.resolve(Proto->getName())) {
+            reportError("'%s' is a type name and cannot be a FUNCTION/PROCEDURE name",
+                    Proto->getName().c_str());
+            return nullptr;
+        }
     }
 
     FuncSig Sig;
-    Sig.Name = Proto->getName();
+    Sig.Name = Proto->getLinkName();
     Sig.ReturnTypeName = Proto->getReturnType();
 
     std::vector<Type*> ArgTypes;
+    if (!Proto->getOwnerClass().empty()) ArgTypes.push_back(PointerType::getUnqual(Context));
     for (const ParamDecl &P : Proto->getParams()) {
         const TypeInfo *Info = Types.resolve(P.TypeName);
         if (Info && Info->isVoid()) {
@@ -153,23 +160,24 @@ Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
 
     Type *RetType = getLLVMType(Proto->getReturnType());
     FunctionType *FT = FunctionType::get(RetType, ArgTypes, false);
-    const FuncSig *Existing = getSignature(Proto->getName());
-    Function *F = Module.getFunction(Existing ? Existing->Name : Proto->getName());
+    const FuncSig *Existing = getSignature(Proto->getLinkName());
+    Function *F = Module.getFunction(Existing ? Existing->Name : Proto->getLinkName());
     if (F && F->getFunctionType() != FT) {
         reportError("Function %s conflicts with an existing declaration of a different type",
                 Proto->getName().c_str());
         return nullptr;
     }
     if (!F) {
-        F = Function::Create(FT, Function::ExternalLinkage, Proto->getName(), &Module);
+        F = Function::Create(FT, Function::ExternalLinkage, Proto->getLinkName(), &Module);
     }
     Exceptions.prepareFunction(F);
 
     // Keep the first registration: on a duplicate definition the first one
     // wins and the second is rejected by emitFunctionDef.
-    Signatures.emplace(Proto->getName(), std::move(Sig));
+    Signatures.emplace(Proto->getLinkName(), std::move(Sig));
 
     Function::arg_iterator AI = F->arg_begin();
+    if (!Proto->getOwnerClass().empty()) (AI++)->setName("this");
     for (const ParamDecl &P : Proto->getParams()) {
         if (AI == F->arg_end()) break;
         (AI++)->setName(P.Name);
@@ -202,7 +210,19 @@ Function *FunctionGen::emitFunctionDef(FunctionDefAST *FuncAST,
         if (Info.ConstantValue && Info.ScopeID == 0) Symbols.emplace(Name, Info);
     }
 
-    createArgumentAllocas(TheFunction, Proto->getParams());
+    bool IsMethod = !Proto->getOwnerClass().empty();
+    if (IsMethod) {
+        const TypeInfo *Class = Types.resolve(Proto->getOwnerClass());
+        Value *Self = TheFunction->getArg(0);
+        if (Class && Class->asRecord()) {
+            for (const auto &[Name, Field] : Class->asRecord()->Fields) {
+                Value *Addr = Builder.CreateStructGEP(Class->LLVMType, Self, Field.Index, Name);
+                Symbols[Name] = {Addr, Field.TypeName, false};
+            }
+        }
+        Symbols["this"] = {Self, Proto->getOwnerClass(), false, nullptr, ScopeID};
+    }
+    createArgumentAllocas(TheFunction, Proto->getParams(), IsMethod);
     for (const auto &P : Proto->getParams()) {
         auto It = Symbols.find(P.Name);
         if (It != Symbols.end()) It->second.ScopeID = ScopeID;
diff --git a/lib/CodeGen/StatementGen.cc b/lib/CodeGen/StatementGen.cc
index b648ec6..c55aebe 100644
--- a/lib/CodeGen/StatementGen.cc
+++ b/lib/CodeGen/StatementGen.cc
@@ -12,7 +12,17 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
     SaveAndRestore<int> Location(DiagnosticLine, Stmt && Stmt->getLine() ? Stmt->getLine() : DiagnosticLine);
     if (!Stmt) return;
 
+    if (auto *Class = dynamic_cast<ClassTypeDeclAST*>(Stmt)) {
+        for (const auto &Method : Class->getMethods()) {
+            // Invalid or duplicate members have already been diagnosed.
+            if (findMethod(Class->getName(), Method->getProto()->getName()) == Method->getProto())
+                emitStmt(Method.get());
+        }
+        return;
+    }
+
     if (auto *FuncDef = dynamic_cast<FunctionDefAST*>(Stmt)) {
+        SaveAndRestore<std::string> ClassContext(CurrentClass, FuncDef->getProto()->getOwnerClass());
         BasicBlock *SavedBlock = Builder->GetInsertBlock();
         // Function bodies get a fresh array scope (metadata SSA values belong
         // to the emitting function), mirroring FunctionGen's Symbols handling.
@@ -54,6 +64,11 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
         return;
     }
 
+    if (auto *Call = dynamic_cast<MemberCallStmtAST*>(Stmt)) {
+        emitMemberCall(Call->getCall());
+        return;
+    }
+
     if (auto *Call = dynamic_cast<CallStmtAST*>(Stmt)) {
         emitCall(Call->getCallee(), Call->getArgs(), Call->getLine());
         return;
diff --git a/lib/CodeGen/ValueTypeGen.cc b/lib/CodeGen/ValueTypeGen.cc
index 1ee4a2d..16c3e0b 100644
--- a/lib/CodeGen/ValueTypeGen.cc
+++ b/lib/CodeGen/ValueTypeGen.cc
@@ -52,7 +52,7 @@ const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) const {
                 const RecordPayload *RP = Cur->asRecord();
                 const RecordFieldInfo *FI = RP ? RP->findField(A.FieldName) : nullptr;
                 if (!FI) return nullptr;
-                if (!FI->IsPublic) return nullptr;
+                if (!FI->IsPublic && !sameIdentifier(CurrentClass, TypeName)) return nullptr;
                 TypeName = FI->TypeName;
             } else if (A.Kind == AccessKind::Deref) {
                 const PointerPayload *PP = Cur->asPointer();
@@ -117,6 +117,13 @@ const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) const {
         return resolveType("INTEGER");
     }
 
+    if (auto *Call = dynamic_cast<MemberCallExprAST*>(Expr)) {
+        const TypeInfo *Receiver = getExprTypeInfo(Call->getReceiver());
+        const PrototypeAST *Method = Receiver ? findMethod(Receiver->Name, Call->getMethod()) : nullptr;
+        if (!Method || (!Method->isPublic() && !sameIdentifier(CurrentClass, Receiver->Name))) return nullptr;
+        return resolveType(Method->getReturnType());
+    }
+
     if (auto *Call = dynamic_cast<CallExprAST*>(Expr)) {
         const std::string &Name = Call->getCallee();
         if (const BuiltinInfo *B = findBuiltin(Name)) {