From dceb1a74fcaa6031148abd2130552c77fc800809 Mon Sep 17 00:00:00 2001
From: Ellie Hermaszewska <ellieh@nvidia.com>
Date: Tue, 8 Aug 2023 00:19:36 +0800
Subject: Allow parsing some SPIR-V enums in intrinsics (#3062)

* Add TokenReader::AdvanceIf overlaod for TokenType

* Add some spirv defs to parser

* Add comment
---
 source/core/slang-token-reader.h        |  9 +++++++
 source/slang/slang-ir-spirv-snippet.cpp | 44 ++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/source/core/slang-token-reader.h b/source/core/slang-token-reader.h
index 264693c81..c4b552d29 100644
--- a/source/core/slang-token-reader.h
+++ b/source/core/slang-token-reader.h
@@ -262,6 +262,15 @@ namespace Misc {
             }
             return false;
         }
+        bool AdvanceIf(TokenType token)
+        {
+            if( LookAhead(token) )
+            {
+                ReadToken();
+                return true;
+            }
+            return false;
+        }
         bool IsEnd()
         {
             return tokenPtr == (int)tokens.getCount();
diff --git a/source/slang/slang-ir-spirv-snippet.cpp b/source/slang/slang-ir-spirv-snippet.cpp
index 3d416a7e3..81e4d0bff 100644
--- a/source/slang/slang-ir-spirv-snippet.cpp
+++ b/source/slang/slang-ir-spirv-snippet.cpp
@@ -37,6 +37,48 @@ SpvSnippet::ASMType parseASMType(Slang::Misc::TokenReader& tokenReader)
     return SpvSnippet::ASMType::None;
 }
 
+// Read an unsigned integer (a SPIR-V word) or a SPIR-V enum (currently those
+// which are coded into this function).
+//
+// This also 'or's together a list of these words/enums separated by '|'
+SpvWord readWordOrWordLiteral(Misc::TokenReader& reader)
+{
+    SpvWord ret = 0;
+    do
+    {
+        switch(reader.NextToken().Type)
+        {
+            case Slang::Misc::TokenType::IntLiteral:
+                ret = reader.ReadUInt();
+                break;
+            case Slang::Misc::TokenType::Identifier:
+                {
+                    const auto i = reader.ReadWord();
+#define GO(x) if(i == #x) ret |= Spv ## x
+                         GO(ScopeWorkgroup);
+                    else GO(ScopeDevice);
+                    else GO(MemorySemanticsMaskNone);
+                    else GO(MemorySemanticsAcquireReleaseMask);
+                    else GO(MemorySemanticsUniformMemoryMask);
+                    else GO(MemorySemanticsImageMemoryMask);
+                    else GO(MemorySemanticsAtomicCounterMemoryMask);
+                    else GO(MemorySemanticsWorkgroupMemoryMask);
+#undef GO
+                    else
+                    {
+                        reader.Back(1);
+                        throw Misc::TextFormatException(
+                            "Text parsing error: Unrecognized SPIR-V enum: " + i);
+                    }
+                }
+                break;
+            default:
+                throw Misc::TextFormatException("Text parsing error: Expected int or SPIR-V enum");
+        }
+    } while(reader.AdvanceIf(Misc::TokenType::OpBitOr));
+    return ret;
+}
+
 RefPtr<SpvSnippet> SpvSnippet::parse(UnownedStringSlice definition)
 {
     RefPtr<SpvSnippet> snippet = new SpvSnippet();
@@ -226,7 +268,7 @@ RefPtr<SpvSnippet> SpvSnippet::parse(UnownedStringSlice definition)
                                     break;
                                     
                                 default:
-                                    constant.intValues[i] = tokenReader.ReadInt();
+                                    constant.intValues[i] = readWordOrWordLiteral(tokenReader);
                                     ++i;
                                     break;
                                 }
-- 
cgit v1.2.3