Allow generics to close with >>

5 files changed, 91 insertions, 32 deletions

diff --git a/source/slang/lexer.cpp b/source/slang/lexer.cpp
index 3d87dd7ea..b6745b15f 100644
--- a/source/slang/lexer.cpp
+++ b/source/slang/lexer.cpp
@@ -12,7 +12,7 @@
 
 namespace Slang
 {
-    static Token GetEndOfFileToken()
+    Token TokenReader::GetEndOfFileToken()
     {
         return Token(TokenType::EndOfFile, UnownedStringSlice::fromLiteral(""), SourceLoc());
     }
@@ -41,31 +41,19 @@ namespace Slang
     {}
 
 
-    Token TokenReader::PeekToken() const
+    Token& TokenReader::PeekToken()
     {
-        if (!mCursor)
-            return GetEndOfFileToken();
-
-        Token token = *mCursor;
-        if (mCursor == mEnd)
-            token.type = TokenType::EndOfFile;
-        return token;
+        return nextToken;
     }
 
     TokenType TokenReader::PeekTokenType() const
     {
-        if (mCursor == mEnd)
-            return TokenType::EndOfFile;
-        SLANG_ASSERT(mCursor);
-        return mCursor->type;
+        return nextToken.type;
     }
 
     SourceLoc TokenReader::PeekLoc() const
     {
-        if (!mCursor)
-            return SourceLoc();
-        SLANG_ASSERT(mCursor);
-        return mCursor->loc;
+        return nextToken.loc;
     }
 
     Token TokenReader::AdvanceToken()
@@ -73,11 +61,14 @@ namespace Slang
         if (!mCursor)
             return GetEndOfFileToken();
 
-        Token token = *mCursor;
-        if (mCursor == mEnd)
-            token.type = TokenType::EndOfFile;
-        else
+        Token token = nextToken;
+        if (mCursor < mEnd)
+        {
             mCursor++;
+            nextToken = *mCursor;
+        }
+        else
+            nextToken.type = TokenType::EndOfFile;
         return token;
     }
 
diff --git a/source/slang/lexer.h b/source/slang/lexer.h
index 593e07c22..8587cc904 100644
--- a/source/slang/lexer.h
+++ b/source/slang/lexer.h
@@ -38,18 +38,37 @@ namespace Slang
 
     struct TokenReader
     {
+        Token nextToken;
         TokenReader();
         explicit TokenReader(TokenSpan const& tokens)
             : mCursor(tokens.begin())
             , mEnd   (tokens.end  ())
+            , nextToken(tokens.begin() ? *tokens.begin() : GetEndOfFileToken())
         {}
         explicit TokenReader(TokenList const& tokens)
             : mCursor(tokens.begin())
             , mEnd   (tokens.end  ())
+            , nextToken(tokens.begin() ? *tokens.begin() : GetEndOfFileToken())
         {}
-
+        struct ParsingCursor
+        {
+            Token nextToken;
+            Token* tokenReaderCursor = nullptr;
+        };
+        ParsingCursor getCursor()
+        {
+            ParsingCursor rs;
+            rs.nextToken = nextToken;
+            rs.tokenReaderCursor = mCursor;
+            return rs;
+        }
+        void setCursor(ParsingCursor cursor)
+        {
+            mCursor = cursor.tokenReaderCursor;
+            nextToken = cursor.nextToken;
+        }
         bool IsAtEnd() const { return mCursor == mEnd; }
-        Token PeekToken() const;
+        Token& PeekToken();
         TokenType PeekTokenType() const;
         SourceLoc PeekLoc() const;
 
@@ -59,6 +78,7 @@ namespace Slang
 
         Token* mCursor;
         Token* mEnd;
+        static Token GetEndOfFileToken();
     };
 
     typedef unsigned int LexerFlags;
diff --git a/source/slang/parser.cpp b/source/slang/parser.cpp
index 076111711..06c7e8a1c 100644
--- a/source/slang/parser.cpp
+++ b/source/slang/parser.cpp
@@ -138,7 +138,6 @@ namespace Slang
         }
 
         RefPtr<ModuleDecl> Parse();
-
         Token ReadToken();
         Token ReadToken(TokenType type);
         Token ReadToken(const char * string);
@@ -1667,11 +1666,15 @@ namespace Slang
         }
         parser->genericDepth--;
 
-        // TODO: probably want to suppot `>>` and `>=` here as well,
-        // and split up those tokens as needed.
-        //
-        parser->ReadToken(TokenType::OpGreater);
-
+        if (parser->tokenReader.PeekToken().type == TokenType::OpRsh)
+        {
+            parser->tokenReader.PeekToken().type = TokenType::OpGreater;
+            parser->tokenReader.PeekToken().loc.setRaw(parser->tokenReader.PeekToken().loc.getRaw() + 1);
+        }
+        else if (parser->LookAheadToken(TokenType::OpGreater))
+            parser->ReadToken(TokenType::OpGreater);
+        else
+            parser->sink->diagnose(parser->tokenReader.PeekToken(), Diagnostics::tokenTypeExpected, "'>'");
         return genericApp;
     }
 
@@ -3444,7 +3447,7 @@ namespace Slang
             //
             // We'll solve this with backtracking for now.
 
-            Token* startPos = tokenReader.mCursor;
+            TokenReader::ParsingCursor startPos = tokenReader.getCursor();
 
             // Try to parse a type (knowing that the type grammar is
             // a subset of the expression grammar, and so this should
@@ -3466,13 +3469,13 @@ namespace Slang
                 // Reset the cursor and try to parse a declaration now.
                 // Note: the declaration will consume any modifiers
                 // that had been in place on the statement.
-                tokenReader.mCursor = startPos;
+                tokenReader.setCursor(startPos);
                 statement = parseVarDeclrStatement(modifiers);
                 return statement;
             }
 
             // Fallback: reset and parse an expression
-            tokenReader.mCursor = startPos;
+            tokenReader.setCursor(startPos);
             statement = ParseExpressionStatement();
         }
         else if (LookAheadToken(TokenType::Semicolon))
diff --git a/tests/compute/generic-closer.slang b/tests/compute/generic-closer.slang
new file mode 100644
index 000000000..f3bb643de
--- /dev/null
+++ b/tests/compute/generic-closer.slang
@@ -0,0 +1,41 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+
+interface IGetter
+{
+    int get();
+}
+struct Gen0 : IGetter
+{
+    int get() { return 1; }
+};
+struct Gen1<TGetter : IGetter> : IGetter
+{
+    TGetter g;
+    int get() { return g.get(); }
+};
+
+RWStructuredBuffer<int> outputBuffer;
+void writeArray(inout float3 a[4])
+{
+    a[0] = float3(1, 1, 1);
+    a[1] = float3(1, 1, 1);
+    a[2] = float3(1, 1, 1);
+    a[3] = float3(1, 1, 1);    
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    Gen1<Gen1<Gen1<Gen0>>> g;
+    Gen1<Gen1<Gen1<Gen1<Gen0>>>> g2;
+    Gen1<Gen1<Gen0>> g3;
+
+    int a = 0; 
+    int b = 5;
+    if (a< b && b > a)
+	    outputBuffer[dispatchThreadID.x] = (g.get() >> 1) + g2.get() + g3.get();
+}
\ No newline at end of file
diff --git a/tests/compute/generic-closer.slang.expected.txt b/tests/compute/generic-closer.slang.expected.txt
new file mode 100644
index 000000000..01dd2b944
--- /dev/null
+++ b/tests/compute/generic-closer.slang.expected.txt
@@ -0,0 +1,4 @@
+2
+2
+2
+2
\ No newline at end of file