Disallow only resources in constant buffers in parameterblocks on metal (#6342)

* Neaten metal parameter block checking

* Disallow only resources in constant buffers in parameterblocks on metal

closes https://github.com/shader-slang/slang/issues/6200

* add unit test for metal parameterblock cbuffer

---------

Co-authored-by: Yong He <yonghe@outlook.com>

2 files changed, 78 insertions, 42 deletions

diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 4f72332ef..ce6217825 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -2560,6 +2560,12 @@ DIAGNOSTIC(
     constantBufferInParameterBlockNotAllowedOnMetal,
     "nested 'ConstantBuffer' inside a 'ParameterBlock' is not supported on Metal, use "
     "'ParameterBlock' instead.")
+DIAGNOSTIC(
+    56101,
+    Error,
+    resourceTypesInConstantBufferInParameterBlockNotAllowedOnMetal,
+    "nesting a 'ConstantBuffer' containing resource types inside a 'ParameterBlock' is not "
+    "supported on Metal, please use 'ParameterBlock' instead.")
 
 DIAGNOSTIC(57001, Warning, spirvOptFailed, "spirv-opt failed. $0")
 DIAGNOSTIC(57002, Error, unknownPatchConstantParameter, "unknown patch constant parameter '$0'.")
diff --git a/source/slang/slang-ir-check-shader-parameter-type.cpp b/source/slang/slang-ir-check-shader-parameter-type.cpp
index 6f3161110..ad61e9540 100644
--- a/source/slang/slang-ir-check-shader-parameter-type.cpp
+++ b/source/slang/slang-ir-check-shader-parameter-type.cpp
@@ -4,63 +4,93 @@
 
 namespace Slang
 {
-void checkForInvalidShaderParameterTypeForMetal(IRModule* module, DiagnosticSink* sink)
+
+template<typename P>
+auto isOrContains(P predicate, IRType* type) -> decltype(predicate(type))
 {
-    HashSet<IRInst*> workListSet;
-    List<IRInst*> workList;
-    for (auto inst : module->getGlobalInsts())
+    HashSet<IRType*> visited;
+
+    auto go = [&visited, &predicate](auto&& self, IRType* type) -> decltype(predicate(type))
     {
-        if (inst->getOp() == kIROp_ParameterBlockType)
+        // Prevent infinite recursion by tracking visited types
+        if (!visited.add(type))
+            return {};
+
+        // Check if the current type matches the predicate
+        if (auto result = predicate(type))
+            return result;
+
+        // Recursively check struct fields
+        if (auto structType = as<IRStructType>(type))
         {
-            auto type = inst->getOperand(0);
-            if (workListSet.add(type))
-                workList.add(type);
-            // Diagnose an error on `ParameterBlock<ConstantBuffer<T>>`.
-            if (type->getOp() == kIROp_ConstantBufferType)
+            for (auto field : structType->getFields())
             {
-                bool foundUseSite = false;
-                for (auto use = inst->firstUse; use; use = use->nextUse)
-                {
-                    auto user = use->getUser();
-                    if (user->sourceLoc.isValid())
-                    {
-                        sink->diagnose(
-                            user,
-                            Diagnostics::constantBufferInParameterBlockNotAllowedOnMetal);
-                        foundUseSite = true;
-                        break;
-                    }
-                }
-                if (!foundUseSite)
-                    sink->diagnose(
-                        inst,
-                        Diagnostics::constantBufferInParameterBlockNotAllowedOnMetal);
+                auto fieldType = field->getFieldType();
+                if (auto result = self(self, fieldType))
+                    return result;
             }
         }
-    }
-    // Diagnose an error any any struct fields whose type is `ConstantBuffer<T>` if the
-    // struct is used inside a `ParameterBlock`.
-    for (Index i = 0; i < workList.getCount(); i++)
+
+        return {};
+    };
+
+    return go(go, type);
+}
+
+void checkForInvalidShaderParameterTypeForMetal(IRModule* module, DiagnosticSink* sink)
+{
+    auto isConstantBufferWithResource = [](IRType* type) -> std::optional<IRType*>
     {
-        auto type = workList[i];
-        if (auto structType = as<IRStructType>(type))
+        if (type->getOp() == kIROp_ConstantBufferType)
         {
-            for (auto field : structType->getFields())
+            // Get the type inside the constant buffer
+            auto innerType = as<IRType>(type->getOperand(0));
+
+            // Check if the inner type contains any resource types
+            auto hasResource = [](IRType* t) -> std::optional<IRType*>
             {
-                auto fieldType = field->getFieldType();
-                if (fieldType->getOp() == kIROp_ConstantBufferType)
+                if (isResourceType(t))
+                    return t;
+                return {};
+            };
+
+            if (auto resourceType = isOrContains(hasResource, innerType))
+                return type; // Return the constant buffer type if it contains a resource
+        }
+        return {};
+    };
+
+    for (auto inst : module->getGlobalInsts())
+    {
+        if (inst->getOp() != kIROp_ParameterBlockType)
+            continue;
+
+        auto type = as<IRType>(inst->getOperand(0));
+        if (auto invalidCBType = isOrContains(isConstantBufferWithResource, type))
+        {
+            // Try to find a valid source location from uses
+            bool foundUseSite = false;
+            for (auto use = inst->firstUse; use; use = use->nextUse)
+            {
+                auto user = use->getUser();
+                if (user->sourceLoc.isValid())
                 {
                     sink->diagnose(
-                        field->getKey(),
-                        Diagnostics::constantBufferInParameterBlockNotAllowedOnMetal);
+                        user,
+                        Diagnostics::
+                            resourceTypesInConstantBufferInParameterBlockNotAllowedOnMetal);
+                    foundUseSite = true;
+                    break;
                 }
-                if (workListSet.add(fieldType))
-                    workList.add(fieldType);
             }
+
+            if (!foundUseSite)
+                sink->diagnose(
+                    inst,
+                    Diagnostics::resourceTypesInConstantBufferInParameterBlockNotAllowedOnMetal);
         }
     }
 }
-
 void checkForInvalidShaderParameterType(
     TargetRequest* target,
     IRModule* module,
@@ -69,4 +99,4 @@ void checkForInvalidShaderParameterType(
     if (isMetalTarget(target))
         checkForInvalidShaderParameterTypeForMetal(module, sink);
 }
-} // namespace Slang
\ No newline at end of file
+} // namespace Slang