Support `[RequirePrelude]` attribute on types. (#3867)

10 files changed, 104 insertions, 9 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index dcb1a159b..ead81fcf5 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -2575,6 +2575,9 @@ attribute_syntax [builtin] : BuiltinAttribute;
 __attributeTarget(DeclBase)
 attribute_syntax [__requiresNVAPI] : RequiresNVAPIAttribute;
 
+__attributeTarget(AggTypeDecl)
+attribute_syntax[RequirePrelude(target, prelude:String)] : RequirePreludeAttribute;
+
 __attributeTarget(DeclBase)
 attribute_syntax [__AlwaysFoldIntoUseSiteAttribute] : AlwaysFoldIntoUseSiteAttribute;
 
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 376307260..847844e5f 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -1256,6 +1256,15 @@ class RequiresNVAPIAttribute : public Attribute
     SLANG_AST_CLASS(RequiresNVAPIAttribute)
 };
 
+    /// A `[RequirePrelude(target, "string")]` attribute indicates that the declaration being modifed
+    /// requires a textual prelude to be injected in the resulting target code.
+class RequirePreludeAttribute : public Attribute
+{
+    SLANG_AST_CLASS(RequirePreludeAttribute)
+
+    CapabilitySet capabilitySet;
+    String prelude;
+};
 
     /// A `[__AlwaysFoldIntoUseSite]` attribute indicates that the calls into the modified
     /// function should always be folded into use sites during source emit.
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index 942afd559..1587fa1b0 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -848,6 +848,35 @@ namespace Slang
             }
             requireCapAttr->capabilitySet = CapabilitySet(capabilityNames);
         }
+        else if (auto requirePreludeAttr = as<RequirePreludeAttribute>(attr))
+        {
+            if (attr->args.getCount() > 2)
+            {
+                getSink()->diagnose(attr, Diagnostics::tooManyArguments, attr->args.getCount(), 0);
+                return false;
+            }
+            else if (attr->args.getCount() < 2)
+            {
+                getSink()->diagnose(attr, Diagnostics::notEnoughArguments, attr->args.getCount(), 2);
+                return false;
+            }
+            CapabilityName capName;
+            if (!checkCapabilityName(attr->args[0], capName))
+            {
+                return false;
+            }
+            requirePreludeAttr->capabilitySet = CapabilitySet(capName);
+            if (auto stringLitExpr = as<StringLiteralExpr>(attr->args[1]))
+            {
+                requirePreludeAttr->prelude = getStringLiteralTokenValue(stringLitExpr->token);
+            }
+            else
+            {
+                getSink()->diagnose(attr->args[1], Diagnostics::expectedAStringLiteral);
+                return false;
+            }
+            return true;
+        }
         else
         {
             if(attr->args.getCount() == 0)
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index b60e0a4ac..5813819b4 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -129,6 +129,16 @@ void CLikeSourceEmitter::emitPreModuleImpl()
 // Types
 //
 
+void CLikeSourceEmitter::ensureTypePrelude(IRType* type)
+{
+    if (auto requirePreludeDecor = as<IRRequirePreludeDecoration>(findBestTargetDecoration<IRRequirePreludeDecoration>(type)))
+    {
+        auto preludeTextInst = as<IRStringLit>(requirePreludeDecor->getOperand(1));
+        if (preludeTextInst)
+            m_requiredPreludes.add(preludeTextInst);
+    }
+}
+
 void CLikeSourceEmitter::emitDeclarator(DeclaratorInfo* declarator)
 {
     if (!declarator) return;
@@ -1733,14 +1743,15 @@ void CLikeSourceEmitter::emitInstResultDecl(IRInst* inst)
     m_writer->emit(" = ");
 }
 
+template<typename T>
 IRTargetSpecificDecoration* CLikeSourceEmitter::findBestTargetDecoration(IRInst* inInst)
 {
-    return Slang::findBestTargetDecoration(inInst, getTargetCaps());
+    return Slang::findBestTargetDecoration<T>(inInst, getTargetCaps());
 }
 
 IRTargetIntrinsicDecoration* CLikeSourceEmitter::_findBestTargetIntrinsicDecoration(IRInst* inInst)
 {
-    return as<IRTargetIntrinsicDecoration>(findBestTargetDecoration(inInst));
+    return as<IRTargetIntrinsicDecoration>(findBestTargetDecoration<IRTargetSpecificDefinitionDecoration>(inInst));
 }
 
 /* static */bool CLikeSourceEmitter::isOrdinaryName(UnownedStringSlice const& name)
@@ -3492,6 +3503,8 @@ void CLikeSourceEmitter::emitFuncDecorationsImpl(IRFunc* func)
 
 void CLikeSourceEmitter::emitStruct(IRStructType* structType)
 {
+    ensureTypePrelude(structType);
+
     // If the selected `struct` type is actually an intrinsic
     // on our target, then we don't want to emit anything at all.
     if(isTargetIntrinsic(structType))
@@ -3549,6 +3562,8 @@ void CLikeSourceEmitter::emitStructDeclarationsBlock(IRStructType* structType, b
 
 void CLikeSourceEmitter::emitClass(IRClassType* classType)
 {
+    ensureTypePrelude(classType);
+
     // If the selected `class` type is actually an intrinsic
     // on our target, then we don't want to emit anything at all.
     if (isTargetIntrinsic(classType))
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index fa99ae061..c559bb135 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -240,6 +240,7 @@ public:
     // Types
     //
 
+    void ensureTypePrelude(IRType* type);
     void emitDeclarator(DeclaratorInfo* declarator);
 
     void emitType(IRType* type, const StringSliceLoc* nameLoc) { emitTypeImpl(type, nameLoc); }
@@ -309,6 +310,7 @@ public:
 
     void emitInstResultDecl(IRInst* inst);
 
+    template<typename T>
     IRTargetSpecificDecoration* findBestTargetDecoration(IRInst* inst);
     IRTargetIntrinsicDecoration* _findBestTargetIntrinsicDecoration(IRInst* inst);
 
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 254a411cd..db1200926 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -691,7 +691,9 @@ INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
     /* TargetSpecificDecoration */
         INST(TargetDecoration,              target,                 1, 0)
         INST(TargetIntrinsicDecoration,     targetIntrinsic,        2, 0)
-    INST_RANGE(TargetSpecificDecoration, TargetDecoration, TargetIntrinsicDecoration)
+        INST_RANGE(TargetSpecificDefinitionDecoration, TargetDecoration, TargetIntrinsicDecoration)
+        INST(RequirePreludeDecoration, requirePrelude, 2, 0)
+    INST_RANGE(TargetSpecificDecoration, TargetDecoration, RequirePreludeDecoration)
     INST(GLSLOuterArrayDecoration,          glslOuterArray,         1, 0)
     
     INST(InterpolationModeDecoration,       interpolationMode,      1, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 51129c28f..165c49f11 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -107,13 +107,18 @@ struct IRTargetSpecificDecoration : IRDecoration
     }
 };
 
-struct IRTargetDecoration : IRTargetSpecificDecoration
+struct IRTargetSpecificDefinitionDecoration : IRTargetSpecificDecoration
+{
+    IR_PARENT_ISA(TargetSpecificDefinitionDecoration)
+};
+
+struct IRTargetDecoration : IRTargetSpecificDefinitionDecoration
 {
     enum { kOp = kIROp_TargetDecoration };
     IR_LEAF_ISA(TargetDecoration)
 };
 
-struct IRTargetIntrinsicDecoration : IRTargetSpecificDecoration
+struct IRTargetIntrinsicDecoration : IRTargetSpecificDefinitionDecoration
 {
     enum { kOp = kIROp_TargetIntrinsicDecoration };
     IR_LEAF_ISA(TargetIntrinsicDecoration)
@@ -126,6 +131,16 @@ struct IRTargetIntrinsicDecoration : IRTargetSpecificDecoration
     }
 };
 
+struct IRRequirePreludeDecoration : IRTargetSpecificDecoration
+{
+    IR_LEAF_ISA(RequirePreludeDecoration)
+
+    UnownedStringSlice getPrelude()
+    {
+        return as<IRStringLit>(getOperand(1))->getStringSlice();
+    }
+};
+
 struct IRIntrinsicOpDecoration : IRDecoration
 {
     enum { kOp = kIROp_IntrinsicOpDecoration };
@@ -4411,6 +4426,11 @@ public:
         addDecoration(value, kIROp_RequireGLSLVersionDecoration, getIntValue(getIntType(), IRIntegerValue(version)));
     }
 
+    void addRequirePreludeDecoration(IRInst* value, const CapabilitySet& caps, UnownedStringSlice prelude)
+    {
+        addDecoration(value, kIROp_RequirePreludeDecoration, getCapabilityValue(caps), getStringValue(prelude));
+    }
+
     void addRequireSPIRVVersionDecoration(IRInst* value, const SemanticVersion& version)
     {
         SemanticVersion::IntegerType intValue = version.toInteger();
@@ -4842,10 +4862,12 @@ void markConstExpr(
 IRTargetIntrinsicDecoration* findAnyTargetIntrinsicDecoration(
         IRInst*                 val);
 
+template<typename T>
 IRTargetSpecificDecoration* findBestTargetDecoration(
         IRInst*                 val,
         CapabilitySet const&    targetCaps);
 
+template<typename T>
 IRTargetSpecificDecoration* findBestTargetDecoration(
         IRInst*         val,
         CapabilityName  targetCapabilityAtom);
@@ -4856,7 +4878,7 @@ inline IRTargetIntrinsicDecoration* findBestTargetIntrinsicDecoration(
     IRInst* inInst,
     CapabilitySet const& targetCaps)
 {
-    return as<IRTargetIntrinsicDecoration>(findBestTargetDecoration(inInst, targetCaps));
+    return as<IRTargetIntrinsicDecoration>(findBestTargetDecoration<IRTargetSpecificDefinitionDecoration>(inInst, targetCaps));
 }
 
 
diff --git a/source/slang/slang-ir-link.cpp b/source/slang/slang-ir-link.cpp
index 177cf4e9d..defca19fd 100644
--- a/source/slang/slang-ir-link.cpp
+++ b/source/slang/slang-ir-link.cpp
@@ -1031,9 +1031,9 @@ static CapabilitySet _getBestSpecializationCaps(
     if(!val->findDecoration<IRTargetDecoration>())
         return CapabilitySet::makeEmpty();
 
-    if( auto targetDecoration = findBestTargetDecoration(inVal, targetCaps) )
+    if (auto targetSpecificDecoration = findBestTargetDecoration<IRTargetSpecificDefinitionDecoration>(inVal, targetCaps))
     {
-        return targetDecoration->getTargetCaps();
+        return targetSpecificDecoration->getTargetCaps();
     }
     else
     {
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 68c6a33b3..95f2a7c75 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -8174,6 +8174,7 @@ namespace Slang
         return inst->findDecoration<IRTargetIntrinsicDecoration>();
     }
 
+    template<typename T>
     IRTargetSpecificDecoration* findBestTargetDecoration(
         IRInst*                 inInst,
         CapabilitySet const&    targetCaps)
@@ -8194,6 +8195,8 @@ namespace Slang
             auto decoration = as<IRTargetSpecificDecoration>(dd);
             if(!decoration)
                 continue;
+            if (!T::isaImpl(decoration->getOp()))
+                continue;
 
             auto decorationCaps = decoration->getTargetCaps();
             if (decorationCaps.isIncompatibleWith(targetCaps))
@@ -8224,13 +8227,19 @@ namespace Slang
         return bestDecoration;
     }
 
+    template<typename T>
     IRTargetSpecificDecoration* findBestTargetDecoration(
             IRInst*         val,
             CapabilityName  targetCapabilityAtom)
     {
-        return findBestTargetDecoration(val, CapabilitySet(targetCapabilityAtom));
+        return findBestTargetDecoration<T>(val, CapabilitySet(targetCapabilityAtom));
     }
 
+    template
+    IRTargetSpecificDecoration* findBestTargetDecoration<IRRequirePreludeDecoration>(
+        IRInst* val,
+        CapabilityName  targetCapabilityAtom);
+
     bool findTargetIntrinsicDefinition(IRInst* callee, CapabilitySet const& targetCaps, UnownedStringSlice& outDefinition)
     {
         if (auto decor = findBestTargetIntrinsicDecoration(callee, targetCaps))
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 799faec88..4372d6e7c 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -9088,6 +9088,10 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         {
             builder->addNVAPIMagicDecoration(irInst, decl->getName()->text.getUnownedSlice());
         }
+        if (const auto requirePrelude = decl->findModifier<RequirePreludeAttribute>())
+        {
+            builder->addRequirePreludeDecoration(irInst, requirePrelude->capabilitySet, requirePrelude->prelude.getUnownedSlice());
+        }
     }
 
     void addBitFieldAccessorDecorations(IRInst* irFunc, Decl* decl)