1 files changed, 80 insertions, 67 deletions

diff --git a/source/slang/slang-ir-com-interface.cpp b/source/slang/slang-ir-com-interface.cpp
index 009d2314d..899596209 100644
--- a/source/slang/slang-ir-com-interface.cpp
+++ b/source/slang/slang-ir-com-interface.cpp
@@ -7,92 +7,105 @@
 namespace Slang
 {
 
-struct ComInterfaceLoweringContext
+static bool _canReplace(IRUse* use)
 {
-    IRModule* module;
-    DiagnosticSink* diagnosticSink;
-
-    ArtifactStyle artifactStyle;
-
-    SharedIRBuilder sharedBuilder;
-
-    void replaceTypeUses(IRInst* inst, IRInst* newValue)
+    switch (use->getUser()->getOp())
     {
-        List<IRUse*> uses;
-        for (auto use = inst->firstUse; use; use = use->nextUse)
+        case kIROp_WitnessTableIDType:
+        case kIROp_WitnessTableType:
+        case kIROp_RTTIPointerType:
+        case kIROp_RTTIHandleType:
         {
-            uses.add(use);
+            // Don't replace
+            return false;
         }
-        for (auto use : uses)
+        case kIROp_ThisType:
         {
-            switch (use->getUser()->getOp())
-            {
-            case kIROp_WitnessTableIDType:
-            case kIROp_WitnessTableType:
-            case kIROp_ThisType:
-            case kIROp_RTTIPointerType:
-            case kIROp_RTTIHandleType:
-            case kIROp_ComPtrType:
-            case kIROp_PtrType:
-                continue;
-            default:
-                break;
-            }
-            use->set(newValue);
+            // Appears replacable.
+            break;
         }
+        case kIROp_ComPtrType:
+        case kIROp_PtrType:
+        {
+            // We can have ** and ComPtr<T>*.
+            // If it's a pointer type it could be because it is a global.
+            break;
+        }
+        default: break;
     }
+    return true;
+}
 
-    IRType* processInterfaceType(IRInterfaceType* type)
+void lowerComInterfaces(IRModule* module, ArtifactStyle artifactStyle, DiagnosticSink* sink)
+{
+    SLANG_UNUSED(sink);
+
+    // Find all of the COM interfaces
+    List<IRInterfaceType*> comInterfaces;
+    for (auto child : module->getGlobalInsts())
     {
-        if (!type->findDecoration<IRComInterfaceDecoration>())
-            return nullptr;
-     
+        auto intf = as<IRInterfaceType>(child);
+        if (intf && intf->findDecoration<IRComInterfaceDecoration>())
+        {
+            comInterfaces.add(intf);
+        }
+    }
+
+    // For all interfaces found replace uses
+    {
+        SharedIRBuilder sharedBuilder;
+        sharedBuilder.init(module);
+
         IRBuilder builder(sharedBuilder);
         builder.setInsertInto(module->getModuleInst());
 
-        IRType* result = (artifactStyle == ArtifactStyle::Kernel) ?
-            static_cast<IRType*>(builder.getPtrType(type)) :
-            static_cast<IRType*>(builder.getComPtrType(type));
+        List<IRUse*> uses;
 
-        replaceTypeUses(type, result);
-        return result;
-    }
+        for (auto comIntf : comInterfaces)
+        {
+            uses.clear();
 
-    void processThisType(IRThisType* type)
-    {
-        auto comPtrType = processInterfaceType(as<IRInterfaceType>(type->getConstraintType()));
-        if (!comPtrType)
-            return;
-        replaceTypeUses(type, comPtrType);
-    }
+            // Find all of the uses *before* doing any replacement
+            // Otherwise we end up replacing the replacement leading 
+            // to it pointing to itself.
+            for (auto use = comIntf->firstUse; use; use = use->nextUse)
+            {
+                // Only store off uses where replacement can be made
+                if (_canReplace(use))
+                {
+                    uses.add(use);
+                }
+            }
 
-    void processModule()
-    {
-        for (auto child : module->getGlobalInsts())
-        {
-            switch (child->getOp())
+            // If there are no uses that can be replaced, then we don't need
+            // to create a replacement result
+            if (uses.getCount() <= 0)
             {
-            case kIROp_InterfaceType:
-                processInterfaceType(as<IRInterfaceType>(child));
-                break;
-            case kIROp_ThisType:
-                processThisType(as<IRThisType>(child));
-                break;
-            default:
-                break;
+                continue;
+            }
+
+            // NOTE! The following code relies on the fact that the builder
+            // *doesn't* dedup in general, and in particular doesn't ptr types. 
+            // This allows the creation a 'new'  pointer type, and subsequent replacment all old uses, 
+            // leading to a `IInterface*` becoming `IInterface**`. 
+            // 
+
+            // TODO(JS): This is a temporary fix, in that whether kernel or not 
+            // shouldn't control the ptr type in general
+            // It's necessary here though because Kernel doesn't have ComPtr<>
+            // so has to be a raw pointer
+            IRType* result = (artifactStyle == ArtifactStyle::Host) ?
+                static_cast<IRType*>(builder.getComPtrType(comIntf)) :
+                static_cast<IRType*>(builder.getPtrType(comIntf));
+
+            // Go through replacing all of the replacable uses
+            for (auto use : uses)
+            {
+                // Do the replacement
+                use->set(result);
             }
         }
     }
-};
-
-void lowerComInterfaces(IRModule* module, ArtifactStyle artifactStyle, DiagnosticSink* sink)
-{
-    ComInterfaceLoweringContext context;
-    context.module = module;
-    context.diagnosticSink = sink;
-    context.artifactStyle = artifactStyle;
-    context.sharedBuilder.init(module);
-    return context.processModule();
 }
 
 }