[HLSL][RootSignature] Correct RootSignatureParser to use correct SourceLocation in diagnostics

clang/include/clang/Lex/LexHLSLRootSignature.h

@@ -31,16 +31,17 @@ struct RootSignatureToken {
 
   Kind TokKind = Kind::invalid;
 
-  // Retain the SouceLocation of the token for diagnostics
-  clang::SourceLocation TokLoc;
+  // Retain the location offset of the token in the Signature
+  // string
+  uint32_t LocOffset;
 
   // Retain spelling of an numeric constant to be parsed later
   StringRef NumSpelling;
 
   // Constructors
-  RootSignatureToken(clang::SourceLocation TokLoc) : TokLoc(TokLoc) {}
-  RootSignatureToken(Kind TokKind, clang::SourceLocation TokLoc)
-      : TokKind(TokKind), TokLoc(TokLoc) {}
+  RootSignatureToken(uint32_t LocOffset) : LocOffset(LocOffset) {}
+  RootSignatureToken(Kind TokKind, uint32_t LocOffset)
+      : TokKind(TokKind), LocOffset(LocOffset) {}
 };
 
 inline const DiagnosticBuilder &
@@ -61,8 +62,7 @@ operator<<(const DiagnosticBuilder &DB, const RootSignatureToken::Kind Kind) {
 
 class RootSignatureLexer {
 public:
-  RootSignatureLexer(StringRef Signature, clang::SourceLocation SourceLoc)
-      : Buffer(Signature), SourceLoc(SourceLoc) {}
+  RootSignatureLexer(StringRef Signature) : Buffer(Signature) {}
 
   /// Consumes and returns the next token.
   RootSignatureToken consumeToken();
@@ -76,23 +76,21 @@ class RootSignatureLexer {
   }
 
 private:
-  // Internal buffer to iterate over
+  // Internal buffer state
   StringRef Buffer;
+  uint32_t LocOffset = 0;
 
   // Current peek state
   std::optional<RootSignatureToken> NextToken = std::nullopt;
 
-  // Passed down parameters from Sema
-  clang::SourceLocation SourceLoc;
-
   /// Consumes the buffer and returns the lexed token.
   RootSignatureToken lexToken();
 
   /// Advance the buffer by the specified number of characters.
   /// Updates the SourceLocation appropriately.
   void advanceBuffer(unsigned NumCharacters = 1) {
     Buffer = Buffer.drop_front(NumCharacters);
-    SourceLoc = SourceLoc.getLocWithOffset(NumCharacters);
+    LocOffset += NumCharacters;
   }
 };
 

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_CLANG_PARSE_PARSEHLSLROOTSIGNATURE_H
 #define LLVM_CLANG_PARSE_PARSEHLSLROOTSIGNATURE_H
 
+#include "clang/AST/Expr.h"
 #include "clang/Basic/DiagnosticParse.h"
 #include "clang/Lex/LexHLSLRootSignature.h"
 #include "clang/Lex/Preprocessor.h"
@@ -29,7 +30,7 @@ class RootSignatureParser {
 public:
   RootSignatureParser(llvm::dxbc::RootSignatureVersion Version,
                       SmallVector<llvm::hlsl::rootsig::RootElement> &Elements,
-                      RootSignatureLexer &Lexer, clang::Preprocessor &PP);
+                      StringLiteral *Signature, Preprocessor &PP);
 
   /// Consumes tokens from the Lexer and constructs the in-memory
   /// representations of the RootElements. Tokens are consumed until an
@@ -187,12 +188,23 @@ class RootSignatureParser {
   bool tryConsumeExpectedToken(RootSignatureToken::Kind Expected);
   bool tryConsumeExpectedToken(ArrayRef<RootSignatureToken::Kind> Expected);
 
+  /// Convert the token's offset in the signature string to its SourceLocation
+  ///
+  /// This allows to currently retrieve the location for multi-token
+  /// StringLiterals
+  SourceLocation getTokenLocation(RootSignatureToken Tok);
+
+  /// Construct a diagnostics at the location of the current token
+  DiagnosticBuilder reportDiag(unsigned DiagID) {
+    return getDiags().Report(getTokenLocation(CurToken), DiagID);
+  }
+
 private:
   llvm::dxbc::RootSignatureVersion Version;
   SmallVector<llvm::hlsl::rootsig::RootElement> &Elements;
-  RootSignatureLexer &Lexer;
-
-  clang::Preprocessor &PP;
+  StringLiteral *Signature;
+  RootSignatureLexer Lexer;
+  Preprocessor &PP;
 
   RootSignatureToken CurToken;
 };

clang/lib/Lex/LexHLSLRootSignature.cpp

@@ -27,10 +27,10 @@ RootSignatureToken RootSignatureLexer::lexToken() {
   advanceBuffer(Buffer.take_while(isspace).size());
 
   if (isEndOfBuffer())
-    return RootSignatureToken(TokenKind::end_of_stream, SourceLoc);
+    return RootSignatureToken(TokenKind::end_of_stream, LocOffset);
 
   // Record where this token is in the text for usage in parser diagnostics
-  RootSignatureToken Result(SourceLoc);
+  RootSignatureToken Result(LocOffset);
 
   char C = Buffer.front();
 
@@ -62,7 +62,7 @@ RootSignatureToken RootSignatureLexer::lexToken() {
 
   // All following tokens require at least one additional character
   if (Buffer.size() <= 1) {
-    Result = RootSignatureToken(TokenKind::invalid, SourceLoc);
+    Result = RootSignatureToken(TokenKind::invalid, LocOffset);
     return Result;
   }
 

clang/lib/Parse/ParseDeclCXX.cpp

@@ -4944,20 +4944,16 @@ void Parser::ParseHLSLRootSignatureAttributeArgs(ParsedAttributes &Attrs) {
   }
 
   // Construct our identifier
-  StringRef Signature = StrLiteral.value()->getString();
+  StringLiteral *Signature = StrLiteral.value();
   auto [DeclIdent, Found] =
-      Actions.HLSL().ActOnStartRootSignatureDecl(Signature);
+      Actions.HLSL().ActOnStartRootSignatureDecl(Signature->getString());
   // If we haven't found an already defined DeclIdent then parse the root
   // signature string and construct the in-memory elements
   if (!Found) {
-    // Offset location 1 to account for '"'
-    SourceLocation SignatureLoc =
-        StrLiteral.value()->getExprLoc().getLocWithOffset(1);
     // Invoke the root signature parser to construct the in-memory constructs
-    hlsl::RootSignatureLexer Lexer(Signature, SignatureLoc);
     SmallVector<llvm::hlsl::rootsig::RootElement> RootElements;
     hlsl::RootSignatureParser Parser(getLangOpts().HLSLRootSigVer, RootElements,
-                                     Lexer, PP);
+                                     Signature, PP);
     if (Parser.parse()) {
       T.consumeClose();
       return;