Fix specialization dictionaries cleanup pass (#2844)

1 files changed, 32 insertions, 13 deletions

diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index d74c8fb37..fb00a960f 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -2346,23 +2346,42 @@ bool specializeModule(
 
 void finalizeSpecialization(IRModule* module)
 {
-    for (auto inst : module->getModuleInst()->getChildren())
+    // Go through all the top-level children of module's module inst,
+    // and remove any specialization dictionary insts.
+    //
+
+    auto moduleInst = module->getModuleInst();
+    IRInst* next = nullptr;
+    for(auto inst = moduleInst->getFirstChild(); inst; inst = next)
     {
-        for (auto decor = inst->getFirstDecoration(); decor; )
+        next = inst->getNextInst();
+
+        switch(inst->getOp())
         {
-            auto next = decor->getNextDecoration();
-            switch (decor->getOp())
+        default:
+            break;
+
+        case kIROp_StructKey:
+            for (auto decor = inst->getFirstDecoration(); decor; )
             {
-            case kIROp_ExistentialFuncSpecializationDictionary:
-            case kIROp_ExistentialTypeSpecializationDictionary:
-            case kIROp_GenericSpecializationDictionary:
-            case kIROp_DispatchFuncDecoration:
-                decor->removeAndDeallocate();
-                break;
-            default:
-                break;
+                auto nextDecor = decor->getNextDecoration();
+                switch (decor->getOp())
+                {
+                case kIROp_DispatchFuncDecoration:
+                    decor->removeAndDeallocate();
+                    break;
+                default:
+                    break;
+                }
+                decor = nextDecor;
             }
-            decor = next;
+            break;
+
+        case kIROp_ExistentialFuncSpecializationDictionary:
+        case kIROp_ExistentialTypeSpecializationDictionary:
+        case kIROp_GenericSpecializationDictionary:
+            inst->removeAndDeallocate();
+            break;
         }
     }
 }