[llvm] a4cc7e7 - [PGO] Avoid potential const_cast UB (NFC)

[Christian Ulmann via llvm-commits]

Author: Christian Ulmann
Date: 2023-05-02T09:18:37Z
New Revision: a4cc7e784f9333baf0f355affca4cc31cacc2a5b

URL: https://github.com/llvm/llvm-project/commit/a4cc7e784f9333baf0f355affca4cc31cacc2a5b
DIFF: https://github.com/llvm/llvm-project/commit/a4cc7e784f9333baf0f355affca4cc31cacc2a5b.diff

LOG: [PGO] Avoid potential const_cast UB (NFC)

This commit removes potential UB in the PGO instrumentation passes that
was caused by casting away constness and then potentially modifying the
object.

Reviewed By: gysit

Differential Revision: https://reviews.llvm.org/D148903

Added: 
    

Modified: 
    llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
    llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h b/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
index a4d2efd821c76..4d31898bb3147 100644
--- a/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
+++ b/llvm/include/llvm/Transforms/Instrumentation/CFGMST.h
@@ -100,7 +100,7 @@ template <class Edge, class BBInfo> class CFGMST {
   void buildEdges() {
     LLVM_DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n");
 
-    const BasicBlock *Entry = &(F.getEntryBlock());
+    BasicBlock *Entry = &(F.getEntryBlock());
     uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2);
     // If we want to instrument the entry count, lower the weight to 0.
     if (InstrumentFuncEntry)
@@ -257,7 +257,7 @@ template <class Edge, class BBInfo> class CFGMST {
   }
 
   // Add an edge to AllEdges with weight W.
-  Edge &addEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W) {
+  Edge &addEdge(BasicBlock *Src, BasicBlock *Dest, uint64_t W) {
     uint32_t Index = BBInfos.size();
     auto Iter = BBInfos.end();
     bool Inserted;

diff  --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 416f559bda343..e7fb036df2b48 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -484,14 +484,14 @@ struct SelectInstVisitor : public InstVisitor<SelectInstVisitor> {
 /// Note that the CFG can be a multi-graph. So there might be multiple edges
 /// with the same SrcBB and DestBB.
 struct PGOEdge {
-  const BasicBlock *SrcBB;
-  const BasicBlock *DestBB;
+  BasicBlock *SrcBB;
+  BasicBlock *DestBB;
   uint64_t Weight;
   bool InMST = false;
   bool Removed = false;
   bool IsCritical = false;
 
-  PGOEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W = 1)
+  PGOEdge(BasicBlock *Src, BasicBlock *Dest, uint64_t W = 1)
       : SrcBB(Src), DestBB(Dest), Weight(W) {}
 
   /// Return the information string of an edge.
@@ -786,8 +786,8 @@ BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   if (E->InMST || E->Removed)
     return nullptr;
 
-  BasicBlock *SrcBB = const_cast<BasicBlock *>(E->SrcBB);
-  BasicBlock *DestBB = const_cast<BasicBlock *>(E->DestBB);
+  BasicBlock *SrcBB = E->SrcBB;
+  BasicBlock *DestBB = E->DestBB;
   // For a fake edge, instrument the real BB.
   if (SrcBB == nullptr)
     return DestBB;
@@ -988,12 +988,11 @@ namespace {
 
 // This class represents a CFG edge in profile use compilation.
 struct PGOUseEdge : public PGOEdge {
+  using PGOEdge::PGOEdge;
+
   bool CountValid = false;
   uint64_t CountValue = 0;
 
-  PGOUseEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W = 1)
-      : PGOEdge(Src, Dest, W) {}
-
   // Set edge count value
   void setEdgeCount(uint64_t Value) {
     CountValue = Value;