Add support for user defined attributes.

11 files changed, 377 insertions, 17 deletions

diff --git a/source/slang/check.cpp b/source/slang/check.cpp
index 42bf396fa..e100cbd1d 100644
--- a/source/slang/check.cpp
+++ b/source/slang/check.cpp
@@ -2389,20 +2389,75 @@ namespace Slang
             // rules to keep us from seeing shadowing variable declarations.
             auto lookupResult = lookUp(getSession(), this, attributeName, scope, LookupMask::Attribute);
 
-            // If we didn't find anything, or the result was overloaded,
+            // If the result was overloaded,
             // then we aren't going to be able to extract a single decl.
-            if(!lookupResult.isValid() || lookupResult.isOverloaded())
+            if(lookupResult.isOverloaded())
                 return nullptr;
 
-            auto decl = lookupResult.item.declRef.getDecl();
-            if( auto attributeDecl = dynamic_cast<AttributeDecl*>(decl) )
+            if (lookupResult.isValid())
             {
-                return attributeDecl;
+                auto decl = lookupResult.item.declRef.getDecl();
+                if (auto attributeDecl = dynamic_cast<AttributeDecl*>(decl))
+                {
+                    return attributeDecl;
+                }
+                else
+                {
+                    return nullptr;
+                }
             }
-            else
+
+            // If we couldn't find a system attribute, try looking up as a user defined attribute
+            // A user defined attribute class is defined as a struct type with a "UserDefinedAttributeAttribute" modifier
+            lookupResult = lookUp(getSession(), this, getSession()->getNameObj(attributeName->text + "Attribute"), scope, LookupMask::type);
+            if (lookupResult.isOverloaded())
             {
-                return nullptr;
+                // see if we have already created an AttributeDecl for this attribute struct
+                for (auto alt : lookupResult.items)
+                {
+                    if (auto adecl = alt.declRef.As<AttributeDecl>())
+                        return adecl.getDecl();
+                }
             }
+            // If we still cannot find any thing, quit
+            if (!lookupResult.isValid() || lookupResult.isOverloaded())
+                return nullptr;
+            // Now construct an AttributeDecl for this user defined attribute class for future lookup
+            auto userDefAttribAttrib = lookupResult.item.declRef.decl->FindModifier<AttributeUsageAttribute>();
+            if (!userDefAttribAttrib)
+                return nullptr;
+            // create an AttributeDecl for the user defined attribute
+            auto structAttribDef = lookupResult.item.declRef.As<StructDecl>().getDecl();
+            RefPtr<AttributeDecl> attribDecl = new AttributeDecl();
+            attribDecl->nameAndLoc = structAttribDef->nameAndLoc;
+            attribDecl->loc = structAttribDef->loc;
+            attribDecl->nextInContainerWithSameName = structAttribDef->nextInContainerWithSameName;
+            // create a __attributeTarget modifier for the attribute class definition
+            RefPtr<AttributeTargetModifier> targetModifier = new AttributeTargetModifier();
+            targetModifier->syntaxClass = userDefAttribAttrib->targetSyntaxClass;
+            targetModifier->loc = structAttribDef->loc;
+            targetModifier->next = attribDecl->modifiers.first;
+            attribDecl->modifiers.first = targetModifier;
+            structAttribDef->nextInContainerWithSameName = attribDecl.Ptr();
+            // we should create UserDefinedAttribute nodes for all user defined attribute instances
+            attribDecl->syntaxClass = getSession()->findSyntaxClass(getSession()->getNameObj("UserDefinedAttribute"));
+            for (auto member : structAttribDef->Members)
+            {
+                if (auto varMember = member.As<VarDecl>())
+                {
+                    RefPtr<ParamDecl> param = new ParamDecl();
+                    param->nameAndLoc = member->nameAndLoc;
+                    param->type = varMember->type;
+                    param->loc = member->loc;
+                    attribDecl->Members.Add(param);
+                }
+            }
+            // add the attribute class definition to the syntax tree, so it can be found
+            structAttribDef->ParentDecl->Members.Add(attribDecl.Ptr());
+            structAttribDef->ParentDecl->memberDictionaryIsValid = false;
+            // do necessary checks on this newly constructed node
+            checkDecl(attribDecl.Ptr());
+            return attribDecl.Ptr();
         }
 
         bool hasIntArgs(Attribute* attr, int numArgs)
@@ -2437,7 +2492,27 @@ namespace Slang
             return true;
         }
 
-        bool validateAttribute(RefPtr<Attribute> attr)
+        bool getAttributeTargetSyntaxClasses(SyntaxClass<RefObject> & cls, uint32_t typeFlags)
+        {
+            if (typeFlags == (int)UserDefinedAttributeTargets::Struct)
+            {
+                cls = getSession()->findSyntaxClass(getSession()->getNameObj("StructDecl"));
+                return true;
+            }
+            if (typeFlags == (int)UserDefinedAttributeTargets::Var)
+            {
+                cls = getSession()->findSyntaxClass(getSession()->getNameObj("VarDecl"));
+                return true;
+            }
+            if (typeFlags == (int)UserDefinedAttributeTargets::Function)
+            {
+                cls = getSession()->findSyntaxClass(getSession()->getNameObj("FuncDecl"));
+                return true;
+            }
+            return false;
+        }
+
+        bool validateAttribute(RefPtr<Attribute> attr, AttributeDecl* attribClassDecl)
         {
                 if(auto numThreadsAttr = attr.As<NumThreadsAttribute>())
                 {
@@ -2529,6 +2604,67 @@ namespace Slang
                     // Has no args
                     SLANG_ASSERT(attr->args.Count() == 0);
                 }
+                else if (auto attrUsageAttr = attr.As<AttributeUsageAttribute>())
+                {
+                    uint32_t targetClassId = (uint32_t)UserDefinedAttributeTargets::None;
+                    if (attr->args.Count() == 1)
+                    {
+                        RefPtr<IntVal> outIntVal;
+                        if (auto cInt = checkConstantIntVal(attr->args[0]))
+                        {
+                            targetClassId = (uint32_t)(cInt->value);
+                        }
+                        else
+                        {
+                            getSink()->diagnose(attr, Diagnostics::expectedSingleIntArg, attr->name);
+                            return false;
+                        }
+                    }
+                    if (!getAttributeTargetSyntaxClasses(attrUsageAttr->targetSyntaxClass, targetClassId))
+                    {
+                        getSink()->diagnose(attr, Diagnostics::invalidAttributeTarget);
+                        return false;
+                    }
+                }
+                else if (auto userDefAttr = attr.As<UserDefinedAttribute>())
+                {
+                    // check arguments against attribute parameters defined in attribClassDecl
+                    uint32_t paramIndex = 0;
+                    auto params = attribClassDecl->getMembersOfType<ParamDecl>();
+                    for (auto paramDecl : params)
+                    {
+                        if (paramIndex < attr->args.Count())
+                        {
+                            auto & arg = attr->args[paramIndex];
+                            bool typeChecked = false;
+                            if (auto basicType = paramDecl->getType()->AsBasicType())
+                            {
+                                if (basicType->baseType == BaseType::Int)
+                                {
+                                    if (auto cint = checkConstantIntVal(arg))
+                                    {
+                                        attr->intArgVals[(uint32_t)paramIndex] = cint;
+                                    }
+                                    typeChecked = true;
+                                }
+                            }
+                            if (!typeChecked)
+                            {
+                                arg = CheckExpr(arg);
+                                arg = Coerce(paramDecl->getType(), arg);
+                            }
+                        }
+                        paramIndex++;
+                    }
+                    if (params.Count() < attr->args.Count())
+                    {
+                        getSink()->diagnose(attr, Diagnostics::tooManyArguments, attr->args.Count(), params.Count());
+                    }
+                    else if (params.Count() > attr->args.Count())
+                    {
+                        getSink()->diagnose(attr, Diagnostics::notEnoughArguments, attr->args.Count(), params.Count());
+                    }
+                }
                 else
                 {
                     if(attr->args.Count() == 0)
@@ -2597,11 +2733,9 @@ namespace Slang
                 UInt paramIndex = paramCounter++;
                 if( paramIndex < argCount )
                 {
-                    auto arg = attr->args[paramIndex];
-
                     // TODO: support checking the argument against the declared
                     // type for the parameter.
-
+                    
                 }
                 else
                 {
@@ -2653,7 +2787,7 @@ namespace Slang
             }
 
             // Now apply type-specific validation to the attribute.
-            if(!validateAttribute(attr))
+            if(!validateAttribute(attr, attrDecl))
             {
                 return uncheckedAttr;
             }
@@ -2778,6 +2912,24 @@ namespace Slang
                     registerExtension(extDecl);
                 }
             }
+            // check user defined attribute classes first
+            for (auto decl : programNode->Members)
+            {
+                if (auto typeMember = decl->As<StructDecl>())
+                {
+                    bool isTypeAttributeClass = false;
+                    for (auto attrib : typeMember->GetModifiersOfType<UncheckedAttribute>())
+                    {
+                        if (attrib->name == getSession()->getNameObj("AttributeUsageAttribute"))
+                        {
+                            isTypeAttributeClass = true;
+                            break;
+                        }
+                    }
+                    if (isTypeAttributeClass)
+                        checkDecl(decl);
+                }
+            }
             // check types
             for (auto & s : programNode->getMembersOfType<TypeDefDecl>())
                 checkDecl(s.Ptr());
@@ -8588,6 +8740,7 @@ namespace Slang
 
         RefPtr<Expr> visitStaticMemberExpr(StaticMemberExpr* /*expr*/)
         {
+            // StaticMemberExpr means it is already checked
             SLANG_UNEXPECTED("should not occur in unchecked AST");
             UNREACHABLE_RETURN(nullptr);
         }
diff --git a/source/slang/compiler.h b/source/slang/compiler.h
index ca82950be..9ca4bbc70 100644
--- a/source/slang/compiler.h
+++ b/source/slang/compiler.h
@@ -587,6 +587,7 @@ namespace Slang
         RootNamePool* getRootNamePool() { return &rootNamePool; }
         NamePool* getNamePool() { return &namePool; }
         Name* getNameObj(String name) { return namePool.getName(name); }
+        Name* tryGetNameObj(String name) { return namePool.tryGetName(name); }
         //
 
         // Generated code for stdlib, etc.
diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index e6bad3f50..50a4cbf96 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -1286,3 +1286,12 @@ attribute_syntax [mutating] : MutatingAttribute;
     /// This is equivalent to the LLVM `readnone` function attribute.
 __attributeTarget(FunctionDeclBase)
 attribute_syntax [__readNone] : ReadNoneAttribute;
+
+enum AttributeTargets
+{
+    Struct = $((int) UserDefinedAttributeTargets::Struct),
+    Var = $((int) UserDefinedAttributeTargets::Var),
+    Function = $((int) UserDefinedAttributeTargets::Function),
+};
+__attributeTarget(StructDecl)
+attribute_syntax [AttributeUsage(target : AttributeTargets)] : AttributeUsageAttribute;
\ No newline at end of file
diff --git a/source/slang/core.meta.slang.h b/source/slang/core.meta.slang.h
index a29fc24ce..746e75bd9 100644
--- a/source/slang/core.meta.slang.h
+++ b/source/slang/core.meta.slang.h
@@ -1304,3 +1304,15 @@ SLANG_RAW("    ///\n")
 SLANG_RAW("    /// This is equivalent to the LLVM `readnone` function attribute.\n")
 SLANG_RAW("__attributeTarget(FunctionDeclBase)\n")
 SLANG_RAW("attribute_syntax [__readNone] : ReadNoneAttribute;\n")
+SLANG_RAW("\n")
+SLANG_RAW("enum AttributeTargets\n")
+SLANG_RAW("{\n")
+
+    sb << "Struct = " << (int)UserDefinedAttributeTargets::Struct << ", ";
+    sb << "Var = " << (int)UserDefinedAttributeTargets::Var << ", ";
+    sb << "Function = " << (int)UserDefinedAttributeTargets::Function;
+
+SLANG_RAW("\n")
+SLANG_RAW("};\n")
+SLANG_RAW("__attributeTarget(StructDecl)\n")
+SLANG_RAW("attribute_syntax [AttributeUsage(target : AttributeTargets)] : AttributeUsageAttribute;")
diff --git a/source/slang/diagnostic-defs.h b/source/slang/diagnostic-defs.h
index 2c4f00dd6..68790d4ab 100644
--- a/source/slang/diagnostic-defs.h
+++ b/source/slang/diagnostic-defs.h
@@ -260,9 +260,8 @@ DIAGNOSTIC(31005, Error, expectedSingleStringArg, "attribute '$0' expects a sing
 DIAGNOSTIC(31006, Error, attributeFunctionNotFound, "Could not find function '$0' for attribute'$1'")
 
 
-DIAGNOSTIC(31100, Error, unknownStageName,  "unknown stage name '$0'")
-
-
+DIAGNOSTIC(31100, Error, unknownStageName, "unknown stage name '$0'")
+DIAGNOSTIC(31120, Error, invalidAttributeTarget, "invalid syntax target for user defined attribute")
 
 // Enums
 
diff --git a/source/slang/hlsl.meta.slang.h b/source/slang/hlsl.meta.slang.h
index 4dc75ab78..ef1a6d273 100644
--- a/source/slang/hlsl.meta.slang.h
+++ b/source/slang/hlsl.meta.slang.h
@@ -418,6 +418,7 @@ SLANG_RAW("\n")
 SLANG_RAW("double asdouble(uint lowbits, uint highbits);\n")
 SLANG_RAW("\n")
 SLANG_RAW("// Reinterpret bits as a float (HLSL SM 4.0)\n")
+SLANG_RAW("\n")
 SLANG_RAW("__target_intrinsic(glsl, \"intBitsToFloat\")\n")
 SLANG_RAW("float asfloat(int x);\n")
 SLANG_RAW("__target_intrinsic(glsl, \"uintBitsToFloat\")\n")
diff --git a/source/slang/modifier-defs.h b/source/slang/modifier-defs.h
index 2276f37f8..164621620 100644
--- a/source/slang/modifier-defs.h
+++ b/source/slang/modifier-defs.h
@@ -315,6 +315,14 @@ END_SYNTAX_CLASS()
 
 // A `[name(arg0, ...)]` style attribute that has been validated.
 SYNTAX_CLASS(Attribute, AttributeBase)
+    FIELD(AttributeArgumentValueDict, intArgVals)
+END_SYNTAX_CLASS()
+
+SYNTAX_CLASS(UserDefinedAttribute, Attribute)
+END_SYNTAX_CLASS()
+
+SYNTAX_CLASS(AttributeUsageAttribute, Attribute)
+    FIELD(SyntaxClass<RefObject>, targetSyntaxClass)
 END_SYNTAX_CLASS()
 
 // An `[unroll]` or `[unroll(count)]` attribute
diff --git a/source/slang/name.cpp b/source/slang/name.cpp
index a5cbb6541..21586e0b6 100644
--- a/source/slang/name.cpp
+++ b/source/slang/name.cpp
@@ -26,4 +26,12 @@ Name* NamePool::getName(String const& text)
     return name;
 }
 
+Name* NamePool::tryGetName(String const& text)
+{
+    RefPtr<Name> name;
+    if (rootPool->names.TryGetValue(text, name))
+        return name;
+    return nullptr;
+}
+
 } // namespace Slang
diff --git a/source/slang/name.h b/source/slang/name.h
index e1a055e60..a144fbb84 100644
--- a/source/slang/name.h
+++ b/source/slang/name.h
@@ -66,7 +66,9 @@ struct NamePool
 {
     // Find or create the `Name` that represents the given `text`.
     Name* getName(String const& text);
-
+    // Try find the `Name` that represents the given `text`.
+    // If the name does not exist, return nullptr
+    Name* tryGetName(String const& text);
     // Set the parent name pool to use for lookup
     void setRootNamePool(RootNamePool* rootNamePool)
     {
diff --git a/source/slang/reflection.cpp b/source/slang/reflection.cpp
index 44920fb9f..9a5a5faf9 100644
--- a/source/slang/reflection.cpp
+++ b/source/slang/reflection.cpp
@@ -3,7 +3,7 @@
 
 #include "compiler.h"
 #include "type-layout.h"
-
+#include "syntax.h"
 #include <assert.h>
 
 // Don't signal errors for stuff we don't implement here,
@@ -18,7 +18,19 @@ using namespace Slang;
 
 
 // Conversion routines to help with strongly-typed reflection API
+static inline Session* convert(SlangSession* session)
+{
+    return (Session*)session;
+}
 
+static inline UserDefinedAttribute* convert(SlangReflectionUserAttribute* attrib)
+{
+    return (UserDefinedAttribute*)attrib;
+}
+static inline SlangReflectionUserAttribute* convert(UserDefinedAttribute* attrib)
+{
+    return (SlangReflectionUserAttribute*)attrib;
+}
 static inline Type* convert(SlangReflectionType* type)
 {
     return (Type*) type;
@@ -85,6 +97,101 @@ static inline SlangReflection* convert(ProgramLayout* program)
     return (SlangReflection*) program;
 }
 
+// user attaribute
+
+unsigned int getUserAttributeCount(Decl* decl)
+{
+    unsigned int count = 0;
+    for (auto x : decl->GetModifiersOfType<UserDefinedAttribute>())
+    {
+        SLANG_UNUSED(x);
+        count++;
+    }
+    return count;
+}
+
+SlangReflectionUserAttribute* findUserAttributeByName(Session* session, Decl* decl, const char* name)
+{
+    auto nameObj = session->tryGetNameObj(name);
+    for (auto x : decl->GetModifiersOfType<UserDefinedAttribute>())
+    {
+        if (x->name == nameObj)
+            return (SlangReflectionUserAttribute*)(x);
+    }
+    return nullptr;
+}
+
+SlangReflectionUserAttribute* getUserAttributeByIndex(Decl* decl, unsigned int index)
+{
+    unsigned int id = 0;
+    for (auto x : decl->GetModifiersOfType<UserDefinedAttribute>())
+    {
+        if (id == index)
+            return convert(x);
+        id++;
+    }
+    return nullptr;
+}
+
+SLANG_API char const* spReflectionUserAttribute_GetName(SlangReflectionUserAttribute* attrib)
+{
+    auto userAttr = convert(attrib);
+    if (!userAttr) return nullptr;
+    return userAttr->getName()->text.Buffer();
+}
+SLANG_API unsigned int spReflectionUserAttribute_GetArgumentCount(SlangReflectionUserAttribute* attrib)
+{
+    auto userAttr = convert(attrib);
+    if (!userAttr) return 0;
+    return (unsigned int)userAttr->args.Count();
+}
+SlangReflectionType* spReflectionUserAttribute_GetArgumentType(SlangReflectionUserAttribute* attrib, unsigned int index)
+{
+    auto userAttr = convert(attrib);
+    if (!userAttr) return nullptr;
+    return convert(userAttr->args[index]->type.type.Ptr());
+}
+SLANG_API SlangResult spReflectionUserAttribute_GetArgumentValueInt(SlangReflectionUserAttribute* attrib, unsigned int index, int * rs)
+{
+    auto userAttr = convert(attrib);
+    if (!userAttr) return SLANG_ERROR_INVALID_PARAMETER;
+    if (index >= userAttr->args.Count()) return SLANG_ERROR_INVALID_PARAMETER;
+    RefPtr<RefObject> val;
+    if (userAttr->intArgVals.TryGetValue(index, val))
+    {
+        *rs = (int)val.As<ConstantIntVal>()->value;
+        return 0;
+    }
+    return SLANG_ERROR_INVALID_PARAMETER;
+}
+SLANG_API SlangResult spReflectionUserAttribute_GetArgumentValueFloat(SlangReflectionUserAttribute* attrib, unsigned int index, float * rs)
+{
+    auto userAttr = convert(attrib);
+    if (!userAttr) return SLANG_ERROR_INVALID_PARAMETER;
+    if (index >= userAttr->args.Count()) return SLANG_ERROR_INVALID_PARAMETER;
+    if (auto cexpr = userAttr->args[index].As<FloatingPointLiteralExpr>())
+    {
+        *rs = (float)cexpr->value;
+        return 0;
+    }
+    return SLANG_ERROR_INVALID_PARAMETER;
+}
+SLANG_API const char* spReflectionUserAttribute_GetArgumentValueString(SlangReflectionUserAttribute* attrib, unsigned int index, size_t* bufLen)
+{
+    auto userAttr = convert(attrib);
+    if (!userAttr) return nullptr;
+    if (index >= userAttr->args.Count()) return nullptr;
+    if (auto cexpr = userAttr->args[index].As<StringLiteralExpr>())
+    {
+        if (bufLen)
+            *bufLen = cexpr->token.Content.size();
+        return cexpr->token.Content.begin();
+    }
+    return nullptr;
+}
+
+
+
 // type Reflection
 
 
@@ -352,6 +459,37 @@ SLANG_API SlangScalarType spReflectionType_GetScalarType(SlangReflectionType* in
     return SLANG_SCALAR_TYPE_NONE;
 }
 
+SLANG_API unsigned int spReflectionType_GetUserAttributeCount(SlangReflectionType* inType)
+{
+    auto type = convert(inType);
+    if (!type) return 0;
+    if (auto declRefType = type->AsDeclRefType())
+    {
+        return getUserAttributeCount(declRefType->declRef.getDecl());
+    }
+    return 0;
+}
+SLANG_API SlangReflectionUserAttribute* spReflectionType_GetUserAttribute(SlangReflectionType* inType, unsigned int index)
+{
+    auto type = convert(inType);
+    if (!type) return 0;
+    if (auto declRefType = type->AsDeclRefType())
+    {
+        return getUserAttributeByIndex(declRefType->declRef.getDecl(), index);
+    }
+    return 0;
+}
+SLANG_API SlangReflectionUserAttribute* spReflectionType_FindUserAttributeByName(SlangReflectionType* inType, char const* name)
+{
+    auto type = convert(inType);
+    if (!type) return 0;
+    if (auto declRefType = type->AsDeclRefType())
+    {
+        return findUserAttributeByName(declRefType->getSession(), declRefType->declRef.getDecl(), name);
+    }
+    return 0;
+}
+
 SLANG_API SlangResourceShape spReflectionType_GetResourceShape(SlangReflectionType* inType)
 {
     auto type = convert(inType);
@@ -757,6 +895,24 @@ SLANG_API SlangReflectionModifier* spReflectionVariable_FindModifier(SlangReflec
     return (SlangReflectionModifier*) modifier;
 }
 
+SLANG_API unsigned int spReflectionVariable_GetUserAttributeCount(SlangReflectionVariable* inVar)
+{
+    auto varDecl = convert(inVar);
+    if (!varDecl) return 0;
+    return getUserAttributeCount(varDecl);
+}
+SLANG_API SlangReflectionUserAttribute* spReflectionVariable_GetUserAttribute(SlangReflectionVariable* inVar, unsigned int index)
+{
+    auto varDecl = convert(inVar);
+    if (!varDecl) return 0;
+    return getUserAttributeByIndex(varDecl, index);
+}
+SLANG_API SlangReflectionUserAttribute* spReflectionVariable_FindUserAttributeByName(SlangReflectionVariable* inVar, SlangSession* session, char const* name)
+{
+    auto varDecl = convert(inVar);
+    if (!varDecl) return 0;
+    return findUserAttributeByName(convert(session), varDecl, name);
+}
 
 // Variable Layout Reflection
 
diff --git a/source/slang/syntax.h b/source/slang/syntax.h
index 74eda66a5..5db762f11 100644
--- a/source/slang/syntax.h
+++ b/source/slang/syntax.h
@@ -1084,6 +1084,8 @@ namespace Slang
         RequirementDictionary requirementDictionary;
     };
 
+    typedef Dictionary<unsigned int, RefPtr<RefObject>> AttributeArgumentValueDict;
+
     // Generate class definition for all syntax classes
 #define SYNTAX_FIELD(TYPE, NAME) TYPE NAME;
 #define FIELD(TYPE, NAME) TYPE NAME;
@@ -1343,6 +1345,15 @@ namespace Slang
         RefPtr<Substitutions>   outerSubst);
 
     RefPtr<GenericSubstitution> findInnerMostGenericSubstitution(Substitutions* subst);
+
+    enum class UserDefinedAttributeTargets
+    {
+        None = 0,
+        Struct = 1,
+        Var = 2,
+        Function = 4,
+        All = 7
+    };
 } // namespace Slang
 
 #endif