Enable exprs for GLSL binding layout qualifiers (#5807)

* Allow glsl set and binding layout qualifiers to be compile time integer exprs

* add new tests

* add comments

* cleanup on asserts

* addressed review comments and cleanup

* fix missing set expr in test

* fixed tests and cleanup

---------

Co-authored-by: Yong He <yonghe@outlook.com>

5 files changed, 286 insertions, 83 deletions

diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 863ffaef2..6d11f0953 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -693,6 +693,38 @@ class UncheckedAttribute : public AttributeBase
     Scope* scope = nullptr;
 };
 
+// A GLSL layout qualifier whose value has not yet been resolved or validated.
+class UncheckedGLSLLayoutAttribute : public AttributeBase
+{
+    SLANG_AST_CLASS(UncheckedGLSLLayoutAttribute)
+
+    SLANG_UNREFLECTED
+};
+
+// GLSL `binding` layout qualifier, does not include `set`.
+class UncheckedGLSLBindingLayoutAttribute : public UncheckedGLSLLayoutAttribute
+{
+    SLANG_AST_CLASS(UncheckedGLSLBindingLayoutAttribute)
+
+    SLANG_UNREFLECTED
+};
+
+// GLSL `set` layout qualifier, does not include `binding`.
+class UncheckedGLSLSetLayoutAttribute : public UncheckedGLSLLayoutAttribute
+{
+    SLANG_AST_CLASS(UncheckedGLSLSetLayoutAttribute)
+
+    SLANG_UNREFLECTED
+};
+
+// GLSL `offset` layout qualifier.
+class UncheckedGLSLOffsetLayoutAttribute : public UncheckedGLSLLayoutAttribute
+{
+    SLANG_AST_CLASS(UncheckedGLSLOffsetLayoutAttribute)
+
+    SLANG_UNREFLECTED
+};
+
 // A `[name(arg0, ...)]` style attribute that has been validated.
 class Attribute : public AttributeBase
 {
@@ -854,6 +886,14 @@ class GLSLOffsetLayoutAttribute : public Attribute
     int64_t offset;
 };
 
+// Implicitly added offset qualifier when no offset is specified.
+class GLSLImplicitOffsetLayoutAttribute : public AttributeBase
+{
+    SLANG_AST_CLASS(GLSLImplicitOffsetLayoutAttribute)
+
+    SLANG_UNREFLECTED
+};
+
 class GLSLSimpleIntegerLayoutAttribute : public Attribute
 {
     SLANG_AST_CLASS(GLSLSimpleIntegerLayoutAttribute)
diff --git a/source/slang/slang-check-impl.h b/source/slang/slang-check-impl.h
index 0dfa75631..45e899132 100644
--- a/source/slang/slang-check-impl.h
+++ b/source/slang/slang-check-impl.h
@@ -614,6 +614,17 @@ struct ImplicitCastMethodKey
     }
 };
 
+/// Used to track offsets for atomic counter storage qualifiers.
+struct GLSLBindingOffsetTracker
+{
+public:
+    void setBindingOffset(int binding, int64_t byteOffset);
+    int64_t getNextBindingOffset(int binding);
+
+private:
+    Dictionary<int, int64_t> bindingToByteOffset;
+};
+
 /// Shared state for a semantics-checking session.
 struct SharedSemanticsContext : public RefObject
 {
@@ -645,6 +656,8 @@ struct SharedSemanticsContext : public RefObject
     List<ModuleDecl*> importedModulesList;
     HashSet<ModuleDecl*> importedModulesSet;
 
+    GLSLBindingOffsetTracker m_glslBindingOffsetTracker;
+
 public:
     SharedSemanticsContext(
         Linkage* linkage,
@@ -705,6 +718,8 @@ public:
         InheritanceCircularityInfo* next = nullptr;
     };
 
+    GLSLBindingOffsetTracker* getGLSLBindingOffsetTracker() { return &m_glslBindingOffsetTracker; }
+
     /// Get the processed inheritance information for `type`, including all its facets
     InheritanceInfo getInheritanceInfo(
         Type* type,
@@ -1055,6 +1070,11 @@ public:
 
     OrderedHashSet<Type*>* getCapturedTypePacks() { return m_capturedTypePacks; }
 
+    GLSLBindingOffsetTracker* getGLSLBindingOffsetTracker()
+    {
+        return m_shared->getGLSLBindingOffsetTracker();
+    }
+
 private:
     SharedSemanticsContext* m_shared = nullptr;
 
@@ -1647,6 +1667,10 @@ public:
         UncheckedAttribute* uncheckedAttr,
         ModifiableSyntaxNode* attrTarget);
 
+    AttributeBase* checkGLSLLayoutAttribute(
+        UncheckedGLSLLayoutAttribute* uncheckedAttr,
+        ModifiableSyntaxNode* attrTarget);
+
     Modifier* checkModifier(
         Modifier* m,
         ModifiableSyntaxNode* syntaxNode,
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index a1e0f7876..aebfe3b96 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -32,6 +32,7 @@ ConstantIntVal* SemanticsVisitor::checkConstantIntVal(Expr* expr)
         IntegerConstantExpressionCoercionType::AnyInteger,
         nullptr,
         ConstantFoldingKind::CompileTime);
+
     if (!intVal)
         return nullptr;
 
@@ -524,17 +525,32 @@ Modifier* SemanticsVisitor::validateAttribute(
         }
 
         // TODO(JS): Prior validation currently doesn't ensure both args are ints (as specified in
-        // core.meta.slang), so check here to make sure they both are
-        auto binding = checkConstantIntVal(attr->args[0]);
-        auto set = checkConstantIntVal(attr->args[1]);
+        // core.meta.slang), so check here to make sure they both are.
+        //
+        // Binding attribute may also be created from GLSL style layout qualifiers where only one of
+        // the args are specified, hence check for each individually.
+        ConstantIntVal* binding = nullptr;
+        if (attr->args[0])
+            binding = checkConstantIntVal(attr->args[0]);
+
+        ConstantIntVal* set = nullptr;
+        if (attr->args[1])
+            set = checkConstantIntVal(attr->args[1]);
 
-        if (binding == nullptr || set == nullptr)
+        if (!binding && !set)
         {
             return nullptr;
         }
 
-        bindingAttr->binding = int32_t(binding->getValue());
-        bindingAttr->set = int32_t(set->getValue());
+        if (binding)
+        {
+            bindingAttr->binding = int32_t(binding->getValue());
+        }
+
+        if (set)
+        {
+            bindingAttr->set = int32_t(set->getValue());
+        }
     }
     else if (auto simpleLayoutAttr = as<GLSLSimpleIntegerLayoutAttribute>(attr))
     {
@@ -1430,6 +1446,97 @@ bool isModifierAllowedOnDecl(bool isGLSLInput, ASTNodeType modifierType, Decl* d
     }
 }
 
+void GLSLBindingOffsetTracker::setBindingOffset(int binding, int64_t byteOffset)
+{
+    bindingToByteOffset.set(binding, byteOffset);
+}
+
+int64_t GLSLBindingOffsetTracker::getNextBindingOffset(int binding)
+{
+    int64_t currentOffset;
+    if (bindingToByteOffset.addIfNotExists(binding, 0))
+        currentOffset = 0;
+    else
+        currentOffset = bindingToByteOffset.getValue(binding) + sizeof(uint32_t);
+
+    bindingToByteOffset.set(binding, currentOffset + sizeof(uint32_t));
+    return currentOffset;
+}
+
+AttributeBase* SemanticsVisitor::checkGLSLLayoutAttribute(
+    UncheckedGLSLLayoutAttribute* uncheckedAttr,
+    ModifiableSyntaxNode* attrTarget)
+{
+    SLANG_ASSERT(uncheckedAttr->args.getCount() == 1);
+
+    Attribute* attr = nullptr;
+
+    // False if the current unchecked attribute node is deleted and does not result in a new checked
+    // attribute.
+    bool addNode = true;
+
+    if (as<UncheckedGLSLBindingLayoutAttribute>(uncheckedAttr) ||
+        as<UncheckedGLSLSetLayoutAttribute>(uncheckedAttr))
+    {
+        // Binding and set are coupled together as a descriptor table slot resource for codegen.
+        // Attempt to retrieve and annotate an existing binding attribute or create a new one.
+        attr = attrTarget->findModifier<GLSLBindingAttribute>();
+        if (!attr)
+        {
+            attr = m_astBuilder->create<GLSLBindingAttribute>();
+        }
+        else
+        {
+            addNode = false;
+        }
+
+        // `validateAttribute`, which will be called to parse the binding arguments, also accepts
+        // modifiers from vk::binding and gl::binding where both set and binding are specified.
+        // Binding is the first and set is the second argument - specify them explicitly here.
+        if (as<UncheckedGLSLBindingLayoutAttribute>(uncheckedAttr))
+        {
+            uncheckedAttr->args.add(nullptr);
+        }
+        else
+        {
+            uncheckedAttr->args.add(uncheckedAttr->args[0]);
+            uncheckedAttr->args[0] = nullptr;
+        }
+
+        SLANG_ASSERT(uncheckedAttr->args.getCount() == 2);
+    }
+    else if (as<UncheckedGLSLOffsetLayoutAttribute>(uncheckedAttr))
+    {
+        attr = m_astBuilder->create<GLSLOffsetLayoutAttribute>();
+    }
+    else
+    {
+        getSink()->diagnose(uncheckedAttr, Diagnostics::unrecognizedGLSLLayoutQualifier);
+    }
+
+    if (attr)
+    {
+        attr->keywordName = uncheckedAttr->keywordName;
+        attr->originalIdentifierToken = uncheckedAttr->originalIdentifierToken;
+        attr->args = uncheckedAttr->args;
+        attr->loc = uncheckedAttr->loc;
+
+        // Offset constant folding computation is deferred until all other modifiers are checked to
+        // ensure bindinig is checked first.
+        if (!as<GLSLOffsetLayoutAttribute>(attr))
+        {
+            validateAttribute(attr, nullptr, attrTarget);
+        }
+    }
+
+    if (!addNode)
+    {
+        attr = nullptr;
+    }
+
+    return attr;
+}
+
 Modifier* SemanticsVisitor::checkModifier(
     Modifier* m,
     ModifiableSyntaxNode* syntaxNode,
@@ -1455,6 +1562,21 @@ Modifier* SemanticsVisitor::checkModifier(
         return checkedAttr;
     }
 
+    if (auto glslLayoutAttribute = as<UncheckedGLSLLayoutAttribute>(m))
+    {
+        return checkGLSLLayoutAttribute(glslLayoutAttribute, syntaxNode);
+    }
+
+    if (const auto glslImplicitOffsetAttribute = as<GLSLImplicitOffsetLayoutAttribute>(m))
+    {
+        auto offsetAttr = m_astBuilder->create<GLSLOffsetLayoutAttribute>();
+        offsetAttr->loc = glslImplicitOffsetAttribute->loc;
+
+        // Offset constant folding computation is deferred until all other modifiers are checked to
+        // ensure bindinig is checked first.
+        return offsetAttr;
+    }
+
     if (auto decl = as<Decl>(syntaxNode))
     {
         auto moduleDecl = getModuleDecl(decl);
@@ -1903,6 +2025,31 @@ void SemanticsVisitor::checkModifiers(ModifiableSyntaxNode* syntaxNode)
     // install the new list of modifiers on the declaration
     syntaxNode->modifiers.first = resultModifiers;
 
+    // GLSL offset layout qualifiers are resolved after all other modifiers are checked to ensure
+    // binding layout qualifiers are processed first.
+    if (auto glslOffsetAttribute = syntaxNode->findModifier<GLSLOffsetLayoutAttribute>())
+    {
+        if (const auto glslBindingAttribute = syntaxNode->findModifier<GLSLBindingAttribute>())
+        {
+            if (glslOffsetAttribute->args.getCount() == 0)
+            {
+                glslOffsetAttribute->offset = getGLSLBindingOffsetTracker()->getNextBindingOffset(
+                    glslBindingAttribute->binding);
+            }
+            else if (const auto constVal = checkConstantIntVal(glslOffsetAttribute->args[0]))
+            {
+                glslOffsetAttribute->offset = uint64_t(constVal->getValue());
+                getGLSLBindingOffsetTracker()->setBindingOffset(
+                    glslBindingAttribute->binding,
+                    glslOffsetAttribute->offset);
+            }
+        }
+        else
+        {
+            getSink()->diagnose(glslOffsetAttribute, Diagnostics::missingLayoutBindingModifier);
+        }
+    }
+
     postProcessingOnModifiers(syntaxNode->modifiers);
 }
 
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index d0f743e5c..d6409e42e 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -1207,6 +1207,13 @@ DIAGNOSTIC(
     Error,
     cannotUseUnsizedTypeInConstantBuffer,
     "cannot use unsized type '$0' in a constant buffer.")
+DIAGNOSTIC(31216, Error, unrecognizedGLSLLayoutQualifier, "GLSL layout qualifier is unrecognized")
+DIAGNOSTIC(
+    31217,
+    Error,
+    unrecognizedGLSLLayoutQualifierOrRequiresAssignment,
+    "GLSL layout qualifier is unrecognized or requires assignment")
+
 
 // Enums
 
diff --git a/source/slang/slang-parser.cpp b/source/slang/slang-parser.cpp
index 4dc962ce6..2fb44921d 100644
--- a/source/slang/slang-parser.cpp
+++ b/source/slang/slang-parser.cpp
@@ -230,26 +230,6 @@ public:
             lastErrorLoc = loc;
         }
     }
-
-public:
-    void setBindingOffset(int binding, int64_t byteOffset)
-    {
-        bindingToByteOffset.set(binding, byteOffset);
-    }
-    int64_t getNextBindingOffset(int binding)
-    {
-        int64_t currentOffset;
-        if (bindingToByteOffset.addIfNotExists(binding, 0))
-            currentOffset = 0;
-        else
-            currentOffset = bindingToByteOffset.getValue(binding) + sizeof(uint32_t);
-
-        bindingToByteOffset.set(binding, currentOffset + sizeof(uint32_t));
-        return currentOffset;
-    }
-
-private:
-    Dictionary<int, int64_t> bindingToByteOffset;
 };
 
 // Forward Declarations
@@ -1035,6 +1015,49 @@ static void ParseSquareBracketAttributes(Parser* parser, Modifier*** ioModifierL
     }
 }
 
+static Modifier* parseUncheckedGLSLLayoutAttribute(Parser* parser, NameLoc& nameLoc)
+{
+    // Only valued GLSL layout qualifiers should be parsed here.
+    if (!AdvanceIf(parser, TokenType::OpAssign))
+    {
+        return parser->astBuilder->create<GLSLUnparsedLayoutModifier>();
+    }
+
+    UncheckedGLSLLayoutAttribute* attr;
+
+    if (nameLoc.name->text == "binding")
+    {
+        // An explicit type for binding is used so that it can be looked up quickly
+        // through the list builder when implicitly injecting an offset qualifier.
+        attr = parser->astBuilder->create<UncheckedGLSLBindingLayoutAttribute>();
+    }
+    else if (nameLoc.name->text == "offset")
+    {
+        // An explicit type for offset is used so that it can be looked up quickly
+        // through the list builder when implicitly injecting an offset qualifier.
+        attr = parser->astBuilder->create<UncheckedGLSLOffsetLayoutAttribute>();
+    }
+    else if (nameLoc.name->text == "set")
+    {
+        attr = parser->astBuilder->create<UncheckedGLSLSetLayoutAttribute>();
+    }
+    else
+    {
+        attr = parser->astBuilder->create<UncheckedGLSLLayoutAttribute>();
+    }
+
+    attr->keywordName = nameLoc.name;
+    attr->loc = nameLoc.loc;
+
+    Expr* arg = parser->ParseArgExpr();
+    if (arg)
+    {
+        attr->args.add(arg);
+    }
+
+    return attr;
+}
+
 static TokenType peekTokenType(Parser* parser)
 {
     return parser->tokenReader.peekTokenType();
@@ -4574,9 +4597,16 @@ static void addSpecialGLSLModifiersBasedOnType(Parser* parser, Decl* decl, Modif
     auto declRefExpr = as<DeclRefExpr>(varDeclBase->type.exp);
     if (!declRefExpr)
         return;
-    auto bindingMod = modifiers->findModifier<GLSLBindingAttribute>();
+
+    AttributeBase* bindingMod = modifiers->findModifier<GLSLBindingAttribute>();
     if (!bindingMod)
+    {
+        bindingMod = modifiers->findModifier<UncheckedGLSLBindingLayoutAttribute>();
+    }
+    if (!bindingMod)
+    {
         return;
+    }
 
     // here is a problem; we link types into a literal in IR stage post parse
     // but, order (top down) mattter when parsing atomic_uint offset
@@ -4587,14 +4617,10 @@ static void addSpecialGLSLModifiersBasedOnType(Parser* parser, Decl* decl, Modif
     {
         if (name->text.equals("atomic_uint"))
         {
-            if (!modifiers->findModifier<GLSLOffsetLayoutAttribute>())
+            if (!modifiers->findModifier<UncheckedGLSLOffsetLayoutAttribute>())
             {
-                const int64_t nextOffset = parser->getNextBindingOffset(bindingMod->binding);
-                GLSLOffsetLayoutAttribute* modifier =
-                    parser->astBuilder->create<GLSLOffsetLayoutAttribute>();
-                modifier->keywordName = NULL;    // no keyword name given
-                modifier->loc = bindingMod->loc; // has no location in file, set to parent binding
-                modifier->offset = nextOffset;
+                auto* modifier = parser->astBuilder->create<GLSLImplicitOffsetLayoutAttribute>();
+                modifier->loc = bindingMod->loc; // has no location in file, set to parent
 
                 Modifiers newModifier;
                 newModifier.first = modifier;
@@ -8441,36 +8467,6 @@ static NodeBase* parseLayoutModifier(Parser* parser, void* /*userData*/)
                 inputAttachmentIndexLayoutAttribute->location = intVal;
             }
         }
-        else if (nameText == "binding" || nameText == "set")
-        {
-            GLSLBindingAttribute* attr = listBuilder.find<GLSLBindingAttribute>();
-            if (!attr)
-            {
-                attr = parser->astBuilder->create<GLSLBindingAttribute>();
-                listBuilder.add(attr);
-            }
-
-            parser->ReadToken(TokenType::OpAssign);
-
-            // If the token asked for is not returned found will put in recovering state, and return
-            // token found
-            Token valToken = parser->ReadToken(TokenType::IntegerLiteral);
-            // If wasn't the desired IntegerLiteral return that couldn't parse
-            if (valToken.type == TokenType::IntegerLiteral)
-            {
-                // Work out the value
-                auto value = getIntegerLiteralValue(valToken);
-
-                if (nameText == "binding")
-                {
-                    attr->binding = int32_t(value);
-                }
-                else
-                {
-                    attr->set = int32_t(value);
-                }
-            }
-        }
         else if (findImageFormatByName(nameText.getUnownedSlice(), &format))
         {
             auto attr = parser->astBuilder->create<FormatAttribute>();
@@ -8494,10 +8490,9 @@ static NodeBase* parseLayoutModifier(Parser* parser, void* /*userData*/)
             CASE(std140, GLSLStd140Modifier)
             CASE(std430, GLSLStd430Modifier)
             CASE(scalar, GLSLScalarModifier)
-            CASE(offset, GLSLOffsetLayoutAttribute)
             CASE(location, GLSLLocationLayoutModifier)
             {
-                modifier = parser->astBuilder->create<GLSLUnparsedLayoutModifier>();
+                modifier = parseUncheckedGLSLLayoutAttribute(parser, nameAndLoc);
             }
             SLANG_ASSERT(modifier);
 #undef CASE
@@ -8513,23 +8508,6 @@ static NodeBase* parseLayoutModifier(Parser* parser, void* /*userData*/)
                 if (AdvanceIf(parser, TokenType::OpAssign))
                     glslModifier->valToken = parser->ReadToken(TokenType::IntegerLiteral);
             }
-            // Special handling for GLSLOffsetLayoutAttribute to add to the byte offset tracker at a
-            // binding location
-            else if (auto glslOffset = as<GLSLOffsetLayoutAttribute>(modifier))
-            {
-                if (auto binding = listBuilder.find<GLSLBindingAttribute>())
-                {
-                    // all GLSLOffsetLayoutAttribute have an OpAssign with value token
-                    parser->ReadToken(TokenType::OpAssign);
-                    glslOffset->offset = int64_t(
-                        getIntegerLiteralValue(parser->ReadToken(TokenType::IntegerLiteral)));
-                    parser->setBindingOffset(binding->binding, glslOffset->offset);
-                }
-                else
-                {
-                    parser->diagnose(modifier->loc, Diagnostics::missingLayoutBindingModifier);
-                }
-            }
             else if (auto specConstAttr = as<VkConstantIdAttribute>(modifier))
             {
                 parser->ReadToken(TokenType::OpAssign);
@@ -8537,6 +8515,13 @@ static NodeBase* parseLayoutModifier(Parser* parser, void* /*userData*/)
                     (int)getIntegerLiteralValue(parser->ReadToken(TokenType::IntegerLiteral));
             }
 
+            if (as<GLSLUnparsedLayoutModifier>(modifier))
+            {
+                parser->diagnose(
+                    modifier,
+                    Diagnostics::unrecognizedGLSLLayoutQualifierOrRequiresAssignment);
+            }
+
             listBuilder.add(modifier);
         }