Add [anyValueSize] attribute to interfaces and propagate that in the IR. (#1469)

Co-authored-by: Tim Foley <tfoleyNV@users.noreply.github.com>

7 files changed, 158 insertions, 31 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index a6a969bed..716a93328 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -1955,3 +1955,6 @@ attribute_syntax [sealed] : SealedAttribute;
 
 __attributeTarget(AggTypeDecl)
 attribute_syntax [open] : OpenAttribute;
+
+__attributeTarget(InterfaceDecl)
+attribute_syntax [anyValueSize(size:int)] : AnyValueSizeAttribute;
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 1f4f6b33f..c986186fd 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -875,4 +875,13 @@ class OpenAttribute : public InheritanceControlAttribute { SLANG_CLASS(OpenAttri
     /// An attribute that marks a type declaration as disallowing the type to be inherited from in other modules.
 class SealedAttribute : public InheritanceControlAttribute { SLANG_CLASS(SealedAttribute) };
 
+    /// An attribute that defines the size of `AnyValue` type to represent a polymoprhic value that conforms to
+    /// the decorated interface type.
+class AnyValueSizeAttribute : public Attribute
+{
+    SLANG_CLASS(AnyValueSizeAttribute)
+
+    int32_t size;
+};
+
 } // namespace Slang
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index ca3ca6dae..887fcdb44 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -26,6 +26,8 @@ INST(Nop, nop, 0, 0)
 
     INST(StringType, String, 0, 0)
 
+    INST(AnyValueType, AnyValueType, 1, 0)
+
     INST(RawPointerType, RawPointerType, 0, 0)
     INST(RTTIPointerType, RTTIPointerType, 1, 0)
     INST(AfterRawPointerTypeBase, AfterRawPointerTypeBase, 0, 0)
@@ -545,7 +547,9 @@ INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
 
     /* Decorations for RTTI objects */
         INST(RTTITypeSizeDecoration, RTTI_typeSize, 1, 0)
+    INST(AnyValueSizeDecoration, AnyValueSize, 1, 0)
 
+    INST(TypeConstraintDecoration, TypeConstraintDecoration, 1, 0)
     INST_RANGE(LinkageDecoration, ImportDecoration, ExportDecoration)
 
     INST(SemanticDecoration, semantic, 2, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index ed9993512..7f8e72121 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -156,6 +156,31 @@ struct IRRTTITypeSizeDecoration : IRDecoration
     }
 };
 
+/// A decoration on `IRInterfaceType` that marks the size of `AnyValue` that should
+/// be used to represent a polymorphic value of the interface.
+struct IRAnyValueSizeDecoration : IRDecoration
+{
+    enum { kOp = kIROp_AnyValueSizeDecoration };
+    IR_LEAF_ISA(AnyValueSizeDecoration)
+
+    IRIntLit* getSizeOperand() { return cast<IRIntLit>(getOperand(0)); }
+    IRIntegerValue getSize()
+    {
+        return getSizeOperand()->getValue();
+    }
+};
+
+/// A decoration on `IRParam`s that represent generic parameters,
+/// marking the interface type that the generic parameter conforms to.
+/// A generic parameter can have more than one `IRTypeConstraintDecoration`s
+struct IRTypeConstraintDecoration : IRDecoration
+{
+    enum { kOp = kIROp_TypeConstraintDecoration };
+    IR_LEAF_ISA(TypeConstraintDecoration)
+
+    IRInst* getConstraintType() { return getOperand(0); }
+};
+
 #define IR_SIMPLE_DECORATION(NAME)      \
     struct IR##NAME : IRDecoration      \
     {                                   \
@@ -1653,12 +1678,13 @@ struct IRBuilder
     IRBasicType* getBoolType();
     IRBasicType* getIntType();
     IRStringType* getStringType();
-    IRAssociatedType* getAssociatedType();
-    IRThisType* getThisType();
+    IRAssociatedType* getAssociatedType(ArrayView<IRInterfaceType*> constraintTypes);
+    IRThisType* getThisType(IRInterfaceType* interfaceType);
     IRRawPointerType* getRawPointerType();
     IRRTTIPointerType* getRTTIPointerType(IRInst* rttiPtr);
     IRRTTIType* getRTTIType();
-
+    IRAnyValueType* getAnyValueType(IRIntegerValue size);
+    IRAnyValueType* getAnyValueType(IRInst* size);
 
     IRBasicBlockType*   getBasicBlockType();
     IRWitnessTableType* getWitnessTableType(IRType* baseType);
@@ -2344,6 +2370,16 @@ struct IRBuilder
     {
         addDecoration(inst, kIROp_RTTITypeSizeDecoration, getIntValue(getIntType(), value));
     }
+
+    void addAnyValueSizeDecoration(IRInst* inst, IRIntegerValue value)
+    {
+        addDecoration(inst, kIROp_AnyValueSizeDecoration, getIntValue(getIntType(), value));
+    }
+
+    void addTypeConstraintDecoration(IRInst* inst, IRInst* constraintType)
+    {
+        addDecoration(inst, kIROp_TypeConstraintDecoration, constraintType);
+    }
 };
 
 void addHoistableInst(
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 32034fec8..fcb039d25 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -2221,14 +2221,16 @@ namespace Slang
         return (IRStringType*)getType(kIROp_StringType);
     }
 
-    IRAssociatedType* IRBuilder::getAssociatedType()
+    IRAssociatedType* IRBuilder::getAssociatedType(ArrayView<IRInterfaceType*> constraintTypes)
     {
-        return (IRAssociatedType*)getType(kIROp_AssociatedType);
+        return (IRAssociatedType*)getType(kIROp_AssociatedType,
+            constraintTypes.getCount(),
+            (IRInst**)constraintTypes.getBuffer());
     }
 
-    IRThisType* IRBuilder::getThisType()
+    IRThisType* IRBuilder::getThisType(IRInterfaceType* interfaceType)
     {
-        return (IRThisType*)getType(kIROp_ThisType);
+        return (IRThisType*)getType(kIROp_ThisType, interfaceType);
     }
 
     IRRawPointerType* IRBuilder::getRawPointerType()
@@ -2246,6 +2248,17 @@ namespace Slang
         return (IRRTTIType*)getType(kIROp_RTTIType);
     }
 
+    IRAnyValueType* IRBuilder::getAnyValueType(IRIntegerValue size)
+    {
+        return (IRAnyValueType*)getType(kIROp_AnyValueType,
+            getIntValue(getIntType(), size));
+    }
+
+    IRAnyValueType* IRBuilder::getAnyValueType(IRInst* size)
+    {
+        return (IRAnyValueType*)getType(kIROp_AnyValueType, size);
+    }
+
     IRBasicBlockType*   IRBuilder::getBasicBlockType()
     {
         return (IRBasicBlockType*)getType(kIROp_BasicBlockType);
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index d95951553..b9668aaa8 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -1252,6 +1252,15 @@ struct IRRTTIType : IRType
     IR_LEAF_ISA(RTTIType);
 };
 
+struct IRAnyValueType : IRType
+{
+    IR_LEAF_ISA(AnyValueType);
+    IRInst* getSize()
+    {
+        return getOperand(0);
+    }
+};
+
 struct IRWitnessTableType : IRType
 {
     IRInst* getConformanceType()
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index b931212fd..36825f828 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -1085,6 +1085,23 @@ IRStructKey* getInterfaceRequirementKey(
     return requirementKey;
 }
 
+void getGenericTypeConformances(IRGenContext* context, ShortList<IRType*>& supTypes, Decl* genericParamDecl)
+{
+    auto parent = genericParamDecl->parentDecl;
+    if (parent)
+    {
+        for (auto typeConstraint : parent->getMembersOfType<GenericTypeConstraintDecl>())
+        {
+            if (auto declRefType = as<DeclRefType>(typeConstraint->sub.type))
+            {
+                if (declRefType->declRef.decl == genericParamDecl)
+                {
+                    supTypes.add(lowerType(context, typeConstraint->getSup().type));
+                }
+            }
+        }
+    }
+}
 
 SubstitutionSet lowerSubstitutions(IRGenContext* context, SubstitutionSet subst);
 //
@@ -1668,7 +1685,6 @@ struct ValLoweringVisitor : ValVisitor<ValLoweringVisitor, LoweredValInfo, Lower
         return emitDeclRef(context, type->interfaceDeclRef, getBuilder()->getTypeKind());
     }
 
-
     // We do not expect to encounter the following types in ASTs that have
     // passed front-end semantic checking.
 #define UNEXPECTED_CASE(NAME) IRType* visit##NAME(NAME*) { SLANG_UNEXPECTED(#NAME); UNREACHABLE_RETURN(nullptr); }
@@ -3465,7 +3481,8 @@ struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
 
     LoweredValInfo visitAssocTypeDecl(AssocTypeDecl* decl)
     {
-        return LoweredValInfo::simple(context->irBuilder->getAssociatedType());
+        SLANG_UNIMPLEMENTED_X("associatedtype expression during code generation");
+        UNREACHABLE_RETURN(LoweredValInfo());
     }
 
     LoweredValInfo visitAssignExpr(AssignExpr* expr)
@@ -5829,6 +5846,22 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         return Slang::getInterfaceRequirementKey(context, requirementDecl);
     }
 
+    LoweredValInfo visitAssocTypeDecl(AssocTypeDecl* decl)
+    {
+        SLANG_ASSERT(decl->parentDecl != nullptr);
+        ShortList<IRInterfaceType*> constraintInterfaces;
+        for (auto inheritanceDecl : decl->getMembersOfType<InheritanceDecl>())
+        {
+            auto baseType = lowerType(context, inheritanceDecl->base.type);
+            SLANG_ASSERT(baseType && baseType->op == kIROp_InterfaceType);
+            constraintInterfaces.add((IRInterfaceType*)baseType);
+        }
+        auto assocType = context->irBuilder->getAssociatedType(
+            constraintInterfaces.getArrayView().arrayView);
+        setValue(context, decl, assocType);
+        return LoweredValInfo::simple(assocType);
+    }
+
     LoweredValInfo visitInterfaceDecl(InterfaceDecl* decl)
     {
         // The members of an interface will turn into the keys that will
@@ -5851,18 +5884,6 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         // Emit any generics that should wrap the actual type.
         auto outerGeneric = emitOuterGenerics(subContext, decl, decl);
 
-        // Setup subContext for proper lowering `ThisType`, associated types and
-        // the interface decl's self reference.
-        subContext->thisType = getBuilder()->getThisType();
-
-        for (auto member : decl->members)
-        {
-            if (as<AssocTypeDecl>(member))
-            {
-                subContext->env->mapDeclToValue[member] = getBuilder()->getAssociatedType();
-            }
-        }
-
         // First, compute the number of requirement entries that will be included in this
         // interface type.
         UInt operandCount = 0;
@@ -5884,6 +5905,17 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         // Add `irInterface` to decl mapping now to prevent cyclic lowering.
         setValue(subContext, decl, LoweredValInfo::simple(irInterface));
 
+        // Setup subContext for proper lowering `ThisType`, associated types and
+        // the interface decl's self reference.
+        auto thisType = getBuilder()->getThisType(irInterface);
+        subContext->thisType = thisType;
+
+        // Lower associated types first, so they can be referred to when lowering functions.
+        for (auto assocTypeDecl : decl->getMembersOfType<AssocTypeDecl>())
+        {
+            ensureDecl(subContext, assocTypeDecl);
+        }
+
         UInt entryIndex = 0;
 
         for (auto requirementDecl : decl->members)
@@ -5891,7 +5923,7 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             auto entry = subBuilder->createInterfaceRequirementEntry(
                 getInterfaceRequirementKey(requirementDecl),
                 nullptr);
-            IRInst* requirementVal = lowerDecl(subContext, requirementDecl).val;
+            IRInst* requirementVal = ensureDecl(subContext, requirementDecl).val;
             if (requirementVal)
             {
                 switch (requirementVal->op)
@@ -5910,7 +5942,7 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
                 default:
                     entry->setRequirementVal(requirementVal);
                     break;
-                }   
+                }
             }
             irInterface->setOperand(entryIndex, entry);
             entryIndex++;
@@ -5927,15 +5959,22 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
                     entryIndex++;
                 }
             }
-            // Add lowered requirement entry to current decl mapping to prevent
-            // the function requirements from being lowered again when we get to
-            // `ensureAllDeclsRec`.
-            setValue(context, requirementDecl, LoweredValInfo::simple(entry));
+            else
+            {
+                // Add lowered requirement entry to current decl mapping to prevent
+                // the function requirements from being lowered again when we get to
+                // `ensureAllDeclsRec`.
+                setValue(context, requirementDecl, LoweredValInfo::simple(entry));
+            }
         }
 
-        
         addNameHint(context, irInterface, decl);
         addLinkageDecoration(context, irInterface, decl);
+        if (auto anyValueSizeAttr = decl->findModifier<AnyValueSizeAttribute>())
+        {
+            subBuilder->addAnyValueSizeDecoration(irInterface, anyValueSizeAttr->size);
+        }
+
         subBuilder->setInsertInto(irInterface);
         // TODO: are there any interface members that should be
         // nested inside the interface type itself?
@@ -6004,7 +6043,7 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
 
         if (as<AssocTypeDecl>(decl))
         {
-            return LoweredValInfo::simple(getBuilder()->getAssociatedType());
+            SLANG_UNREACHABLE("associatedtype should have been handled by visitAssocTypeDecl.");
         }
 
         // Given a declaration of a type, we need to make sure
@@ -6192,10 +6231,20 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             {
                 // TODO: use a `WitnessTableKind` to represent the
                 // classifier of the parameter.
-                auto param = subBuilder->emitParam(subBuilder->getWitnessTableType(
-                    lowerType(context, constraintDecl->sup.type)));
+                auto supType = lowerType(context, constraintDecl->sup.type);
+                auto param = subBuilder->emitParam(subBuilder->getWitnessTableType(supType));
                 addNameHint(context, param, constraintDecl);
                 setValue(subContext, constraintDecl, LoweredValInfo::simple(param));
+
+                // Attach the constraint interface type as a decoration to the IRParam value
+                // representing the generic parameter, to provide downstream passes knowledge
+                // of the correspondence.
+                if (auto declRefType = as<DeclRefType>(constraintDecl->sub.type))
+                {
+                    auto typeParamDeclVal = subContext->findLoweredDecl(declRefType->declRef.decl);
+                    SLANG_ASSERT(typeParamDeclVal && typeParamDeclVal->val);
+                    subBuilder->addTypeConstraintDecoration(typeParamDeclVal->val, supType);
+                }
             }
         }
 
@@ -6856,6 +6905,10 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         {
             return ensureDecl(context, extensionDecl);
         }
+        else if (auto interfaceDecl = as<InterfaceDecl>(genDecl->inner))
+        {
+            return ensureDecl(context, interfaceDecl);
+        }
         SLANG_RELEASE_ASSERT(false);
         UNREACHABLE_RETURN(LoweredValInfo());
     }