Refactor code as per review comments

Author: mkeshavaNV

source/slang/slang-ir-glsl-legalize.cpp

@@ -329,10 +329,6 @@ struct GLSLLegalizationContext
 
     IRBuilder* builder;
     IRBuilder* getBuilder() { return builder; }
-
-    // For ray tracing shaders, we need to consolidate all parameters into a single structure
-    Dictionary<IRFunc*, IRInst*> rayTracingConsolidatedVars;
-    Dictionary<IRFunc*, List<IRParam*>> rayTracingProcessedParams;
 };
 
 // This examines the passed type and determines the GLSL mesh shader indices
@@ -2306,7 +2302,7 @@ IRInst* materializeValue(IRBuilder* builder, ScalarizedVal const& val)
     }
 }
 
-static void handleSingleParam(
+void handleSingleParam(
     GLSLLegalizationContext* context,
     IRFunc* func,
     IRParam* pp,
@@ -2360,7 +2356,6 @@ static void handleSingleParam(
 
 static void consolidateParameters(
     GLSLLegalizationContext* context,
-    IRFunc* func,
     List<IRParam*>& params)
 {
     auto builder = context->getBuilder();
@@ -2381,7 +2376,8 @@ static void consolidateParameters(
             valueType = inOutType->getValueType();
 
         auto key = builder->createStructKey();
-        builder->addNameHintDecoration(key, UnownedStringSlice("field"));
+        if (auto nameDecor = _param->findDecoration<IRNameHintDecoration>())
+            builder->addNameHintDecoration(key, nameDecor->getName());
         auto field = builder->createStructField(structType, key, valueType);
         field->removeFromParent();
         field->insertAtEnd(structType);
@@ -2396,9 +2392,6 @@ static void consolidateParameters(
     // Add the ray payload decoration and assign location 0.
     builder->addVulkanRayPayloadDecoration(consolidatedVar, 0);
 
-    // Store the consolidated variable for this function
-    context->rayTracingConsolidatedVars[func] = consolidatedVar;
-
     // Replace each parameter with a field in the consolidated struct
     for (Index i = 0; i < params.getCount(); ++i)
     {
@@ -2441,64 +2434,54 @@ static void consolidateParameters(
     }
 }
 
-static void handleMultipleParams(GLSLLegalizationContext* context, IRFunc* func, IRParam* pp)
+// Consolidate ray tracing parameters for an entry point function
+void consolidateRayTracingParameters(
+    GLSLLegalizationContext* context,
+    IRFunc* func)
 {
+    auto builder = context->getBuilder();
     auto firstBlock = func->getFirstBlock();
+    if (!firstBlock)
+        return;
 
-    // Now we run the consolidation step, but if we've already
-    // processed this parameter, skip it.
-    List<IRParam*>* processedParams = nullptr;
-    if (auto foundList = context->rayTracingProcessedParams.tryGetValue(func))
-    {
-        processedParams = foundList;
-        if (processedParams->contains(pp))
-            return;
-    }
-    else
-    {
-        context->rayTracingProcessedParams[func] = List<IRParam*>();
-        processedParams = &context->rayTracingProcessedParams[func];
-    }
-
-    // Collect all parameters that need to be consolidated
-    List<IRParam*> params;
-    List<IRVarLayout*> paramLayouts;
+    // Collect all out/inout parameters that need to be consolidated
+    List<IRParam*> outParams;
+    List<IRParam*> otherParams;
 
-    for (auto _param = firstBlock->getFirstParam(); _param; _param = _param->getNextParam())
+    for (auto param = firstBlock->getFirstParam(); param; param = param->getNextParam())
     {
-        auto pLayoutDecoration = _param->findDecoration<IRLayoutDecoration>();
-        SLANG_ASSERT(pLayoutDecoration);
-        auto pLayout = as<IRVarLayout>(pLayoutDecoration->getLayout());
-        SLANG_ASSERT(pLayout);
-
-        // Only include parameters that haven't been processed yet
-        auto _paramType = _param->getDataType();
-        bool needsConsolidation = (as<IROutType>(_paramType) || as<IRInOutType>(_paramType));
-        if (!processedParams->contains(_param) && needsConsolidation)
+        builder->setInsertBefore(firstBlock->getFirstOrdinaryInst());
+        if (as<IROutTypeBase>(param->getDataType()))
+        {
+            outParams.add(param);
+        }
+        else
         {
-            params.add(_param);
-            paramLayouts.add(pLayout);
-            processedParams->add(_param);
+            otherParams.add(param);
         }
     }
 
-    consolidateParameters(context, func, params);
-}
-
-void legalizeRayTracingEntryPointParameterForGLSL(
-    GLSLLegalizationContext* context,
-    IRFunc* func,
-    IRParam* pp,
-    IRVarLayout* paramLayout,
-    bool hasSingleOutOrInOutParam)
-{
-    if (hasSingleOutOrInOutParam)
+    // We don't need consolidation here.
+    if (outParams.getCount() <= 1)
     {
-        handleSingleParam(context, func, pp, paramLayout);
+        // We have one out/inout param, so add it as part of otherParams so we can
+        // just do one pass of handleSingleParam().
+        if (outParams.getCount() == 1) 
+        { 
+            otherParams.add(outParams[0]); 
+        }
+        for (auto param : otherParams)
+        {
+            auto paramLayoutDecoration = param->findDecoration<IRLayoutDecoration>();
+            SLANG_ASSERT(paramLayoutDecoration);
+            auto paramLayout = as<IRVarLayout>(paramLayoutDecoration->getLayout());
+            handleSingleParam(context, func, param, paramLayout);
+        }
         return;
     }
 
-    handleMultipleParams(context, func, pp);
+    // Consolidate the parameters
+    consolidateParameters(context, outParams);
 }
 
 static void legalizeMeshPayloadInputParam(
@@ -3344,28 +3327,6 @@ void legalizeEntryPointParameterForGLSL(
     case Stage::Intersection:
     case Stage::Miss:
     case Stage::RayGeneration:
-        {
-            // Count the number of inout or out parameters
-            int inoutOrOutParamCount = 0;
-            auto firstBlock = func->getFirstBlock();
-            for (auto _param = firstBlock->getFirstParam(); _param; _param = _param->getNextParam())
-            {
-                auto _paramType = _param->getDataType();
-                if (as<IROutType>(_paramType) || as<IRInOutType>(_paramType))
-                {
-                    inoutOrOutParamCount++;
-                }
-            }
-
-            // If we have just one inout or out param, we don't need consolidation.
-            bool hasSingleOutOrInOutParam = inoutOrOutParamCount <= 1;
-            legalizeRayTracingEntryPointParameterForGLSL(
-                context,
-                func,
-                pp,
-                paramLayout,
-                hasSingleOutOrInOutParam);
-        }
         return;
     }
 
@@ -3941,12 +3902,33 @@ void legalizeEntryPointForGLSL(
         invokePathConstantFuncInHullShader(&context, codeGenContext, scalarizedGlobalOutput);
     }
 
+    // Special handling for ray tracing shaders
+    bool isRayTracingShader = false;
+    switch (stage)
+    {
+    case Stage::AnyHit:
+    case Stage::Callable:
+    case Stage::ClosestHit:
+    case Stage::Intersection:
+    case Stage::Miss:
+    case Stage::RayGeneration:
+        isRayTracingShader = true;
+        consolidateRayTracingParameters(&context, func);
+        break;
+    default:
+        break;
+    }
+    
     // Next we will walk through any parameters of the entry-point function,
     // and turn them into global variables.
     if (auto firstBlock = func->getFirstBlock())
     {
         for (auto pp = firstBlock->getFirstParam(); pp; pp = pp->getNextParam())
         {
+            if (isRayTracingShader)
+            {
+                continue;
+            }
             // Any initialization code we insert for parameters needs
             // to be at the start of the "ordinary" instructions in the block:
             builder.setInsertBefore(firstBlock->getFirstOrdinaryInst());
@@ -3962,7 +3944,7 @@ void legalizeEntryPointForGLSL(
 
             legalizeEntryPointParameterForGLSL(&context, codeGenContext, func, pp, paramLayout);
         }
-
+        
         // At this point we should have eliminated all uses of the
         // parameters of the entry block. Also, our control-flow
         // rules mean that the entry block cannot be the target