feat(overload): select overloads by return type as well as parameters

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

8 files changed, 117 insertions, 51 deletions

diff --git a/lib/CodeGen/CTypeGen.cc b/lib/CodeGen/CTypeGen.cc
index 90f0aa3..b12537a 100644
--- a/lib/CodeGen/CTypeGen.cc
+++ b/lib/CodeGen/CTypeGen.cc
@@ -49,6 +49,7 @@ Value *CodeGen::emitCast(CastExprAST *Cast) {
         reportError("Invalid cast target %s", Cast->getTypeName().c_str());
         return nullptr;
     }
+    expectResult(Operand, Target, /*Exact=*/false);
     if (isWholeArrayVar(Operand)) {
         reportError("Cannot cast a whole array; index it");
         return nullptr;
diff --git a/lib/CodeGen/CallGen.cc b/lib/CodeGen/CallGen.cc
index 43c0dd4..30f0d61 100644
--- a/lib/CodeGen/CallGen.cc
+++ b/lib/CodeGen/CallGen.cc
@@ -171,7 +171,8 @@ Value *CodeGen::emitVariadicArgument(ExprAST *Expr, const std::string &Callee, s
 // Resolve using static source types only; argument evaluation belongs to marshalling.
 const FuncSig *CodeGen::resolveCall(const std::string &Name,
                                   const std::vector<std::unique_ptr<ExprAST>> &Args,
-                                  const std::vector<std::string> *Explicit) {
+                                  const std::vector<std::string> *Explicit,
+                                  const ResultContext *Context) {
     auto Overloads = FuncGen->getOverloads(Name);
     bool HasTemplates = hasFunctionTemplate(Name);
     bool HasExplicit = Explicit && !Explicit->empty();
@@ -202,52 +203,77 @@ const FuncSig *CodeGen::resolveCall(const std::string &Name,
         if (Rank) Key += "ARRAY" + std::to_string(Rank) + " OF ";
         Key += Type ? Type->Name : "?";
     }
-    const FuncSig *Ordinary = nullptr;
-    for (const FuncSig *Sig : Overloads) {
-        if (HasExplicit) break;
-        if (Sig->Params.size() != Actual.size()) continue;
-        bool Match = true;
-        for (size_t I = 0; I < Actual.size(); ++I) {
-            const ParamSig &P = Sig->Params[I];
-            if (Actual[I].Rank != (P.IsArray ? P.Rank : 0) || !Actual[I].Type ||
-                !Actual[I].Type->sameType(Types->resolve(P.TypeName))) {
-                Match = false;
-                break;
-            }
-        }
-        if (Match) {
-            if (Ordinary && HasTemplates) {
-                reportError("Ambiguous call to %s(%s)", Name.c_str(), Key.c_str());
-                return nullptr;
-            }
-            Ordinary = Sig;
-            if (!HasTemplates) return Sig;
-        }
-    }
-    if (Ordinary) return Ordinary;
-    // Failing an exact match, CHAR and STRING arguments convert into one
-    // another; the overload must be unique.
-    if (!HasTemplates && !HasExplicit) {
-        auto Scalar = [](const TypeInfo *T) { return T && (T->Kind == TypeKind::Char || T->isString()); };
-        const FuncSig *Converted = nullptr;
+    // A context with a type prefers candidates that return it; a strict
+    // context requires one. Without a context (or after a cast's preference
+    // finds nothing) selection stays parameter-only.
+    const bool PreferResult = Context && Context->Type;
+    auto ReturnsWanted = [&](const FuncSig *Sig) {
+        const TypeInfo *Ret = resolveType(Sig->ReturnTypeName);
+        return Ret && Ret->sameType(Context->Type);
+    };
+    auto Scalar = [](const TypeInfo *T) { return T && (T->Kind == TypeKind::Char || T->isString()); };
+    bool ParamMatched = false;
+    int Count = 0;
+    // One layer of parameter matching. Lenient adds CHAR/STRING conversions;
+    // FilterReturn keeps only candidates whose result is the wanted type.
+    auto Scan = [&](bool Lenient, bool FilterReturn) -> const FuncSig * {
+        const FuncSig *First = nullptr;
+        Count = 0;
         for (const FuncSig *Sig : Overloads) {
+            if (HasExplicit) break;
             if (Sig->Params.size() != Actual.size()) continue;
             bool Match = true;
             for (size_t I = 0; I < Actual.size() && Match; ++I) {
                 const ParamSig &P = Sig->Params[I];
                 const TypeInfo *Formal = Types->resolve(P.TypeName);
-                Match = Actual[I].Rank == (P.IsArray ? P.Rank : 0) && Actual[I].Type &&
-                        (Actual[I].Type->sameType(Formal) ||
-                         (!P.IsByRef && !P.IsArray && Scalar(Actual[I].Type) && Scalar(Formal)));
+                bool SameRank = Actual[I].Rank == (P.IsArray ? P.Rank : 0);
+                bool Same = SameRank && Actual[I].Type && Actual[I].Type->sameType(Formal);
+                Match = Same ||
+                    (Lenient && SameRank && Actual[I].Type && !P.IsByRef && !P.IsArray &&
+                     Scalar(Actual[I].Type) && Scalar(Formal));
             }
             if (!Match) continue;
-            if (Converted) {
-                reportError("Ambiguous call to %s(%s)", Name.c_str(), Key.c_str());
-                return nullptr;
-            }
-            Converted = Sig;
+            ParamMatched = true;
+            if (FilterReturn && !ReturnsWanted(Sig)) continue;
+            ++Count;
+            if (!First) First = Sig;
         }
-        if (Converted) return Converted;
+        return First;
+    };
+    auto ReportAmbiguous = [&](bool FilterReturn) {
+        if (FilterReturn)
+            reportError("Ambiguous call to %s(%s)", Name.c_str(), Key.c_str());
+        else
+            reportError("Ambiguous call to %s(%s); the return type cannot be determined from context",
+                        Name.c_str(), Key.c_str());
+    };
+    // Exact parameters first, then CHAR/STRING conversions.
+    bool Ambiguous = false;
+    auto Try = [&](bool FilterReturn, const FuncSig *&Selected) {
+        Selected = Scan(false, FilterReturn);
+        if (Count > 1) {
+            ReportAmbiguous(FilterReturn);
+            Ambiguous = true;
+            return;
+        }
+        if (Selected || HasTemplates || HasExplicit) return;
+        Selected = Scan(true, FilterReturn);
+        if (Count > 1) {
+            ReportAmbiguous(FilterReturn);
+            Ambiguous = true;
+        }
+    };
+
+    const FuncSig *Selected = nullptr;
+    if (PreferResult) {
+        Try(true, Selected);
+        if (Ambiguous) return nullptr;
+        if (Selected) return Selected;
+    }
+    if (!Context || !Context->Exact) {
+        Try(false, Selected);
+        if (Ambiguous) return nullptr;
+        if (Selected) return Selected;
     }
     if (HasTemplates) {
         // A matching non-variadic template wins before any variadic candidate.
@@ -323,6 +349,11 @@ const FuncSig *CodeGen::resolveCall(const std::string &Name,
             if (Selected) return instantiateFunction(*Selected, SelectedBindings);
         }
     }
+    if (Context && Context->Exact && ParamMatched) {
+        reportError("No matching overload for %s(%s) returning %s", Name.c_str(), Key.c_str(),
+                    Context->Type->Name.c_str());
+        return nullptr;
+    }
     reportError("No matching overload for %s(%s)", Name.c_str(), Key.c_str());
     return nullptr;
 }
diff --git a/lib/CodeGen/ConstructorGen.cc b/lib/CodeGen/ConstructorGen.cc
index e956f9d..f598494 100644
--- a/lib/CodeGen/ConstructorGen.cc
+++ b/lib/CodeGen/ConstructorGen.cc
@@ -20,6 +20,7 @@ void CodeGen::transferObjects(const std::vector<size_t> &Objects) {
 
 Value *CodeGen::emitOwnedValue(ExprAST *Expr, const TypeInfo *Type,
                                std::vector<size_t> &Transfers) {
+    expectResult(Expr, Type);
     if (!Type || (!Type->IsClass && !containsString(Type))) return emitCoercedExpr(Expr, Type);
     const TypeInfo *Source = getExprTypeInfo(Expr);
     // A CHAR argument becomes a fresh STRING temporary, owned like any other.
@@ -146,6 +147,7 @@ Value *CodeGen::symbolAddress(const SymbolInfo &Info) {
 
 void CodeGen::emitAssignment(const LValueInfo &Target, ExprAST *Expr, const SymbolInfo *Symbol) {
     const TypeInfo *Info = resolveType(Target.TypeName);
+    expectResult(Expr, Info);
     if (auto *New = dynamic_cast<NewExprAST*>(Expr)) {
         SaveAndRestore<int> Location(DiagnosticLine, New->getLine());
         const TypeInfo *Class = resolveType(New->getClassName());
diff --git a/lib/CodeGen/Designator.cc b/lib/CodeGen/Designator.cc
index 6737937..d790d27 100644
--- a/lib/CodeGen/Designator.cc
+++ b/lib/CodeGen/Designator.cc
@@ -96,16 +96,17 @@ bool CodeGen::isRegisteredMethod(const std::string &Class, const PrototypeAST *M
 }
 
 const FuncSig *CodeGen::resolveMethod(const std::string &Class, const std::string &Name,
-                                    const std::vector<std::unique_ptr<ExprAST>> &Args) {
+                                    const std::vector<std::unique_ptr<ExprAST>> &Args,
+                                    const ResultContext *Context) {
     const auto *Methods = findMethod(Class, Name);
     if (!Methods) {
         reportError("Type %s has no method %s", Class.c_str(), Name.c_str());
         return nullptr;
     }
-    return resolveCall(Methods->front()->getProto()->getLinkName(), Args);
+    return resolveCall(Methods->front()->getProto()->getLinkName(), Args, nullptr, Context);
 }
 
-Value *CodeGen::emitMemberCall(MemberCallExprAST *Call) {
+Value *CodeGen::emitMemberCall(MemberCallExprAST *Call, const ResultContext *Context) {
     SaveAndRestore<int> Location(DiagnosticLine, Call->getLine());
     LValueInfo Receiver;
     if (isObjectLValue(Call->getReceiver())) {
@@ -120,14 +121,15 @@ Value *CodeGen::emitMemberCall(MemberCallExprAST *Call) {
         if (!Receiver.Addr) return nullptr;
     }
     return emitMethodCall(Receiver.TypeName, Receiver.Addr, Call->getMethod(),
-                          Call->getArgs(), Call->getLine());
+                          Call->getArgs(), Call->getLine(), Context);
 }
 
 Value *CodeGen::emitMethodCall(const std::string &Class, Value *Receiver,
                                 const std::string &Name,
-                                const std::vector<std::unique_ptr<ExprAST>> &ArgExprs, int Line) {
+                                const std::vector<std::unique_ptr<ExprAST>> &ArgExprs, int Line,
+                                const ResultContext *Context) {
     SaveAndRestore<int> Location(DiagnosticLine, Line);
-    const FuncSig *Sig = resolveMethod(Class, Name, ArgExprs);
+    const FuncSig *Sig = resolveMethod(Class, Name, ArgExprs, Context);
     if (!Sig) return nullptr;
     const PrototypeAST *Method = Sig->Prototype;
     if (!Method->isPublic() && !sameIdentifier(CurrentClass, Class)) {
diff --git a/lib/CodeGen/ExpressionGen.cc b/lib/CodeGen/ExpressionGen.cc
index 45583d3..38155e4 100644
--- a/lib/CodeGen/ExpressionGen.cc
+++ b/lib/CodeGen/ExpressionGen.cc
@@ -154,10 +154,11 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
     }
 
     if (auto *Call = dynamic_cast<MemberCallExprAST*>(Expr))
-        return emitMemberCall(Call);
+        return emitMemberCall(Call, resultContext(Expr));
 
     if (auto *Call = dynamic_cast<CallExprAST*>(Expr))
-        return emitCall(Call->getCallee(), Call->getArgs(), Call->getLine(), &Call->getTemplateArgs());
+        return emitCall(Call->getCallee(), Call->getArgs(), Call->getLine(),
+                        &Call->getTemplateArgs(), resultContext(Expr));
 
     if (auto *D = dynamic_cast<DesignatorExprAST*>(Expr)) {
         LValueInfo LV;
@@ -263,7 +264,8 @@ Value *CodeGen::emitBinaryOp(int Op, Value *L, const TypeInfo *LT,
 
 Value *CodeGen::emitCall(const std::string &Callee,
                          const std::vector<std::unique_ptr<ExprAST>> &ArgExprs, int Line,
-                         const std::vector<std::string> *Explicit) {
+                         const std::vector<std::string> *Explicit,
+                         const ResultContext *Context) {
     SaveAndRestore<int> Location(DiagnosticLine, Line ? Line : DiagnosticLine);
     std::string Name = Callee;
     if (const BuiltinInfo *B = findBuiltin(Name)) {
@@ -303,7 +305,7 @@ Value *CodeGen::emitCall(const std::string &Callee,
 
     std::vector<Value*> Args;
     std::vector<const TypeInfo *> Variadic;
-    const FuncSig *Sig = resolveCall(Name, ArgExprs, Explicit);
+    const FuncSig *Sig = resolveCall(Name, ArgExprs, Explicit, Context);
     if (!Sig || !marshalCallArgs(*Sig, ArgExprs, Line, Args, &Variadic)) return nullptr;
     return FuncGen->emitCall(*Sig, Args, Variadic);
 }
diff --git a/lib/CodeGen/FunctionGen.cc b/lib/CodeGen/FunctionGen.cc
index c233205..2e927f9 100644
--- a/lib/CodeGen/FunctionGen.cc
+++ b/lib/CodeGen/FunctionGen.cc
@@ -229,6 +229,10 @@ Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
     }
 
     Type *RetType = getLLVMType(Proto->getReturnType());
+    const TypeInfo *ReturnInfo = Types.resolve(Proto->getReturnType());
+    Sig.ReturnKey = ReturnInfo
+        ? (ReturnInfo->isC() ? "C:" + ReturnInfo->CIdentity : ReturnInfo->Name)
+        : Proto->getReturnType();
     FunctionType *FT = FunctionType::get(RetType, ArgTypes, false);
     auto &Overloads = Signatures[Sig.Name];
     for (const FuncSig &Existing : Overloads) {
@@ -241,7 +245,8 @@ Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
                         Proto->getName().c_str(), Proto->getName().c_str());
             return nullptr;
         }
-        if (sameIdentifier(Existing.ParamKey, Sig.ParamKey)) {
+        if (sameIdentifier(Existing.ParamKey, Sig.ParamKey) &&
+            sameIdentifier(Existing.ReturnKey, Sig.ReturnKey)) {
             if (Existing.OwnerUnit != CurrentUnit)
                 reportError("Duplicate definition %s in %s and %s", Sig.Name.c_str(),
                             UnitPaths[Existing.OwnerUnit].c_str(), UnitPaths[CurrentUnit].c_str());
@@ -250,6 +255,12 @@ Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
         }
     }
     Sig.LinkName = Sig.Name + (Sig.IsOverload ? "(" + Sig.ParamKey + ")" : "");
+    // Only the return type distinguishes non-void overloads that share a
+    // parameter list, so it enters their link name. PROCEDURE, NEW and
+    // OPERATOR<- keep the established names.
+    if (Sig.IsOverload && ReturnInfo && !ReturnInfo->isVoid() &&
+        !sameIdentifier(Proto->getName(), "NEW") && Proto->getOperator() != tok_assign)
+        Sig.LinkName += ":" + Sig.ReturnKey;
     Function *F = Function::Create(FT,
         Proto->getOwnerClass().empty() && sameIdentifier(Proto->getName(), "MAIN")
             ? Function::InternalLinkage : Function::ExternalLinkage,
diff --git a/lib/CodeGen/TypeCheckGen.cc b/lib/CodeGen/TypeCheckGen.cc
index 0f94859..6e9cac2 100644
--- a/lib/CodeGen/TypeCheckGen.cc
+++ b/lib/CodeGen/TypeCheckGen.cc
@@ -14,6 +14,7 @@ using namespace cps;
 Value *CodeGen::emitCoercedExpr(ExprAST *E, const TypeInfo *Target) {
     SaveAndRestore<int> Location(DiagnosticLine, E && E->getLine() ? E->getLine() : DiagnosticLine);
     if (!E || !Target) return nullptr;
+    expectResult(E, Target);
 
     // A whole-array reference evaluates to its raw data pointer, and
     // getExprTypeInfo reports the ELEMENT type for it — both would slip
diff --git a/lib/CodeGen/ValueTypeGen.cc b/lib/CodeGen/ValueTypeGen.cc
index af32d74..dddb1b0 100644
--- a/lib/CodeGen/ValueTypeGen.cc
+++ b/lib/CodeGen/ValueTypeGen.cc
@@ -8,6 +8,19 @@ using namespace cps;
 
 using namespace cps::detail;
 
+// The first context registered for a node wins: it is the outermost place the
+// result is consumed, and later queries must agree with that selection.
+void CodeGen::expectResult(const ExprAST *E, const TypeInfo *Type, bool Exact) {
+    if (!E || !Type) return;
+    ResultContexts.emplace(E, ResultContext{Type, Exact});
+}
+
+const CodeGen::ResultContext *CodeGen::resultContext(const ExprAST *E) const {
+    if (!E) return nullptr;
+    auto It = ResultContexts.find(E);
+    return It == ResultContexts.end() ? nullptr : &It->second;
+}
+
 const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) {
     if (!Expr) return nullptr;
 
@@ -128,7 +141,9 @@ const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) {
 
     if (auto *Call = dynamic_cast<MemberCallExprAST*>(Expr)) {
         const TypeInfo *Receiver = getExprTypeInfo(Call->getReceiver());
-        const FuncSig *Sig = Receiver ? resolveMethod(Receiver->Name, Call->getMethod(), Call->getArgs()) : nullptr;
+        const FuncSig *Sig = Receiver
+            ? resolveMethod(Receiver->Name, Call->getMethod(), Call->getArgs(), resultContext(Expr))
+            : nullptr;
         if (!Sig || (!Sig->Prototype->isPublic() && !sameIdentifier(CurrentClass, Receiver->Name))) return nullptr;
         return resolveType(Sig->ReturnTypeName);
     }
@@ -142,7 +157,8 @@ const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) {
             }
             return resolveType(B->ReturnTypeName);
         }
-        if (const FuncSig *Sig = resolveCall(Name, Call->getArgs(), &Call->getTemplateArgs())) {
+        if (const FuncSig *Sig = resolveCall(Name, Call->getArgs(), &Call->getTemplateArgs(),
+                                            resultContext(Expr))) {
             return resolveType(Sig->ReturnTypeName);
         }