Support extension on generic type. (#4968)

1 files changed, 44 insertions, 11 deletions

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 190433e2f..5654ac7a6 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -1075,12 +1075,13 @@ namespace Slang
     DeclRef<ExtensionDecl> applyExtensionToType(
         SemanticsVisitor*       semantics,
         ExtensionDecl*          extDecl,
-        Type*  type)
+        Type*  type,
+        Dictionary<Type*, SubtypeWitness*>* additionalSubtypeWitness)
     {
         if(!semantics)
             return DeclRef<ExtensionDecl>();
 
-        return semantics->applyExtensionToType(extDecl, type);
+        return semantics->applyExtensionToType(extDecl, type, additionalSubtypeWitness);
     }
 
     bool SemanticsVisitor::isDeclUsableAsStaticMember(
@@ -6320,7 +6321,7 @@ namespace Slang
             {
                 // Force add IDefaultInitializable to any struct missing (transitively) `IDefaultInitializable`.
                 auto* defaultInitializableType = m_astBuilder->getDefaultInitializableType();
-                if(!isSubtype(DeclRefType::create(m_astBuilder, decl), defaultInitializableType, IsSubTypeOptions::NotReadyForLookup))
+                if(!isSubtype(DeclRefType::create(m_astBuilder, decl), defaultInitializableType, IsSubTypeOptions::NoCaching))
                 {
                     InheritanceDecl* conformanceDecl = m_astBuilder->create<InheritanceDecl>();
                     conformanceDecl->parentDecl = decl;
@@ -8253,7 +8254,31 @@ namespace Slang
 
                 return;
             }
+            else if (auto genericTypeParamDecl = targetDeclRefType->getDeclRef().as<GenericTypeParamDecl>())
+            {
+                // If we are extending a generic type parameter as in `extension<T:IFoo> T`,
+                // we want to register the extension with the interface type `IFoo` instead.
+                auto genericDecl = as<GenericDecl>(genericTypeParamDecl.getDecl()->parentDecl);
+                if (!genericDecl)
+                    goto error;
+                if (genericDecl != decl->parentDecl)
+                    goto error;
+                bool isTypeConstrained = false;
+                for (auto constraintDecl : genericDecl->getMembersOfType<GenericTypeConstraintDecl>())
+                {
+                    ensureDecl(constraintDecl, DeclCheckState::ReadyForReference);
+                    if (targetDeclRefType == constraintDecl->sub.type)
+                    {
+                        auto supTypeDeclRef = isDeclRefTypeOf<AggTypeDecl>(constraintDecl->sup.type);
+                        getShared()->registerCandidateExtension(supTypeDeclRef.getDecl(), decl);
+                        isTypeConstrained = true;
+                    }
+                }
+                if (isTypeConstrained)
+                    return;
+            }
         }
+    error:;
         if (!as<ErrorType>(decl->targetType.type))
         {
             getSink()->diagnose(decl->targetType.exp, Diagnostics::invalidExtensionOnType, decl->targetType);
@@ -8357,6 +8382,15 @@ namespace Slang
             //
             return DeclRefType::create(m_astBuilder, aggTypeDeclRef);
         }
+        else if (auto genTypeParam = declRef.as<GenericTypeParamDecl>())
+        {
+            // We will reach here when we are checking `extension<T> T {...}`,
+            // where inside the extension, `This` type is the target type
+            // of the extension, in this case this is a DeclRefType to
+            // a GenericTypeParamDecl.
+            //
+            return DeclRefType::create(m_astBuilder, declRef);
+        }
         else if (auto extDeclRef = declRef.as<ExtensionDecl>())
         {
             // In the body of an `extension`, the `This`
@@ -8661,7 +8695,8 @@ namespace Slang
 
     DeclRef<ExtensionDecl> SemanticsVisitor::applyExtensionToType(
         ExtensionDecl*  extDecl,
-        Type*    type)
+        Type*    type,
+        Dictionary<Type*, SubtypeWitness*>* additionalSubtypeWitnessesForType)
     {
         DeclRef<ExtensionDecl> extDeclRef = makeDeclRef(extDecl);
 
@@ -8674,6 +8709,11 @@ namespace Slang
             ConstraintSystem constraints;
             constraints.loc = extDecl->loc;
             constraints.genericDecl = extGenericDecl;
+            if (additionalSubtypeWitnessesForType)
+            {
+                constraints.subTypeForAdditionalWitnesses = type;
+                constraints.additionalSubtypeWitnesses = additionalSubtypeWitnessesForType;
+            }
 
             // Inside the body of an extension declaration, we may end up trying to apply that
             // extension to its own target type.
@@ -8681,13 +8721,6 @@ namespace Slang
             // without any additional substitutions.
             if (extDecl->targetType->equals(type))
             {
-                /*
-                auto subst = trySolveConstraintSystem(
-                    &constraints,
-                    DeclRef<Decl>(extGenericDecl, nullptr).as<GenericDecl>(),
-                    as<GenericSubstitution>(as<DeclRefType>(type)->declRef.substitutions.substitutions));
-                return DeclRef<Decl>(extDecl, subst).as<ExtensionDecl>();
-                */
                 return createDefaultSubstitutionsIfNeeded(m_astBuilder, this, extDeclRef).as<ExtensionDecl>();
             }