Fix bug in legalizeFuncType that leads to invalid IR. (#2902)

* Fix bug in legalizeFuncType that leads to invalid IR.

* Diagnose on functions that never returns when differentiate it.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 7 insertions, 2 deletions

diff --git a/source/slang/slang-ir-single-return.cpp b/source/slang/slang-ir-single-return.cpp
index 10d8972bc..0b61e5065 100644
--- a/source/slang/slang-ir-single-return.cpp
+++ b/source/slang/slang-ir-single-return.cpp
@@ -89,7 +89,7 @@ void convertFuncToSingleReturnForm(IRModule* irModule, IRGlobalValueWithCode* fu
     context.processFunc(func);
 }
 
-bool isSingleReturnFunc(IRGlobalValueWithCode* func)
+int getReturnCount(IRGlobalValueWithCode* func)
 {
     int returnCount = 0;
     for (auto block : func->getBlocks())
@@ -102,7 +102,12 @@ bool isSingleReturnFunc(IRGlobalValueWithCode* func)
             }
         }
     }
-    return returnCount <= 1;
+    return returnCount;
+}
+
+bool isSingleReturnFunc(IRGlobalValueWithCode* func)
+{
+    return getReturnCount(func) == 1;
 }
 
 } // namespace Slang