1 files changed, 17 insertions, 3 deletions
diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index dd7107b18..085c3d933 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -179,14 +179,27 @@ static LegalVal maybeMaterializeWrappedValue(IRTypeLegalizationContext* context,
 static LegalVal legalizeOperand(IRTypeLegalizationContext* context, IRInst* irValue)
 {
     LegalVal legalVal;
+
+    // Special handling for type operands
+    if (auto oldType = as<IRType>(irValue))
+    {
+        // e.g. ParameterBlock<Struct>, the inst. ParameterBlockType holds the operand `StructType`,
+        // if we don't legalize it here and the same structType is legalized somewhere else, the
+        // operand of ParameterBlockType might not get updated, and it would result in a type
+        // mismatch.
+        auto legalType = legalizeType(context, oldType);
+        if (legalType.flavor == LegalType::Flavor::simple)
+            return LegalVal::simple(legalType.getSimple());
+        // legalType is not simple, fallback to the original value
+    }
+
     if (context->mapValToLegalVal.tryGetValue(irValue, legalVal))
     {
         return maybeMaterializeWrappedValue(context, legalVal);
     }
 
     // For now, assume that anything not covered
-    // by the mapping is legal as-is.
-
+    // by the type legalization or val mapping is legal as-is.
     return LegalVal::simple(irValue);
 }
 
@@ -2386,7 +2399,8 @@ static LegalVal legalizeInst(IRTypeLegalizationContext* context, IRInst* inst)
     for (UInt aa = 0; aa < argCount; ++aa)
     {
         auto oldArg = inst->getOperand(aa);
-        auto legalArg = legalizeOperand(context, oldArg);
+
+        LegalVal legalArg = legalizeOperand(context, oldArg);
         legalArgs.add(legalArg);
 
         if (legalArg.flavor != LegalVal::Flavor::simple)