[llvm] [SandboxVectorizer][NFCI] Fix use of possibly-uninitialized scalar. (PR #122201)

[Tyler Lanphear via llvm-commits]

https://github.com/tylanphear updated https://github.com/llvm/llvm-project/pull/122201

>From 807217e0e987ca370b8892a3d9a03fe7e2ae184b Mon Sep 17 00:00:00 2001
From: Tyler Lanphear <tyler.lanphear at intel.com>
Date: Wed, 8 Jan 2025 17:14:09 -0800
Subject: [PATCH 1/3] [SandboxVectorizer][NFCI] Fix use of
 possibly-uninitialized scalar.

The `EraseCallbackID` field is not always initialized in the ctor for
SeedCollector, if not, it will be used uninitialized by its dtor. This
could potentially lead to the erasure of a random callback, leading to a
bug.

Fixed by initializing to a purposefully invalid ID
(`std::numeric_limits<ContextID>::max()`)
---
 .../Transforms/Vectorize/SandboxVectorizer/SeedCollector.h    | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
index 6e16a84d832e5e..cc4ebb5c508414 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
@@ -292,7 +292,9 @@ class SeedCollector {
   SeedContainer StoreSeeds;
   SeedContainer LoadSeeds;
   Context &Ctx;
-  Context::CallbackID EraseCallbackID;
+  Context::CallbackID EraseCallbackID =
+      std::numeric_limits<Context::CallbackID>::max();
+
   /// \Returns the number of SeedBundle groups for all seed types.
   /// This is to be used for limiting compilation time.
   unsigned totalNumSeedGroups() const {

>From 9344694553b4dc72849eb8ed7069155efd813464 Mon Sep 17 00:00:00 2001
From: Tyler Lanphear <tyler.lanphear at intel.com>
Date: Thu, 9 Jan 2025 14:57:14 -0800
Subject: [PATCH 2/3] Make CallbackID an opaque type, and hide its
 implementation details in Context

---
 llvm/include/llvm/SandboxIR/Context.h         | 51 ++++++++++++++++---
 .../SandboxVectorizer/SeedCollector.h         |  3 +-
 llvm/lib/SandboxIR/Context.cpp                |  6 +--
 3 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/Context.h b/llvm/include/llvm/SandboxIR/Context.h
index b0d6f8335d9e0e..917a281f2db981 100644
--- a/llvm/include/llvm/SandboxIR/Context.h
+++ b/llvm/include/llvm/SandboxIR/Context.h
@@ -18,7 +18,8 @@
 
 #include <cstdint>
 
-namespace llvm::sandboxir {
+namespace llvm {
+namespace sandboxir {
 
 class Argument;
 class BBIterator;
@@ -37,10 +38,25 @@ class Context {
   using MoveInstrCallback =
       std::function<void(Instruction *, const BBIterator &)>;
 
-  /// An ID for a registered callback. Used for deregistration. Using a 64-bit
-  /// integer so we don't have to worry about the unlikely case of overflowing
-  /// a 32-bit counter.
-  using CallbackID = uint64_t;
+  /// An ID for a registered callback. Used for deregistration. A dedicated type
+  /// is employed so as to keep IDs opaque to the end user; only Context should
+  /// deal with its underlying representation.
+  class CallbackID {
+  public:
+    // Uses a 64-bit integer so we don't have to worry about the unlikely case
+    // of overflowing a 32-bit counter.
+    using ReprTy = uint64_t;
+
+  private:
+    // Default initialization results in an invalid ID.
+    ReprTy Val = std::numeric_limits<ReprTy>::max();
+    explicit CallbackID(ReprTy Val) : Val{Val} {}
+
+  public:
+    CallbackID() = default;
+    friend class Context;
+    friend struct DenseMapInfo<CallbackID>;
+  };
 
 protected:
   LLVMContext &LLVMCtx;
@@ -83,7 +99,7 @@ class Context {
   /// A counter used for assigning callback IDs during registration. The same
   /// counter is used for all kinds of callbacks so we can detect mismatched
   /// registration/deregistration.
-  CallbackID NextCallbackID = 0;
+  uint64_t NextCallbackID = 0;
 
   /// Remove \p V from the maps and returns the unique_ptr.
   std::unique_ptr<Value> detachLLVMValue(llvm::Value *V);
@@ -263,6 +279,27 @@ class Context {
   // TODO: Add callbacks for instructions inserted/removed if needed.
 };
 
-} // namespace llvm::sandboxir
+} // namespace sandboxir
+
+// DenseMap info for CallbackIDs
+template <> struct DenseMapInfo<sandboxir::Context::CallbackID> {
+  using CallbackID = sandboxir::Context::CallbackID;
+  using ReprInfo = DenseMapInfo<CallbackID::ReprTy>;
+
+  static CallbackID getEmptyKey() {
+    return CallbackID{ReprInfo::getEmptyKey()};
+  }
+  static CallbackID getTombstoneKey() {
+    return CallbackID{ReprInfo::getTombstoneKey()};
+  }
+  static unsigned getHashValue(const CallbackID &ID) {
+    return ReprInfo::getHashValue(ID.Val);
+  }
+  static bool isEqual(const CallbackID &LHS, const CallbackID &RHS) {
+    return ReprInfo::isEqual(LHS.Val, RHS.Val);
+  }
+};
+
+} // namespace llvm
 
 #endif // LLVM_SANDBOXIR_CONTEXT_H
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
index cc4ebb5c508414..8e9a8973720dc1 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
@@ -292,8 +292,7 @@ class SeedCollector {
   SeedContainer StoreSeeds;
   SeedContainer LoadSeeds;
   Context &Ctx;
-  Context::CallbackID EraseCallbackID =
-      std::numeric_limits<Context::CallbackID>::max();
+  Context::CallbackID EraseCallbackID;
 
   /// \Returns the number of SeedBundle groups for all seed types.
   /// This is to be used for limiting compilation time.
diff --git a/llvm/lib/SandboxIR/Context.cpp b/llvm/lib/SandboxIR/Context.cpp
index b86ed5864c1ac1..42ca456881fd0d 100644
--- a/llvm/lib/SandboxIR/Context.cpp
+++ b/llvm/lib/SandboxIR/Context.cpp
@@ -686,7 +686,7 @@ void Context::runMoveInstrCallbacks(Instruction *I, const BBIterator &WhereIt) {
 Context::CallbackID Context::registerEraseInstrCallback(EraseInstrCallback CB) {
   assert(EraseInstrCallbacks.size() <= MaxRegisteredCallbacks &&
          "EraseInstrCallbacks size limit exceeded");
-  CallbackID ID = NextCallbackID++;
+  CallbackID ID{NextCallbackID++};
   EraseInstrCallbacks[ID] = CB;
   return ID;
 }
@@ -700,7 +700,7 @@ Context::CallbackID
 Context::registerCreateInstrCallback(CreateInstrCallback CB) {
   assert(CreateInstrCallbacks.size() <= MaxRegisteredCallbacks &&
          "CreateInstrCallbacks size limit exceeded");
-  CallbackID ID = NextCallbackID++;
+  CallbackID ID{NextCallbackID++};
   CreateInstrCallbacks[ID] = CB;
   return ID;
 }
@@ -713,7 +713,7 @@ void Context::unregisterCreateInstrCallback(CallbackID ID) {
 Context::CallbackID Context::registerMoveInstrCallback(MoveInstrCallback CB) {
   assert(MoveInstrCallbacks.size() <= MaxRegisteredCallbacks &&
          "MoveInstrCallbacks size limit exceeded");
-  CallbackID ID = NextCallbackID++;
+  CallbackID ID{NextCallbackID++};
   MoveInstrCallbacks[ID] = CB;
   return ID;
 }

>From 3c8349c259407b26cb55866f597b48853b85b276 Mon Sep 17 00:00:00 2001
From: Tyler Lanphear <tyler.lanphear at intel.com>
Date: Thu, 9 Jan 2025 15:00:04 -0800
Subject: [PATCH 3/3] Fix accidental explicit usage of uint64_t

---
 llvm/include/llvm/SandboxIR/Context.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/include/llvm/SandboxIR/Context.h b/llvm/include/llvm/SandboxIR/Context.h
index 917a281f2db981..4c37c4c1f58e5f 100644
--- a/llvm/include/llvm/SandboxIR/Context.h
+++ b/llvm/include/llvm/SandboxIR/Context.h
@@ -99,7 +99,7 @@ class Context {
   /// A counter used for assigning callback IDs during registration. The same
   /// counter is used for all kinds of callbacks so we can detect mismatched
   /// registration/deregistration.
-  uint64_t NextCallbackID = 0;
+  CallbackID::ReprTy NextCallbackID = 0;
 
   /// Remove \p V from the maps and returns the unique_ptr.
   std::unique_ptr<Value> detachLLVMValue(llvm::Value *V);