Fix crash on overloaded custom derivative function (#2782)

* Fix issue with resolving overloaded custom forward derivative methods.

* Add test

2 files changed, 26 insertions, 11 deletions

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 3ffb6c100..63c7d9741 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -6923,20 +6923,34 @@ namespace Slang
         auto ctx = visitor->withExprLocalScope(&scope);
         auto subVisitor = SemanticsVisitor(ctx);
         auto checkedFuncExpr = visitor->dispatchExpr(attr->funcExpr, ctx);
-        if (auto derivFuncDeclRef = as<DeclRefExpr>(checkedFuncExpr)->declRef)
+        if (auto declRefExpr = as<DeclRefExpr>(checkedFuncExpr))
         {
-            visitor->ensureDecl(derivFuncDeclRef, DeclCheckState::TypesFullyResolved);
-            auto invokeExpr = subVisitor.constructUncheckedInvokeExpr(checkedFuncExpr, imaginaryArguments);
-            auto resolved = subVisitor.ResolveInvoke(invokeExpr);
-            if (auto resolvedInvoke = as<InvokeExpr>(resolved))
+            visitor->ensureDecl(declRefExpr->declRef, DeclCheckState::TypesFullyResolved);
+        }
+        else if (auto overloadedExpr = as<OverloadedExpr>(checkedFuncExpr))
+        {
+            for (auto candidate : overloadedExpr->lookupResult2.items)
             {
-                if (auto calleeDeclRef = as<DeclRefExpr>(resolvedInvoke->functionExpr))
-                {
-                    attr->funcExpr = calleeDeclRef;
-                    return;
-                }
+                visitor->ensureDecl(candidate.declRef, DeclCheckState::TypesFullyResolved);
             }
         }
+        else
+        {   
+            visitor->getSink()->diagnose(attr, Diagnostics::cannotResolveDerivativeFunction);
+            return;
+        }
+        
+        auto invokeExpr = subVisitor.constructUncheckedInvokeExpr(checkedFuncExpr, imaginaryArguments);
+        auto resolved = subVisitor.ResolveInvoke(invokeExpr);
+        if (auto resolvedInvoke = as<InvokeExpr>(resolved))
+        {
+            if (auto calleeDeclRef = as<DeclRefExpr>(resolvedInvoke->functionExpr))
+            {
+                attr->funcExpr = calleeDeclRef;
+                return;
+            }
+        }
+        
         visitor->getSink()->diagnose(attr, Diagnostics::invalidCustomDerivative);
     }
 
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 4a9b83c6c..128142d84 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -356,8 +356,9 @@ DIAGNOSTIC(31143, Error, missingOriginalDefintionOfExternDecl, "no original defi
 DIAGNOSTIC(31145, Error, invalidCustomDerivative, "invalid custom derivative attribute.")
 DIAGNOSTIC(31146, Error, declAlreadyHasAttribute, "'$0' already has attribute '[$1]'.")
 DIAGNOSTIC(31147, Error, cannotResolveOriginalFunctionForDerivative, "cannot resolve the original function for the the custom derivative.")
+DIAGNOSTIC(31148, Error, cannotResolveDerivativeFunction, "cannot resolve the custom derivative function")
 
-DIAGNOSTIC(31148, Error, differentiableGenericInterfaceMethodNotSupported, "`[ForwardDifferentiable] and [BackwardDifferentiable] are not supported on generic interface requirements.")
+DIAGNOSTIC(31149, Error, differentiableGenericInterfaceMethodNotSupported, "`[ForwardDifferentiable] and [BackwardDifferentiable] are not supported on generic interface requirements.")
 
 // Enums