Fix bug related to findAndCheckEntrypoint. (#5241)

1 files changed, 33 insertions, 17 deletions

diff --git a/source/slang/slang-ir-link.cpp b/source/slang/slang-ir-link.cpp
index 01b1c20de..b44c0cd5e 100644
--- a/source/slang/slang-ir-link.cpp
+++ b/source/slang/slang-ir-link.cpp
@@ -908,8 +908,8 @@ static void maybeCopyLayoutInformationToParameters(
 
 IRFunc* specializeIRForEntryPoint(
     IRSpecContext*      context,
-    String const&       mangledName,
-    String const&       nameOverride)
+    EntryPoint*         entryPoint,
+    UnownedStringSlice  nameOverride)
 {
     // We start by looking up the IR symbol that
     // matches the mangled name given to the
@@ -921,6 +921,7 @@ IRFunc* specializeIRForEntryPoint(
     // not the same as the mangled name of the decl.
     //
     RefPtr<IRSpecSymbol> sym;
+    auto mangledName = entryPoint->getEntryPointMangledName(0);
     if (!context->getSymbols().tryGetValue(mangledName, sym))
     {
         String hashedName = getHashedName(mangledName.getUnownedSlice());
@@ -948,20 +949,6 @@ IRFunc* specializeIRForEntryPoint(
     //
     auto clonedVal = cloneGlobalValue(context, originalVal);
 
-    if (nameOverride.getLength())
-    {
-        if (auto entryPointDecor = clonedVal->findDecoration<IREntryPointDecoration>())
-        {
-            IRInst* operands[] = {
-                entryPointDecor->getProfileInst(),
-                context->builder->getStringValue(nameOverride.getUnownedSlice()),
-                entryPointDecor->getModuleName()};
-            context->builder->addDecoration(
-                clonedVal, IROp::kIROp_EntryPointDecoration, operands, 3);
-            entryPointDecor->removeAndDeallocate();
-        }
-    }
-
     // In the case where the user is requesting a specialization
     // of a generic entry point, we have a bit of a problem.
     //
@@ -1023,6 +1010,34 @@ IRFunc* specializeIRForEntryPoint(
         context->builder->addKeepAliveDecoration(clonedFunc);
     }
 
+    // If an IREntryPointDecoration already exist in the function,
+    // check if we need to update its name with nameOverride.
+    // If the decoration doesn't exist, create it with the desired name.
+    if (auto entryPointDecor = clonedFunc->findDecoration<IREntryPointDecoration>())
+    {
+        if (nameOverride.getLength())
+        {
+                IRInst* operands[] = {
+                    entryPointDecor->getProfileInst(),
+                    context->builder->getStringValue(nameOverride),
+                    entryPointDecor->getModuleName() };
+                context->builder->addDecoration(
+                    clonedFunc, IROp::kIROp_EntryPointDecoration, operands, 3);
+                entryPointDecor->removeAndDeallocate();
+        }
+    }
+    else
+    {
+        if (!nameOverride.getLength())
+            nameOverride = getUnownedStringSliceText(entryPoint->getName());
+        IRInst* operands[] = {
+                   context->builder->getIntValue(context->builder->getIntType(), entryPoint->getProfile().raw),
+                   context->builder->getStringValue(nameOverride),
+                   context->builder->getStringValue(UnownedStringSlice(entryPoint->getModule()->getName())) };
+        context->builder->addDecoration(
+            clonedFunc, IROp::kIROp_EntryPointDecoration, operands, 3);
+    }
+
     // We will also go on and attach layout information
     // to the function parameters, so that we have it
     // available directly on the parameters, rather
@@ -1803,7 +1818,8 @@ LinkedIR linkIR(
     {
         auto entryPointMangledName = program->getEntryPointMangledName(entryPointIndex);
         auto nameOverride = program->getEntryPointNameOverride(entryPointIndex);
-        irEntryPoints.add(specializeIRForEntryPoint(context, entryPointMangledName, nameOverride));
+        auto entryPoint = program->getEntryPoint(entryPointIndex).get();
+        irEntryPoints.add(specializeIRForEntryPoint(context, entryPoint, nameOverride.getUnownedSlice()));
     }
 
     // Layout information for global shader parameters is also required,