Merge pull request #365 from csyonghe/extension

Allow extension of a concrete type to implement additional interface

4 files changed, 107 insertions, 33 deletions

diff --git a/source/slang/check.cpp b/source/slang/check.cpp
index 6c484b493..3103a7908 100644
--- a/source/slang/check.cpp
+++ b/source/slang/check.cpp
@@ -1076,8 +1076,12 @@ namespace Slang
                 {
                     // The user is asking for us to actually perform the conversion,
                     // so we need to generate an appropriate expression here.
-                    
-                    throw "foo bar baz";
+
+                    // YONGH: I am confused why we are not hitting this case before
+                    //throw "foo bar baz";
+                    // YONGH: temporary work around, may need to create the actual
+                    // invocation expr to the constructor call
+                    *outToExpr = fromExpr;
                 }
 
                 return true;
@@ -1794,7 +1798,7 @@ namespace Slang
         // `requiredMemberDeclRef` is a required member of
         // the interface.
         RefPtr<Decl> findWitnessForInterfaceRequirement(
-            DeclRef<AggTypeDecl>    typeDeclRef,
+            DeclRef<AggTypeDeclBase>    typeDeclRef,
             InheritanceDecl*        inheritanceDecl,
             DeclRef<InterfaceDecl>  interfaceDeclRef,
             DeclRef<Decl>           requiredMemberDeclRef,
@@ -1833,11 +1837,9 @@ namespace Slang
 
             // Make sure that by-name lookup is possible.
             buildMemberDictionary(typeDeclRef.getDecl());
-
-            Decl* firstMemberOfName = nullptr;
-            typeDeclRef.getDecl()->memberDictionary.TryGetValue(name, firstMemberOfName);
-
-            if (!firstMemberOfName)
+            auto lookupResult = lookUpLocal(getSession(), this, name, typeDeclRef);
+            
+            if (!lookupResult.isValid())
             {
                 getSink()->diagnose(inheritanceDecl, Diagnostics::typeDoesntImplementInterfaceRequirement, typeDeclRef, requiredMemberDeclRef);
                 return nullptr;
@@ -1845,12 +1847,12 @@ namespace Slang
 
             // Iterate over the members and look for one that matches
             // the expected signature for the requirement.
-            for (auto memberDecl = firstMemberOfName; memberDecl; memberDecl = memberDecl->nextInContainerWithSameName)
+            for (auto member : lookupResult)
             {
-                if (doesMemberSatisfyRequirement(DeclRef<Decl>(memberDecl, typeDeclRef.substitutions), requiredMemberDeclRef, requirementWitness))
-                    return memberDecl;
+                if (doesMemberSatisfyRequirement(member.declRef, requiredMemberDeclRef, requirementWitness))
+                    return member.declRef.getDecl();
             }
-
+            
             // No suitable member found, although there were candidates.
             //
             // TODO: Eventually we might want something akin to the current
@@ -1867,7 +1869,7 @@ namespace Slang
         // (via the given `inheritanceDecl`) actually provides
         // members to satisfy all the requirements in the interface.
         bool checkInterfaceConformance(
-            DeclRef<AggTypeDecl>        typeDeclRef,
+            DeclRef<AggTypeDeclBase>        typeDeclRef,
             InheritanceDecl*    inheritanceDecl,
             DeclRef<InterfaceDecl>  interfaceDeclRef)
         {
@@ -1925,7 +1927,7 @@ namespace Slang
         }
 
         bool checkConformanceToType(
-            DeclRef<AggTypeDecl> typeDeclRef,
+            DeclRef<AggTypeDeclBase> typeDeclRef,
             InheritanceDecl*     inheritanceDecl,
             Type*                baseType)
         {
@@ -1953,7 +1955,7 @@ namespace Slang
         // `inheritanceDecl` actually does what it needs to
         // for that inheritance to be valid.
         bool checkConformance(
-            DeclRef<AggTypeDecl>        typeDecl,
+            DeclRef<AggTypeDeclBase>        typeDecl,
             InheritanceDecl*            inheritanceDecl)
         {
             // Look at the type being inherited from, and validate
@@ -1963,10 +1965,10 @@ namespace Slang
         }
 
         bool checkConformance(
-            AggTypeDecl*        typeDecl,
+            AggTypeDeclBase*        typeDecl,
             InheritanceDecl*    inheritanceDecl)
         {
-            return checkConformance(DeclRef<AggTypeDecl>(typeDecl, SubstitutionSet()), inheritanceDecl);
+            return checkConformance(DeclRef<AggTypeDeclBase>(typeDecl, SubstitutionSet()), inheritanceDecl);
         }
 
         void visitAggTypeDecl(AggTypeDecl* decl)
@@ -3479,10 +3481,11 @@ namespace Slang
 
             // TODO: need to check that the target type names a declaration...
 
+            DeclRef<AggTypeDecl> aggTypeDeclRef;
             if (auto targetDeclRefType = decl->targetType->As<DeclRefType>())
             {
                 // Attach our extension to that type as a candidate...
-                if (auto aggTypeDeclRef = targetDeclRefType->declRef.As<AggTypeDecl>())
+                if (aggTypeDeclRef = targetDeclRefType->declRef.As<AggTypeDecl>())
                 {
                     auto aggTypeDecl = aggTypeDeclRef.getDecl();
                     decl->nextCandidateExtension = aggTypeDecl->candidateExtensions;
@@ -3516,6 +3519,14 @@ namespace Slang
                 EnsureDecl(m);
             }
 
+            if (aggTypeDeclRef)
+            {
+                for (auto inheritanceDecl : decl->getMembersOfType<InheritanceDecl>())
+                {
+                    checkConformance(aggTypeDeclRef.getDecl(), inheritanceDecl);
+                }
+            }
+            
             decl->SetCheckState(DeclCheckState::Checked);
         }
 
@@ -3802,7 +3813,7 @@ namespace Slang
 
                 if( auto aggTypeDeclRef = declRef.As<AggTypeDecl>() )
                 {
-                    for( auto inheritanceDeclRef : getMembersOfType<InheritanceDecl>(aggTypeDeclRef))
+                    for( auto inheritanceDeclRef : getMembersOfTypeWithExt<InheritanceDecl>(aggTypeDeclRef))
                     {
                         EnsureDecl(inheritanceDeclRef.getDecl());
 
diff --git a/source/slang/lookup.cpp b/source/slang/lookup.cpp
index 376d9bd6d..19503d63f 100644
--- a/source/slang/lookup.cpp
+++ b/source/slang/lookup.cpp
@@ -345,7 +345,19 @@ void DoLookupImpl(
 
             // Now perform "local" lookup in the context of the container,
             // as if we were looking up a member directly.
-            //
+            
+            // if we are currently in an extension decl, perform local lookup
+            // in the target decl we are extending
+            if (auto extDeclRef = containerDeclRef.As<ExtensionDecl>())
+            {
+                if (auto targetDeclRef = extDeclRef.getDecl()->targetType->AsDeclRefType())
+                {
+                    if (auto aggDeclRef = targetDeclRef->declRef.As<AggTypeDecl>())
+                    {
+                        containerDeclRef = extDeclRef.Substitute(aggDeclRef);
+                    }
+                }
+            }
             DoLocalLookupImpl(
                 session,
                 name, containerDeclRef, request, result, breadcrumbs);
diff --git a/source/slang/lower-to-ir.cpp b/source/slang/lower-to-ir.cpp
index 61ca53278..498783f4b 100644
--- a/source/slang/lower-to-ir.cpp
+++ b/source/slang/lower-to-ir.cpp
@@ -2788,19 +2788,30 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         // TODO: if this inheritance declaration is under an extension,
         // then we should construct the type that is being extended,
         // and not a reference to the extension itself.
-        auto parentDecl = inheritanceDecl->ParentDecl;
-        RefPtr<Type> type = DeclRefType::Create(
-            context->getSession(),
-            makeDeclRef(parentDecl));
 
+        auto parentDecl = inheritanceDecl->ParentDecl;
+        RefPtr<Type> type;
+        if (auto extParentDecl = dynamic_cast<ExtensionDecl*>(parentDecl))
+        {
+            type = extParentDecl->targetType.type;
+            if (auto declRefType = type.As<DeclRefType>())
+            {
+                if (auto aggTypeDecl = declRefType->declRef.As<AggTypeDecl>())
+                    parentDecl = aggTypeDecl.getDecl();
+            }
+        }
+        else
+        {
+            type = DeclRefType::Create(
+                context->getSession(),
+                makeDeclRef(parentDecl));
+        }
         // What is the super-type that we have declared we inherit from?
         RefPtr<Type> superType = inheritanceDecl->base.type;
 
         // Construct the mangled name for the witness table, which depends
         // on the type that is conforming, and the type that it conforms to.
-        String mangledName = getMangledNameForConformanceWitness(
-            makeDeclRef(parentDecl),
-            superType);
+        String mangledName = getMangledNameForConformanceWitness(type, superType);
 
         // Build an IR level witness table, which will represent the
         // conformance of the type to its super-type.
diff --git a/source/slang/syntax.h b/source/slang/syntax.h
index 375eb5f1c..ab26b1f6d 100644
--- a/source/slang/syntax.h
+++ b/source/slang/syntax.h
@@ -978,6 +978,30 @@ namespace Slang
         {
             return items.Count() > 1 ? items[0].declRef.GetName() : item.declRef.GetName();
         }
+        LookupResultItem* begin()
+        {
+            if (isValid())
+            {
+                if (isOverloaded())
+                    return items.begin();
+                else
+                    return &item;
+            }
+            else
+                return nullptr;
+        }
+        LookupResultItem* end()
+        {
+            if (isValid())
+            {
+                if (isOverloaded())
+                    return items.end();
+                else
+                    return &item + 1;
+            }
+            else
+                return nullptr;
+        }
     };
 
     struct SemanticsVisitor;
@@ -1085,6 +1109,27 @@ namespace Slang
         return FilteredMemberRefList<T>(declRef.getDecl()->Members, declRef.substitutions);
     }
 
+    inline ExtensionDecl* GetCandidateExtensions(DeclRef<AggTypeDecl> const& declRef)
+    {
+        return declRef.getDecl()->candidateExtensions;
+    }
+
+    template<typename T>
+    inline FilteredMemberRefList<T> getMembersOfTypeWithExt(DeclRef<ContainerDecl> const& declRef)
+    {
+        auto rs = getMembersOfType<T>(declRef);
+        if (auto aggDeclRef = declRef.As<AggTypeDecl>())
+        {
+            for (auto ext = GetCandidateExtensions(aggDeclRef); ext; ext = ext->nextCandidateExtension)
+            {
+                auto extMembers = getMembersOfType<T>(DeclRef<ContainerDecl>(ext, declRef.substitutions));
+                const_cast<List<RefPtr<Decl>>&>(rs.decls).AddRange(extMembers.decls);
+            }
+        }
+        return rs;
+    }
+
+
     inline RefPtr<Type> GetType(DeclRef<VarDeclBase> const& declRef)
     {
         return declRef.Substitute(declRef.getDecl()->type.Ptr());
@@ -1099,12 +1144,7 @@ namespace Slang
     {
         return declRef.Substitute(declRef.getDecl()->targetType.Ptr());
     }
-
-    inline ExtensionDecl* GetCandidateExtensions(DeclRef<AggTypeDecl> const& declRef)
-    {
-        return declRef.getDecl()->candidateExtensions;
-    }
-
+    
     inline FilteredMemberRefList<StructField> GetFields(DeclRef<StructDecl> const& declRef)
     {
         return getMembersOfType<StructField>(declRef);