feat(operator): select operator overloads by the expected result type

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

4 files changed, 78 insertions, 22 deletions

diff --git a/lib/CodeGen/ControlFlowGen.cc b/lib/CodeGen/ControlFlowGen.cc
index a351162..fe9a954 100644
--- a/lib/CodeGen/ControlFlowGen.cc
+++ b/lib/CodeGen/ControlFlowGen.cc
@@ -229,7 +229,12 @@ void CodeGen::emitCaseStmt(CaseStmtAST *Stmt) {
             }
             if (containsString(LabelType) && (LabelConverted || !isObjectLValue(Label)))
                 materializeOperatorReceiver(LabelValue, LabelType);
-            auto Overload = resolveOperator(Op, LeftType, LabelType);
+            ResultContext BoolContext{resolveType("BOOLEAN"), true};
+            auto Overload = resolveOperator(Op, LeftType, LabelType, &BoolContext);
+            if (Overload.WrongResult) {
+                reportError("CASE comparison must return BOOLEAN");
+                return nullptr;
+            }
             if (Overload.Signature && !Overload.Ambiguous) {
                 const TypeInfo *Result = resolveType(Overload.Signature->ReturnTypeName);
                 if (!Result || Result->Kind != TypeKind::Boolean) {
@@ -237,7 +242,8 @@ void CodeGen::emitCaseStmt(CaseStmtAST *Stmt) {
                     return nullptr;
                 }
             }
-            Value *Result = emitBinaryOp(Op, LeftValue, LeftType, LabelValue, LabelType, Branch.Line);
+            Value *Result = emitBinaryOp(Op, LeftValue, LeftType, LabelValue, LabelType, Branch.Line,
+                                         &BoolContext);
             cleanupTemporaries(Temporaries);
             return Result;
         };
diff --git a/lib/CodeGen/ExpressionGen.cc b/lib/CodeGen/ExpressionGen.cc
index 38155e4..18c934d 100644
--- a/lib/CodeGen/ExpressionGen.cc
+++ b/lib/CodeGen/ExpressionGen.cc
@@ -75,11 +75,19 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
             reportError("A whole array cannot be used as an operand; index it");
             return nullptr;
         }
-        auto Overload = resolveOperator(Unary->getOp(), getExprTypeInfo(Unary->getOperand()));
+        const ResultContext *Context = resultContext(Expr);
+        const TypeInfo *OperandType = getExprTypeInfo(Unary->getOperand());
+        auto Overload = resolveOperator(Unary->getOp(), OperandType, nullptr, Context);
         if (Overload.Ambiguous) {
             reportError("Ambiguous overload for %s", operatorName(Unary->getOp()).c_str());
             return nullptr;
         }
+        if (Overload.WrongResult) {
+            reportError("No matching overload for %s(%s) returning %s",
+                        operatorName(Unary->getOp()).c_str(),
+                        OperandType ? OperandType->Name.c_str() : "?", Context->Type->Name.c_str());
+            return nullptr;
+        }
         if (Overload.Signature) return emitOperatorCall(*Overload.Signature, Unary->getOperand());
         if (Unary->getOp() == tok_not) {
             const TypeInfo *Info = getExprTypeInfo(Unary->getOperand());
@@ -135,13 +143,20 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
         // '&' path: their admissible operations are decided by name here.
         const TypeInfo *LT = getExprTypeInfo(Bin->getLHS());
         const TypeInfo *RT = getExprTypeInfo(Bin->getRHS());
+        const ResultContext *Context = resultContext(Expr);
         if (LT && RT) {
-            auto Overload = resolveOperator(Bin->getOp(), LT, RT);
+            auto Overload = resolveOperator(Bin->getOp(), LT, RT, Context);
             if (Overload.Ambiguous) {
                 reportError("Ambiguous overload for %s(%s, %s)", operatorName(Bin->getOp()).c_str(),
                             LT->Name.c_str(), RT->Name.c_str());
                 return nullptr;
             }
+            if (Overload.WrongResult) {
+                reportError("No matching overload for %s(%s, %s) returning %s",
+                            operatorName(Bin->getOp()).c_str(), LT->Name.c_str(),
+                            RT->Name.c_str(), Context->Type->Name.c_str());
+                return nullptr;
+            }
             if (Overload.Signature)
                 return emitOperatorCall(*Overload.Signature, Bin->getLHS(), Bin->getRHS());
         }
@@ -150,7 +165,7 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
             return emitShortCircuit(Bin, LT, RT);
         Value *L = emitExpr(Bin->getLHS());
         Value *R = emitExpr(Bin->getRHS());
-        return emitBinaryOp(Bin->getOp(), L, LT, R, RT, Bin->getLine());
+        return emitBinaryOp(Bin->getOp(), L, LT, R, RT, Bin->getLine(), Context);
     }
 
     if (auto *Call = dynamic_cast<MemberCallExprAST*>(Expr))
@@ -219,18 +234,24 @@ Value *CodeGen::emitShortCircuit(BinaryExprAST *Bin, const TypeInfo *LT, const T
 
 // Ordinary binary expressions and CASE use the same typed operations.
 Value *CodeGen::emitBinaryOp(int Op, Value *L, const TypeInfo *LT,
-                             Value *R, const TypeInfo *RT, int Line) {
+                             Value *R, const TypeInfo *RT, int Line,
+                             const ResultContext *Context) {
     SaveAndRestore<int> Location(DiagnosticLine, Line ? Line : DiagnosticLine);
     if (!L || !R) return nullptr;
     if (!LT || !RT) {
         reportError("Cannot determine operand types for %s", operatorName(Op).c_str());
         return nullptr;
     }
-    auto Overload = resolveOperator(Op, LT, RT);
+    auto Overload = resolveOperator(Op, LT, RT, Context);
     if (Overload.Ambiguous) {
         reportError("Ambiguous overload for %s(%s, %s)", operatorName(Op).c_str(), LT->Name.c_str(), RT->Name.c_str());
         return nullptr;
     }
+    if (Overload.WrongResult) {
+        reportError("No matching overload for %s(%s, %s) returning %s", operatorName(Op).c_str(),
+                    LT->Name.c_str(), RT->Name.c_str(), Context->Type->Name.c_str());
+        return nullptr;
+    }
     if (Overload.Signature) return emitOperatorValues(*Overload.Signature, L, LT, R);
     if ((LT && LT->isC()) || (RT && RT->isC()))
         return emitCBinaryOp(Op, L, LT, R, RT, Line);
diff --git a/lib/CodeGen/OperatorGen.cc b/lib/CodeGen/OperatorGen.cc
index a840408..560b167 100644
--- a/lib/CodeGen/OperatorGen.cc
+++ b/lib/CodeGen/OperatorGen.cc
@@ -8,7 +8,7 @@ using namespace llvm;
 // Probing candidates must not evaluate operands or diagnose absent overloads:
 // both static typing and emission use this same selection.
 CodeGen::OperatorResolution CodeGen::resolveOperator(
-    int Op, const TypeInfo *Left, const TypeInfo *Right) {
+    int Op, const TypeInfo *Left, const TypeInfo *Right, const ResultContext *Context) {
     OperatorResolution Result;
     if (!Left) return Result;
     auto Candidates = FuncGen->getOverloads(operatorName(Op));
@@ -25,17 +25,46 @@ CodeGen::OperatorResolution CodeGen::resolveOperator(
     };
     // Two CHARs compare as CHARs; only '&' always builds a STRING.
     bool CanConvert = Right && (Op == '&' || Left->isString() || Right->isString());
-    for (bool Lenient : {false, true}) {
-        if (Lenient && !CanConvert) break;
-        for (const FuncSig *Sig : Candidates) {
-            bool Member = !Sig->Prototype->getOwnerClass().empty();
-            if (Sig->Prototype->getOperator() != Op || Sig->Params.size() != (Right ? 2u : 1u) - Member) continue;
-            if (!Member && !Accepts(Left, Sig->Params[0].TypeName, Lenient)) continue;
-            if (Right && !Accepts(Right, Sig->Params[Member ? 0 : 1].TypeName, Lenient)) continue;
-            if (Result.Signature) Result.Ambiguous = true;
-            else Result.Signature = Sig;
+    const bool PreferResult = Context && Context->Type;
+    auto ReturnsWanted = [&](const FuncSig *Sig) {
+        const TypeInfo *Ret = resolveType(Sig->ReturnTypeName);
+        return Ret && Ret->sameType(Context->Type);
+    };
+    bool ParamMatched = false;
+    auto RunLayers = [&](bool FilterReturn) {
+        Result.Signature = nullptr;
+        Result.Ambiguous = false;
+        for (bool Lenient : {false, true}) {
+            if (Lenient && !CanConvert) break;
+            const FuncSig *Selected = nullptr;
+            int Count = 0;
+            for (const FuncSig *Sig : Candidates) {
+                bool Member = !Sig->Prototype->getOwnerClass().empty();
+                if (Sig->Prototype->getOperator() != Op ||
+                    Sig->Params.size() != (Right ? 2u : 1u) - Member) continue;
+                if (!Member && !Accepts(Left, Sig->Params[0].TypeName, Lenient)) continue;
+                if (Right && !Accepts(Right, Sig->Params[Member ? 0 : 1].TypeName, Lenient)) continue;
+                ParamMatched = true;
+                if (FilterReturn && !ReturnsWanted(Sig)) continue;
+                ++Count;
+                if (!Selected) Selected = Sig;
+            }
+            if (Selected) {
+                Result.Signature = Selected;
+                Result.Ambiguous = Count > 1;
+                break;
+            }
         }
-        if (Result.Signature) break;
+    };
+
+    if (PreferResult) {
+        RunLayers(true);
+        if (Result.Signature || Result.Ambiguous) return Result;
+    }
+    if (!Context || !Context->Exact) {
+        RunLayers(false);
+    } else if (ParamMatched) {
+        Result.WrongResult = true;
     }
     return Result;
 }
diff --git a/lib/CodeGen/ValueTypeGen.cc b/lib/CodeGen/ValueTypeGen.cc
index dddb1b0..6026821 100644
--- a/lib/CodeGen/ValueTypeGen.cc
+++ b/lib/CodeGen/ValueTypeGen.cc
@@ -86,8 +86,8 @@ const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) {
     if (auto *Unary = dynamic_cast<UnaryExprAST*>(Expr)) {
         if (isWholeArrayVar(Unary->getOperand())) return nullptr;
         const TypeInfo *Operand = getExprTypeInfo(Unary->getOperand());
-        auto Overload = resolveOperator(Unary->getOp(), Operand);
-        if (Overload.Ambiguous) return nullptr;
+        auto Overload = resolveOperator(Unary->getOp(), Operand, nullptr, resultContext(Expr));
+        if (Overload.Ambiguous || Overload.WrongResult) return nullptr;
         if (Overload.Signature) return resolveType(Overload.Signature->ReturnTypeName);
         if (Unary->getOp() == tok_not) {
             return resolveType("BOOLEAN");
@@ -102,8 +102,8 @@ const TypeInfo *CodeGen::getExprTypeInfo(ExprAST *Expr) {
         const TypeInfo *L = getExprTypeInfo(Bin->getLHS());
         const TypeInfo *R = getExprTypeInfo(Bin->getRHS());
         if (!L || !R || isWholeArrayVar(Bin->getLHS()) || isWholeArrayVar(Bin->getRHS())) return nullptr;
-        auto Overload = resolveOperator(Bin->getOp(), L, R);
-        if (Overload.Ambiguous) return nullptr;
+        auto Overload = resolveOperator(Bin->getOp(), L, R, resultContext(Expr));
+        if (Overload.Ambiguous || Overload.WrongResult) return nullptr;
         if (Overload.Signature) return resolveType(Overload.Signature->ReturnTypeName);
         switch (Bin->getOp()) {
             case tok_eq: