Fix spirv generation for using output stream in a function. (#3806)

* Fix spirv generation for using output stream in a function.

* polish.

2 files changed, 35 insertions, 56 deletions

diff --git a/source/slang/slang-ir-glsl-legalize.cpp b/source/slang/slang-ir-glsl-legalize.cpp
index f51ae7215..68bd4920f 100644
--- a/source/slang/slang-ir-glsl-legalize.cpp
+++ b/source/slang/slang-ir-glsl-legalize.cpp
@@ -2495,76 +2495,53 @@ void legalizeEntryPointParameterForGLSL(
             stage,
             pp);
 
-        // TODO: a GS output stream might be passed into other
+        // A GS output stream might be passed into other
         // functions, so that we should really be modifying
         // any function that has one of these in its parameter
-        // list (and in the limit we should be leagalizing any
-        // type that nests these...).
+        // list.
         //
-        // For now we will just try to deal with `Append` calls
-        // directly in this function.
-
-        for( auto bb = func->getFirstBlock(); bb; bb = bb->getNextBlock() )
-        {
-            for( auto ii = bb->getFirstInst(); ii; ii = ii->getNextInst() )
+        HashSet<IRInst*> workListSet;
+        List<IRFunc*> workList;
+        workList.add(func);
+        workListSet.add(func);
+        for (Index i = 0; i < workList.getCount(); i++)
+        {
+            auto f = workList[i];
+            for (auto bb = f->getFirstBlock(); bb; bb = bb->getNextBlock())
             {
-                // Is it a call?
-                if(ii->getOp() != kIROp_Call)
-                    continue;
-
-                // Is it calling the append operation?
-                auto callee = ii->getOperand(0);
-                for(;;)
+                for (auto ii = bb->getFirstInst(); ii; ii = ii->getNextInst())
                 {
-                    // If the instruction is `specialize(X,...)` then
-                    // we want to look at `X`, and if it is `generic { ... return R; }`
-                    // then we want to look at `R`. We handle this
-                    // iteratively here.
-                    //
-                    // TODO: This idiom seems to come up enough that we
-                    // should probably have a dedicated convenience routine
-                    // for this.
-                    //
-                    // Alternatively, we could switch the IR encoding so
-                    // that decorations are added to the generic instead of the
-                    // value it returns.
-                    //
-                    switch(callee->getOp())
-                    {
-                    case kIROp_Specialize:
-                        {
-                            callee = cast<IRSpecialize>(callee)->getOperand(0);
-                            continue;
-                        }
+                    // Is it a call?
+                    if (ii->getOp() != kIROp_Call)
+                        continue;
+
+                    // Is it calling the append operation?
+                    auto callee = getResolvedInstForDecorations(ii->getOperand(0));
+                    if (callee->getOp() != kIROp_Func)
+                        continue;
 
-                    case kIROp_Generic:
+                    if (getBuiltinFuncName(callee) != UnownedStringSlice::fromLiteral("GeometryStreamAppend"))
+                    {
+                        // If we are calling a function that takes a output stream as a parameter,
+                        // we need to add it to the work list to be processed.
+                        for (UInt a = 1; a < ii->getOperandCount(); a++)
                         {
-                            auto genericResult = findGenericReturnVal(cast<IRGeneric>(callee));
-                            if(genericResult)
+                            if (as<IRHLSLStreamOutputType>(ii->getOperand(a)->getDataType()))
                             {
-                                callee = genericResult;
-                                continue;
+                                if (workListSet.add(callee))
+                                    workList.add(as<IRFunc>(callee));
+                                break;
                             }
                         }
-
-                    default:
-                        break;
+                        continue;
                     }
-                    break;
-                }
-                if(callee->getOp() != kIROp_Func)
-                    continue;
 
-                if (getBuiltinFuncName(callee) != UnownedStringSlice::fromLiteral("GeometryStreamAppend"))
-                {
-                    continue;
-                }
+                    // Okay, we have a declaration, and we want to modify it!
 
-                // Okay, we have a declaration, and we want to modify it!
+                    builder->setInsertBefore(ii);
 
-                builder->setInsertBefore(ii);
-
-                assign(builder, globalOutputVal, ScalarizedVal::value(ii->getOperand(2)));
+                    assign(builder, globalOutputVal, ScalarizedVal::value(ii->getOperand(2)));
+                }
             }
         }
 
diff --git a/source/slang/slang-ir-specialize-resources.cpp b/source/slang/slang-ir-specialize-resources.cpp
index af7f41207..a2bf14f4d 100644
--- a/source/slang/slang-ir-specialize-resources.cpp
+++ b/source/slang/slang-ir-specialize-resources.cpp
@@ -1236,6 +1236,8 @@ bool isIllegalGLSLParameterType(IRType* type)
     }
     if (as<IRMeshOutputType>(type))
         return true;
+    if (as<IRHLSLStreamOutputType>(type))
+        return true;
     return false;
 }