[HLSL][RootSignature] Correct RootSignatureParser to use correct SourceLocation in diagnostics (#147084)

The `SourceLocation` of a `RootSignatureToken` is incorrectly set to be
the "offset" into the concatenated string that denotes the
rootsignature. This causes an issue when the `StringLiteral` is a
multi-line expansion macro, since the offset will not account for the
characters between `StringLiteral` tokens.

This pr resolves this by retaining the `SourceLocation` information that
is kept in `StringLiteral` and then converting the offset in the
concatenated string into the proper `SourceLocation` using the
`StringLiteral::getLocationOfByte` interface. To do so, we will need to
adjust the `RootSignatureToken` to only hold its offset into the root
signature string. Then when the parser will use the token, it will need
to compute its actual `SourceLocation`.

See linked issue for more context.

For example:

```
#define DemoRootSignature \
 "CBV(b0)," \
 "RootConstants(num32BitConstants = 3, b0, invalid)"
  expected caret location ---------------^
  actual caret location ------------^
```

The caret points 5 characters early because the current offset did not
account for the characters:
```
'"' ' ' '\' ' ' '"'
 1   2   3   4   5
```

- Updates `RootSignatureParser` to retain `SourceLocation` information
by retaining the `StringLiteral` and passing the underlying `StringRef`
to the `Lexer`
- Updates `RootSignatureLexer` so that the constructed tokens only
reflect an offset into the `StringRef`
- Updates `RootSignatureParser` to directly construct its used `Lexer`
so that the `StringLiteral` is directly tied with the string used in the
`RootSignatureLexer`
- Updates `RootSignatureParser` to use
`StringLiteral::getLocationOfByte` to get the actual token location for
diagnostics
- Updates `ParseHLSLRootSignatureTest` to construct a phony
`AST`/`StringLiteral` for the test cases
- Adds a test to `RootSignature-err.hlsl` showing that the
`SourceLocation` is correctly set for diagnostics in a multi-line macro
expansion

Resolves: #146967

Author: inbelic

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

@@ -4941,20 +4941,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;