Rework differential conformance dictionary checking. (#2483)

* Rework differential conformance dictionary checking.

* Revert space changes.

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 16 insertions, 14 deletions

diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index 251849ede..ad199300a 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -899,6 +899,16 @@ namespace Slang
         return result;
     }
 
+    void SemanticsVisitor::registerDifferentiableType(DeclRefType* type, SubtypeWitness* witness)
+    {
+        SLANG_RELEASE_ASSERT(m_parentDifferentiableAttr);
+        if (witness)
+        {
+            m_parentDifferentiableAttr->m_mapTypeToIDifferentiableWitness.AddIfNotExists(type->declRef, witness);
+        }
+    }
+
+
     void SemanticsVisitor::maybeRegisterDifferentiableType(ASTBuilder* builder, Type* type)
     {
         if (!builder->isDifferentiableInterfaceAvailable())
@@ -906,6 +916,11 @@ namespace Slang
             return;
         }
 
+        if (!m_parentDifferentiableAttr)
+        {
+            return;
+        }
+
         // Check for special cases such as PtrTypeBase<T> or Array<T>
         // This could potentially be handled later by simply defining extensions
         // for Ptr<T:IDifferentiable> etc..
@@ -927,10 +942,8 @@ namespace Slang
             if (auto subtypeWitness = as<SubtypeWitness>(
                 tryGetInterfaceConformanceWitness(type, getASTBuilder()->getDifferentiableInterface())))
             {
-                auto diffTypeContext = this->getShared()->innermostDiffTypeContext();
-                diffTypeContext->registerDifferentiableType((DeclRefType*)type, subtypeWitness);
+                registerDifferentiableType((DeclRefType*)type, subtypeWitness);
             }
-
             return;
         }
     }
@@ -2007,20 +2020,9 @@ namespace Slang
 
     Expr* SemanticsExprVisitor::visitForwardDifferentiateExpr(ForwardDifferentiateExpr* expr)
     {
-        this->getShared()->getDiffTypeContext()->requireDifferentiableTypeDictionary();
-        
         // Check/Resolve inner function declaration.
         expr->baseFunction = CheckTerm(expr->baseFunction);
 
-        // Register parameter types.
-        if (auto funcType = as<FuncType>(expr->baseFunction->type.type))
-        {
-            for (UInt i = 0; i < funcType->getParamCount(); i++)
-            {
-                maybeRegisterDifferentiableType(m_astBuilder, funcType->getParamType(i));
-            }
-        }
-
         // For now we only support using higher order expr as callee in an invoke expr.
         // The actual type of the higher order function will be derived during resolve invoke.
         expr->type = m_astBuilder->getBottomType();