[HLSL][RootSignature] Retain SourceLocation of RootElement for SemaHLSL diagnostics

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -13098,6 +13098,7 @@ def err_hlsl_resource_range_overlap: Error<
   "resource ranges %sub{subst_hlsl_format_ranges}0,1,2,3 and %sub{subst_hlsl_format_ranges}4,5,6,7 "
   "overlap within space = %8 and visibility = "
   "%select{All|Vertex|Hull|Domain|Geometry|Pixel|Amplification|Mesh}9">;
+def note_hlsl_resource_range_here: Note<"overlapping resource range here">;
 
 // Layout randomization diagnostics.
 def err_non_designated_init_used : Error<

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -17,6 +17,7 @@
 #include "clang/Basic/DiagnosticParse.h"
 #include "clang/Lex/LexHLSLRootSignature.h"
 #include "clang/Lex/Preprocessor.h"
+#include "clang/Sema/SemaHLSL.h"
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -29,7 +30,7 @@ namespace hlsl {
 class RootSignatureParser {
 public:
   RootSignatureParser(llvm::dxbc::RootSignatureVersion Version,
-                      SmallVector<llvm::hlsl::rootsig::RootElement> &Elements,
+                      SmallVector<RootSignatureElement> &Elements,
                       StringLiteral *Signature, Preprocessor &PP);
 
   /// Consumes tokens from the Lexer and constructs the in-memory
@@ -201,7 +202,7 @@ class RootSignatureParser {
 
 private:
   llvm::dxbc::RootSignatureVersion Version;
-  SmallVector<llvm::hlsl::rootsig::RootElement> &Elements;
+  SmallVector<RootSignatureElement> &Elements;
   StringLiteral *Signature;
   RootSignatureLexer Lexer;
   Preprocessor &PP;

clang/include/clang/Sema/SemaHLSL.h

@@ -32,6 +32,25 @@ class ParsedAttr;
 class Scope;
 class VarDecl;
 
+namespace hlsl {
+
+// Introduce a wrapper struct around the underlying RootElement. This structure
+// will retain extra clang diagnostic information that is not available in llvm.
+struct RootSignatureElement {
+  RootSignatureElement(SourceLocation Loc,
+                       llvm::hlsl::rootsig::RootElement Element)
+      : Loc(Loc), Element(Element) {}
+
+  const llvm::hlsl::rootsig::RootElement &getElement() const { return Element; }
+  const SourceLocation &getLocation() const { return Loc; }
+
+private:
+  SourceLocation Loc;
+  llvm::hlsl::rootsig::RootElement Element;
+};
+
+} // namespace hlsl
+
 using llvm::dxil::ResourceClass;
 
 // FIXME: This can be hidden (as static function in SemaHLSL.cpp) once we no
@@ -130,12 +149,14 @@ class SemaHLSL : public SemaBase {
 
   /// Creates the Root Signature decl of the parsed Root Signature elements
   /// onto the AST and push it onto current Scope
-  void ActOnFinishRootSignatureDecl(
-      SourceLocation Loc, IdentifierInfo *DeclIdent,
-      SmallVector<llvm::hlsl::rootsig::RootElement> &Elements);
-
-  // Returns true when D is invalid and a diagnostic was produced
-  bool handleRootSignatureDecl(HLSLRootSignatureDecl *D, SourceLocation Loc);
+  void
+  ActOnFinishRootSignatureDecl(SourceLocation Loc, IdentifierInfo *DeclIdent,
+                               ArrayRef<hlsl::RootSignatureElement> Elements);
+
+  // Returns true if any RootSignatureElement is invalid and a diagnostic was
+  // produced
+  bool
+  handleRootSignatureElements(ArrayRef<hlsl::RootSignatureElement> Elements);
   void handleRootSignatureAttr(Decl *D, const ParsedAttr &AL);
   void handleNumThreadsAttr(Decl *D, const ParsedAttr &AL);
   void handleWaveSizeAttr(Decl *D, const ParsedAttr &AL);

clang/lib/Parse/ParseDeclCXX.cpp

@@ -4948,7 +4948,7 @@ void Parser::ParseHLSLRootSignatureAttributeArgs(ParsedAttributes &Attrs) {
   // signature string and construct the in-memory elements
   if (!Found) {
     // Invoke the root signature parser to construct the in-memory constructs
-    SmallVector<llvm::hlsl::rootsig::RootElement> RootElements;
+    SmallVector<hlsl::RootSignatureElement> RootElements;
     hlsl::RootSignatureParser Parser(getLangOpts().HLSLRootSigVer, RootElements,
                                      Signature, PP);
     if (Parser.parse()) {

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -19,7 +19,7 @@ using TokenKind = RootSignatureToken::Kind;
 
 RootSignatureParser::RootSignatureParser(
     llvm::dxbc::RootSignatureVersion Version,
-    SmallVector<RootElement> &Elements, StringLiteral *Signature,
+    SmallVector<RootSignatureElement> &Elements, StringLiteral *Signature,
     Preprocessor &PP)
     : Version(Version), Elements(Elements), Signature(Signature),
       Lexer(Signature->getString()), PP(PP), CurToken(0) {}
@@ -29,31 +29,36 @@ bool RootSignatureParser::parse() {
   // end of the stream
   while (!peekExpectedToken(TokenKind::end_of_stream)) {
     if (tryConsumeExpectedToken(TokenKind::kw_RootFlags)) {
+      SourceLocation ElementLoc = getTokenLocation(CurToken);
       auto Flags = parseRootFlags();
       if (!Flags.has_value())
         return true;
-      Elements.push_back(*Flags);
+      Elements.emplace_back(ElementLoc, *Flags);
     } else if (tryConsumeExpectedToken(TokenKind::kw_RootConstants)) {
+      SourceLocation ElementLoc = getTokenLocation(CurToken);
       auto Constants = parseRootConstants();
       if (!Constants.has_value())
         return true;
-      Elements.push_back(*Constants);
+      Elements.emplace_back(ElementLoc, *Constants);
     } else if (tryConsumeExpectedToken(TokenKind::kw_DescriptorTable)) {
+      SourceLocation ElementLoc = getTokenLocation(CurToken);
       auto Table = parseDescriptorTable();
       if (!Table.has_value())
         return true;
-      Elements.push_back(*Table);
+      Elements.emplace_back(ElementLoc, *Table);
     } else if (tryConsumeExpectedToken(
                    {TokenKind::kw_CBV, TokenKind::kw_SRV, TokenKind::kw_UAV})) {
+      SourceLocation ElementLoc = getTokenLocation(CurToken);
       auto Descriptor = parseRootDescriptor();
       if (!Descriptor.has_value())
         return true;
-      Elements.push_back(*Descriptor);
+      Elements.emplace_back(ElementLoc, *Descriptor);
     } else if (tryConsumeExpectedToken(TokenKind::kw_StaticSampler)) {
+      SourceLocation ElementLoc = getTokenLocation(CurToken);
       auto Sampler = parseStaticSampler();
       if (!Sampler.has_value())
         return true;
-      Elements.push_back(*Sampler);
+      Elements.emplace_back(ElementLoc, *Sampler);
     }
 
     // ',' denotes another element, otherwise, expected to be at end of stream
@@ -245,10 +250,11 @@ std::optional<DescriptorTable> RootSignatureParser::parseDescriptorTable() {
     if (tryConsumeExpectedToken({TokenKind::kw_CBV, TokenKind::kw_SRV,
                                  TokenKind::kw_UAV, TokenKind::kw_Sampler})) {
       // DescriptorTableClause - CBV, SRV, UAV, or Sampler
+      SourceLocation ElementLoc = getTokenLocation(CurToken);
       auto Clause = parseDescriptorTableClause();
       if (!Clause.has_value())
         return std::nullopt;
-      Elements.push_back(*Clause);
+      Elements.emplace_back(ElementLoc, *Clause);
       Table.NumClauses++;
     } else if (tryConsumeExpectedToken(TokenKind::kw_visibility)) {
       // visibility = SHADER_VISIBILITY

clang/lib/Sema/SemaHLSL.cpp

@@ -1064,27 +1064,33 @@ SemaHLSL::ActOnStartRootSignatureDecl(StringRef Signature) {
 
 void SemaHLSL::ActOnFinishRootSignatureDecl(
     SourceLocation Loc, IdentifierInfo *DeclIdent,
-    SmallVector<llvm::hlsl::rootsig::RootElement> &Elements) {
+    ArrayRef<hlsl::RootSignatureElement> RootElements) {
+
+  if (handleRootSignatureElements(RootElements))
+    return;
+
+  SmallVector<llvm::hlsl::rootsig::RootElement> Elements;
+  for (auto &RootSigElement : RootElements)
+    Elements.push_back(RootSigElement.getElement());
 
   auto *SignatureDecl = HLSLRootSignatureDecl::Create(
       SemaRef.getASTContext(), /*DeclContext=*/SemaRef.CurContext, Loc,
       DeclIdent, SemaRef.getLangOpts().HLSLRootSigVer, Elements);
 
-  if (handleRootSignatureDecl(SignatureDecl, Loc))
-    return;
-
   SignatureDecl->setImplicit();
   SemaRef.PushOnScopeChains(SignatureDecl, SemaRef.getCurScope());
 }
 
-bool SemaHLSL::handleRootSignatureDecl(HLSLRootSignatureDecl *D,
-                                       SourceLocation Loc) {
+bool SemaHLSL::handleRootSignatureElements(
+    ArrayRef<hlsl::RootSignatureElement> Elements) {
   using RangeInfo = llvm::hlsl::rootsig::RangeInfo;
   using OverlappingRanges = llvm::hlsl::rootsig::OverlappingRanges;
+  using InfoPairT = std::pair<RangeInfo, const hlsl::RootSignatureElement *>;
 
   // 1. Collect RangeInfos
-  llvm::SmallVector<RangeInfo> Infos;
-  for (const llvm::hlsl::rootsig::RootElement &Elem : D->getRootElements()) {
+  llvm::SmallVector<InfoPairT> InfoPairs;
+  for (const hlsl::RootSignatureElement &RootSigElem : Elements) {
+    const llvm::hlsl::rootsig::RootElement &Elem = RootSigElem.getElement();
     if (const auto *Descriptor =
             std::get_if<llvm::hlsl::rootsig::RootDescriptor>(&Elem)) {
       RangeInfo Info;
@@ -1095,7 +1101,8 @@ bool SemaHLSL::handleRootSignatureDecl(HLSLRootSignatureDecl *D,
           llvm::dxil::ResourceClass(llvm::to_underlying(Descriptor->Type));
       Info.Space = Descriptor->Space;
       Info.Visibility = Descriptor->Visibility;
-      Infos.push_back(Info);
+
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Constants =
                    std::get_if<llvm::hlsl::rootsig::RootConstants>(&Elem)) {
       RangeInfo Info;
@@ -1105,7 +1112,8 @@ bool SemaHLSL::handleRootSignatureDecl(HLSLRootSignatureDecl *D,
       Info.Class = llvm::dxil::ResourceClass::CBuffer;
       Info.Space = Constants->Space;
       Info.Visibility = Constants->Visibility;
-      Infos.push_back(Info);
+
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Sampler =
                    std::get_if<llvm::hlsl::rootsig::StaticSampler>(&Elem)) {
       RangeInfo Info;
@@ -1115,7 +1123,8 @@ bool SemaHLSL::handleRootSignatureDecl(HLSLRootSignatureDecl *D,
       Info.Class = llvm::dxil::ResourceClass::Sampler;
       Info.Space = Sampler->Space;
       Info.Visibility = Sampler->Visibility;
-      Infos.push_back(Info);
+
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Clause =
                    std::get_if<llvm::hlsl::rootsig::DescriptorTableClause>(
                        &Elem)) {
@@ -1129,38 +1138,83 @@ bool SemaHLSL::handleRootSignatureDecl(HLSLRootSignatureDecl *D,
 
       Info.Class = Clause->Type;
       Info.Space = Clause->Space;
+
       // Note: Clause does not hold the visibility this will need to
-      Infos.push_back(Info);
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Table =
                    std::get_if<llvm::hlsl::rootsig::DescriptorTable>(&Elem)) {
       // Table holds the Visibility of all owned Clauses in Table, so iterate
       // owned Clauses and update their corresponding RangeInfo
-      assert(Table->NumClauses <= Infos.size() && "RootElement");
+      assert(Table->NumClauses <= InfoPairs.size() && "RootElement");
       // The last Table->NumClauses elements of Infos are the owned Clauses
       // generated RangeInfo
       auto TableInfos =
-          MutableArrayRef<RangeInfo>(Infos).take_back(Table->NumClauses);
-      for (RangeInfo &Info : TableInfos)
-        Info.Visibility = Table->Visibility;
+          MutableArrayRef<InfoPairT>(InfoPairs).take_back(Table->NumClauses);
+      for (InfoPairT &Pair : TableInfos)
+        Pair.first.Visibility = Table->Visibility;
     }
   }
 
-  // Helper to report diagnostics
-  auto ReportOverlap = [this, Loc](OverlappingRanges Overlap) {
+  // 2. Sort with the RangeInfo <operator to prepare it for findOverlapping
+  llvm::sort(InfoPairs,
+             [](InfoPairT A, InfoPairT B) { return A.first < B.first; });
+
+  llvm::SmallVector<RangeInfo> Infos;
+  for (const InfoPairT &Pair : InfoPairs)
+    Infos.push_back(Pair.first);
+
+  // Helpers to report diagnostics
+  uint32_t DuplicateCounter = 0;
+  using ElemPair = std::pair<const hlsl::RootSignatureElement *,
+                             const hlsl::RootSignatureElement *>;
+  auto GetElemPair = [&Infos, &InfoPairs, &DuplicateCounter](
+                         OverlappingRanges Overlap) -> ElemPair {
+    // Given we sorted the InfoPairs (and by implication) Infos, and,
+    // that Overlap.B is the item retrieved from the ResourceRange. Then it is
+    // guarenteed that Overlap.B <= Overlap.A.
+    //
+    // So we will find Overlap.B first and then continue to find Overlap.A
+    // after
+    auto InfoB = std::lower_bound(Infos.begin(), Infos.end(), *Overlap.B);
+    auto DistB = std::distance(Infos.begin(), InfoB);
+    auto PairB = InfoPairs.begin();
+    std::advance(PairB, DistB);
+
+    auto InfoA = std::lower_bound(InfoB, Infos.end(), *Overlap.A);
+    // Similarily, from the property that we have sorted the RangeInfos,
+    // all duplicates will be processed one after the other. So
+    // DuplicateCounter can be re-used for each set of duplicates we
+    // encounter as we handle incoming errors
+    DuplicateCounter = InfoA == InfoB ? DuplicateCounter + 1 : 0;
+    auto DistA = std::distance(InfoB, InfoA) + DuplicateCounter;
+    auto PairA = PairB;
+    std::advance(PairA, DistA);
+
+    return {PairA->second, PairB->second};
+  };
+
+  auto ReportOverlap = [this, &GetElemPair](OverlappingRanges Overlap) {
+    auto Pair = GetElemPair(Overlap);
     const RangeInfo *Info = Overlap.A;
+    const hlsl::RootSignatureElement *Elem = Pair.first;
     const RangeInfo *OInfo = Overlap.B;
+
     auto CommonVis = Info->Visibility == llvm::dxbc::ShaderVisibility::All
                          ? OInfo->Visibility
                          : Info->Visibility;
-    this->Diag(Loc, diag::err_hlsl_resource_range_overlap)
+    this->Diag(Elem->getLocation(), diag::err_hlsl_resource_range_overlap)
         << llvm::to_underlying(Info->Class) << Info->LowerBound
         << /*unbounded=*/(Info->UpperBound == RangeInfo::Unbounded)
         << Info->UpperBound << llvm::to_underlying(OInfo->Class)
         << OInfo->LowerBound
         << /*unbounded=*/(OInfo->UpperBound == RangeInfo::Unbounded)
         << OInfo->UpperBound << Info->Space << CommonVis;
+
+    const hlsl::RootSignatureElement *OElem = Pair.second;
+    this->Diag(OElem->getLocation(), diag::note_hlsl_resource_range_here);
   };
 
+  // 3. Invoke find overlapping ranges
   llvm::SmallVector<OverlappingRanges> Overlaps =
       llvm::hlsl::rootsig::findOverlappingRanges(Infos);
   for (OverlappingRanges Overlap : Overlaps)