[Chkpnt] TakeName and DestroyName

[vporpo]

Summary

This patch adds support for tracking Value::takeName() and name destruction.
This means that a Chkpnt.restore() can now revert these actions.

llvm/include/llvm/IR/CheckpointChanges.h

static bool classof(const ChangeBase *Other) {
     return Other->getID() == ChangeID::SetNameID;
   }
   ~SetName() {}
 #ifndef NDEBUG
   void dump(raw_ostream &OS) const override;
   LLVM_DUMP_METHOD void dump() const override;
 #endif // NDEBUG
 };
+
+class TakeName : public ChangeBase {
+  std::string OrigName;
+  Value *FromV;
+
+public:
+  TakeName(Value *Val, Value *FromV, CheckpointTracker *CT);
+  void revert() override;
+  void apply() override{};
+  static bool classof(const ChangeBase *Other) {
+    return Other->getID() == ChangeID::TakeNameID;
+  }
+  ~TakeName() {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const override;
+  LLVM_DUMP_METHOD void dump() const override;
+#endif // NDEBUG
+};
+
+class DestroyName : public ChangeBase {
+public:
+  DestroyName(Value *Val, CheckpointTracker *CT);
+  void revert() override {}
+  void apply() override;
+  static bool classof(const ChangeBase *Other) {
+    return Other->getID() == ChangeID::DestroyNameID;
+  }
+  ~DestroyName() {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const override;
+  LLVM_DUMP_METHOD void dump() const override;
+#endif // NDEBUG
+};
 } // namespace llvm
 
 #endif // LLVM_IR_CHECKPOINTCHANGES_H

llvm/include/llvm/IR/CheckpointTracker.h

#endif                     // NDEBUG
 
   /// Clears the state.
   void clear();
   /// \Returns the parent component of \p V. Can handle detached instructions.
   std::optional<ChkpntComponent> getParentComponent(Value *V) const;
 
   /// To be called when \p V is about to get its name updated.
   void setName(Value *V);
+  /// Called by V->takeName(FromV).
+  void takeName(Value *V, Value *FromV);
+  /// Gets called by ~Value() before destroyValueName().
+  void destroyName(Value *V);
 
   // Main API functions. These are called by `Checkpoint`.
 
   /// Start tracking IR changes for \p Component from this point on.
   void trackComponent(ChkpntComponent Component);
   /// Override the maximum number of tracked changes.
   void setMaxNumOfTrackedChanges(uint32_t MaxNumOfTrackedChanges);
   /// Accept all changes.

llvm/lib/IR/CheckpointChanges.cpp

 #ifndef NDEBUG
 void SetName::dump(raw_ostream &OS) const {
   dumpCommon(OS);
   OS << "SetName: " << V << " OrigName='" << OrigName << "'\n";
 }
 
 void SetName::dump() const { dump(dbgs()); }
 #endif // NDEBUG
+
+TakeName::TakeName(Value *Val, Value *FromV, CheckpointTracker *CT)
+    : ChangeBase(Val, ChangeID::TakeNameID, CT), OrigName(Val->getName()),
+      FromV(FromV) {}
+
+void TakeName::revert() {
+  std::string CurrName(V->getName());
+  V->setName(OrigName);
+  FromV->setName(CurrName);
+}
+
+#ifndef NDEBUG
+void TakeName::dump(raw_ostream &OS) const {
+  dumpCommon(OS);
+  OS << "TakeName\n";
+}
+
+void TakeName::dump() const { dump(dbgs()); }
+#endif // NDEBUG
+
+DestroyName::DestroyName(Value *Val, CheckpointTracker *CT)
+    : ChangeBase(Val, ChangeID::DestroyNameID, CT) {}
+
+void DestroyName::apply() { V->destroyValueName(); }
+
+#ifndef NDEBUG
+void DestroyName::dump(raw_ostream &OS) const {
+  dumpCommon(OS);
+  OS << "DestroyName\n";
+}
+
+void DestroyName::dump() const { dump(dbgs()); }
+#endif // NDEBUG

llvm/lib/IR/CheckpointTracker.cpp

if (!Parent)
     return;
   // If changing name of a BB, we consider its parent to be the BB itself.
   if (isa<BasicBlock>(V) && trackingComponent(cast<BasicBlock>(V)))
     Parent = cast<BasicBlock>(V);
   if (trackingComponent(*Parent))
     Changes.push_back(make_unique<SetName>(V, this));
 }
 
+void CheckpointTracker::takeName(Value *V, Value *FromV) {
+  std::optional<ChkpntComponent> Parent = getParentComponent(V);
+  if (!Parent)
+    return;
+  // If changing name of a BB, we consider its parent to be the BB itself.
+  if (isa<BasicBlock>(V) && trackingComponent(cast<BasicBlock>(V)))
+    Parent = cast<BasicBlock>(V);
+  if (trackingComponent(*Parent))
+    Changes.push_back(make_unique<TakeName>(V, FromV, this));
+}
+
+void CheckpointTracker::destroyName(Value *V) {
+  std::optional<ChkpntComponent> Parent = getParentComponent(V);
+  if (!Parent)
+    return;
+  // If changing name of a BB, we consider its parent to be the BB itself.
+  if (isa<BasicBlock>(V) && trackingComponent(cast<BasicBlock>(V)))
+    Parent = cast<BasicBlock>(V);
+  if (trackingComponent(*Parent))
+    Changes.push_back(make_unique<DestroyName>(V, this));
+}
+
 void CheckpointTracker::trackComponent(ChkpntComponent Component) {
   assert(!trackingComponent(Component) && "Already tracking component");
   ComponentsTracked.insert(Component);
   Active = true;
 }
 
 void CheckpointTracker::setMaxNumOfTrackedChanges(
     uint32_t MaxNumOfTrackedChanges) {

llvm/lib/IR/Value.cpp

if (!materialized_use_empty()) {
     for (auto *U : users())
       dbgs() << "Use still stuck around after Def is destroyed:" << *U << "\n";
   }
 #endif
   assert(materialized_use_empty() && "Uses remain when a value is destroyed!");
 
   // If this value is named, destroy the name.  This should not be in a symtab
   // at this point.
+  auto &Chkpnt = getContext().getChkpntTracker();
+  if (LLVM_UNLIKELY(Chkpnt.isActive()))
+    Chkpnt.destroyName(this);
+  else
     destroyValueName();
 }
 
 void Value::deleteValue() {
   switch (getValueID()) {
 #define HANDLE_VALUE(Name)                                                     \
   case Value::Name##Val:                                                       \
     delete static_cast<Name *>(this);                                          \
     break;
if (LLVM_UNLIKELY(ChkpntTracker.isActive()))
     ChkpntTracker.setName(this);
   setNameImpl(NewName);
   if (Function *F = dyn_cast<Function>(this))
     F->recalculateIntrinsicID();
 }
 
 void Value::takeName(Value *V) {
   assert(V != this && "Illegal call to this->takeName(this)!");
+  CheckpointTracker &ChkpntTracker = getContext().getChkpntTracker();
+  if (LLVM_UNLIKELY(ChkpntTracker.isActive()))
+    ChkpntTracker.takeName(this, V);
+
   ValueSymbolTable *ST = nullptr;
   // If this value has a name, drop it.
   if (hasName()) {
     // Get the symtab this is in.
     if (getSymTab(this, ST)) {
       // We can't set a name on this value, but we need to clear V's name if
       // it has one.
       if (V->hasName()) V->setName("");

llvm/unittests/IR/CheckpointTest.cpp

)");
   Checkpoint Chkpnt = C.getCheckpoint();
   Chkpnt.track(BB0);
   Instr->setName("new");
   EXPECT_FALSE(Chkpnt.empty());
   Chkpnt.accept();
   EXPECT_EQ(Instr->getName(), "new");
 }
 
+TEST(CheckpointTest, BB_TakeNameInstr) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"(
+define void @foo(i32 %a, i32 %b) {
+bb0:
+  %instr1 = add i32 %a, %b
+  %instr2 = add i32 %a, %b
+  ret void
+}
+)");
+  Function *F = &*M->begin();
+  BasicBlock *BB0 = &*F->begin();
+  auto It = BB0->begin();
+  auto *Instr1 = &*It++;
+  auto *Instr2 = &*It++;
+
+  Checkpoint Chkpnt = C.getCheckpoint();
+  Chkpnt.track(BB0);
+  Instr1->takeName(Instr2);
+  EXPECT_EQ(Instr1->getName(), "instr2");
+  EXPECT_FALSE(Instr2->hasName());
+  EXPECT_FALSE(Chkpnt.empty());
+  Chkpnt.restore();
+
+  EXPECT_EQ(Instr1->getName(), "instr1");
+  EXPECT_EQ(Instr2->getName(), "instr2");
+}
+
 TEST(CheckpointTest, BB_SetNameBB) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(
 define void @foo(i32 %a, i32 %b) {
 bb0:
   %instr = add i32 %a, %b
   ret void
 }