proto: using std::pair mapping with sorting outside of function

Author: inbelic

clang/lib/Sema/SemaHLSL.cpp

@@ -1085,9 +1085,10 @@ bool SemaHLSL::handleRootSignatureElements(
     ArrayRef<hlsl::RootSignatureElement> Elements) {
   using RangeInfo = llvm::hlsl::rootsig::RangeInfo;
   using OverlappingRanges = llvm::hlsl::rootsig::OverlappingRanges;
+  using InfoPair = std::pair<RangeInfo, const hlsl::RootSignatureElement *>;
 
   // 1. Collect RangeInfos
-  llvm::SmallVector<RangeInfo> Infos;
+  llvm::SmallVector<InfoPair> InfoPairs;
   for (const hlsl::RootSignatureElement &RootSigElem : Elements) {
     const llvm::hlsl::rootsig::RootElement &Elem = RootSigElem.getElement();
     if (const auto *Descriptor =
@@ -1101,8 +1102,7 @@ bool SemaHLSL::handleRootSignatureElements(
       Info.Space = Descriptor->Space;
       Info.Visibility = Descriptor->Visibility;
 
-      Info.Cookie = static_cast<void *>(&RootSigElem);
-      Infos.push_back(Info);
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Constants =
                    std::get_if<llvm::hlsl::rootsig::RootConstants>(&Elem)) {
       RangeInfo Info;
@@ -1113,8 +1113,7 @@ bool SemaHLSL::handleRootSignatureElements(
       Info.Space = Constants->Space;
       Info.Visibility = Constants->Visibility;
 
-      Info.Cookie = static_cast<void *>(&RootSigElem);
-      Infos.push_back(Info);
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Sampler =
                    std::get_if<llvm::hlsl::rootsig::StaticSampler>(&Elem)) {
       RangeInfo Info;
@@ -1125,8 +1124,7 @@ bool SemaHLSL::handleRootSignatureElements(
       Info.Space = Sampler->Space;
       Info.Visibility = Sampler->Visibility;
 
-      Info.Cookie = static_cast<void *>(&RootSigElem);
-      Infos.push_back(Info);
+      InfoPairs.push_back({Info, &RootSigElem});
     } else if (const auto *Clause =
                    std::get_if<llvm::hlsl::rootsig::DescriptorTableClause>(
                        &Elem)) {
@@ -1142,30 +1140,59 @@ bool SemaHLSL::handleRootSignatureElements(
       Info.Space = Clause->Space;
 
       // Note: Clause does not hold the visibility this will need to
-      Info.Cookie = static_cast<void *>(&RootSigElem);
-      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<InfoPair>(InfoPairs).take_back(Table->NumClauses);
+      for (InfoPair &Pair : TableInfos)
+        Pair.first.Visibility = Table->Visibility;
     }
   }
 
-  // Helper to report diagnostics
-  auto ReportOverlap = [this](OverlappingRanges Overlap) {
+  // Sort as specified
+  auto ComparePairs = [](InfoPair A, InfoPair B) { return A.first < B.first; };
+
+  std::sort(InfoPairs.begin(), InfoPairs.end(), ComparePairs);
+
+  llvm::SmallVector<RangeInfo> Infos;
+  for (const InfoPair &Pair : InfoPairs)
+    Infos.push_back(Pair.first);
+
+  // Helpers to report diagnostics
+  using ElemPair = std::pair<const hlsl::RootSignatureElement *,
+                             const hlsl::RootSignatureElement *>;
+  auto GetElemPair = [&Infos, &InfoPairs](
+                         OverlappingRanges Overlap) -> ElemPair {
+    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);
+    if (InfoA == InfoB)
+      InfoA++;
+    auto DistA = std::distance(InfoB, InfoA);
+    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;
-    auto Elem = static_cast<const hlsl::RootSignatureElement *>(Info->Cookie);
     this->Diag(Elem->getLocation(), diag::err_hlsl_resource_range_overlap)
         << llvm::to_underlying(Info->Class) << Info->LowerBound
         << /*unbounded=*/(Info->UpperBound == RangeInfo::Unbounded)
@@ -1174,7 +1201,7 @@ bool SemaHLSL::handleRootSignatureElements(
         << /*unbounded=*/(OInfo->UpperBound == RangeInfo::Unbounded)
         << OInfo->UpperBound << Info->Space << CommonVis;
 
-    auto OElem = static_cast<const hlsl::RootSignatureElement *>(OInfo->Cookie);
+    const hlsl::RootSignatureElement *OElem = Pair.second;
     this->Diag(OElem->getLocation(), diag::note_hlsl_resource_range_here);
   };
 

llvm/include/llvm/Frontend/HLSL/RootSignatureValidations.h

@@ -52,8 +52,17 @@ struct RangeInfo {
   uint32_t Space;
   llvm::dxbc::ShaderVisibility Visibility;
 
-  // Retain information for diagnostic reporting
-  void *Cookie;
+  bool operator==(const RangeInfo &RHS) {
+    return std::tie(LowerBound, UpperBound, Class, Space, Visibility) ==
+           std::tie(RHS.LowerBound, RHS.UpperBound, RHS.Class, RHS.Space,
+                    RHS.Visibility);
+  }
+
+  bool operator<(const RangeInfo &RHS) {
+    return std::tie(Class, Space, LowerBound, UpperBound, Visibility) <
+           std::tie(RHS.Class, RHS.Space, RHS.LowerBound, RHS.UpperBound,
+                    RHS.Visibility);
+  }
 };
 
 class ResourceRange {
@@ -137,7 +146,7 @@ struct OverlappingRanges {
 ///   ResourceRange
 ///      B: Check for overlap with any overlapping Visibility ResourceRange
 llvm::SmallVector<OverlappingRanges>
-findOverlappingRanges(llvm::SmallVector<RangeInfo> &Infos);
+findOverlappingRanges(ArrayRef<RangeInfo> Infos);
 
 } // namespace rootsig
 } // namespace hlsl

llvm/lib/Frontend/HLSL/RootSignatureValidations.cpp

@@ -244,16 +244,11 @@ std::optional<const RangeInfo *> ResourceRange::insert(const RangeInfo &Info) {
 }
 
 llvm::SmallVector<OverlappingRanges>
-findOverlappingRanges(llvm::SmallVector<RangeInfo> &Infos) {
+findOverlappingRanges(ArrayRef<RangeInfo> Infos) {
   // 1. The user has provided the corresponding range information
   llvm::SmallVector<OverlappingRanges> Overlaps;
   using GroupT = std::pair<dxil::ResourceClass, /*Space*/ uint32_t>;
 
-  // 2. Sort the RangeInfo's by their GroupT to form groupings
-  std::sort(Infos.begin(), Infos.end(), [](RangeInfo A, RangeInfo B) {
-    return std::tie(A.Class, A.Space) < std::tie(B.Class, B.Space);
-  });
-
   // 3. First we will init our state to track:
   if (Infos.size() == 0)
     return Overlaps; // No ranges to overlap