Make `DifferentialPair` able to nest. (#2477)

1 files changed, 12 insertions, 2 deletions

diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 3a59eb6c9..382f7be5e 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -3629,6 +3629,7 @@ namespace Slang
 
     IRInst* IRBuilder::findDifferentiableTypeEntry(IRInst* irType, IRInst* scope)
     {
+        IRInst* foundResult = nullptr;
         for (auto child = scope->getFirstChild(); child; child = child->getNextInst()) 
         {
             if (child->getOp() == kIROp_DifferentiableTypeDictionary)
@@ -3640,13 +3641,20 @@ namespace Slang
 
                     if (irType == entryType)
                     {
-                        return entryConformanceWitness;
+                        foundResult = entryConformanceWitness;
+                        // If the found witness table is not a trivial one (i.e. DifferentialBottom:IDifferential),
+                        // return immediately. Otherwise, continue the search to see if we can find a better one.
+                        if (auto witness = as<IRWitnessTable>(foundResult))
+                        {
+                            if (witness->getConcreteType()->getOp() != kIROp_DifferentialBottomType)
+                                return foundResult;
+                        }
                     }
                 }
             }
         }
 
-        return nullptr;
+        return foundResult;
     }
 
     IRInst* IRBuilder::findDifferentiableTypeEntry(IRInst* irType)
@@ -6282,6 +6290,8 @@ namespace Slang
         case kIROp_MakeOptionalNone:
         case kIROp_OptionalHasValue:
         case kIROp_GetOptionalValue:
+        case kIROp_MakeTuple:
+        case kIROp_GetTupleElement:
         case kIROp_Load:    // We are ignoring the possibility of loads from bad addresses, or `volatile` loads
         case kIROp_ImageSubscript:
         case kIROp_FieldExtract: