Combine lookupWitness lowering with specialization. (#2794)

1 files changed, 11 insertions, 3 deletions

diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index d2e042363..eb3677653 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -5,6 +5,7 @@
 #include "slang-ir-clone.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-ssa-simplification.h"
+#include "slang-ir-lower-witness-lookup.h"
 
 namespace Slang
 {
@@ -43,6 +44,7 @@ struct SpecializationContext
     // For convenience, we will keep a pointer to the module
     // we are specializing.
     IRModule* module;
+    DiagnosticSink* sink;
 
     bool changed = false;
 
@@ -932,7 +934,11 @@ struct SpecializationContext
             }
             else
             {
-                break;
+                // If we run out of specialization opportunities, consider
+                // lower lookupWitnessMethod insts into dynamic dispatch calls.
+                iterChanged = lowerWitnessLookup(module, sink);
+                if (!iterChanged || sink->getErrorCount())
+                    break;
             }
         }
 
@@ -1323,7 +1329,6 @@ struct SpecializationContext
             IRInst* curInst = localWorkList.getLast();
 
             localWorkList.removeLast();
-            processedInsts.Remove(curInst);
 
             switch (curInst->getOp())
             {
@@ -2329,10 +2334,12 @@ struct SpecializationContext
 };
 
 bool specializeModule(
-    IRModule*   module)
+    IRModule*   module,
+    DiagnosticSink* sink)
 {
     SpecializationContext context;
     context.module = module;
+    context.sink = sink;
     context.processModule();
     return context.changed;
 }
@@ -2349,6 +2356,7 @@ void finalizeSpecialization(IRModule* module)
             case kIROp_ExistentialFuncSpecializationDictionary:
             case kIROp_ExistentialTypeSpecializationDictionary:
             case kIROp_GenericSpecializationDictionary:
+            case kIROp_DispatchFuncDecoration:
                 decor->removeAndDeallocate();
                 break;
             default: