Reapply "[SandboxIR][Tracker] Implement accept(/*AcceptAll*/) and revert(/*RevertAll*/)" (#199776) (#199805)

This reverts commit a7aceff0b1.

llvm/docs/SandboxIR.md

@@ -108,7 +108,8 @@ To start tracking again, the user needs to call `save()`.
 
 Sandbox IR supports nested checkpoints, meaning that you can save more than once and revert more than once.
 Conceptually each `save()` adds a new checkpoint to a stack and each `revert()` rolls back the IR state to that of the checkpoint at the top of the stack and pops the checkpoint off the stack.
-A call to `accept()` will clear the stack.
+A call to `accept()` pops the last checkpoint from the stack.
+Reverting or accepting all can be done with `revert(/*RevertAll=*/true)` and `accept(/*AcceptAll=*/true)`.
 
 ## Users of Sandbox IR
 - [The Sandbox Vectorizer](project:SandboxVectorizer.md)

llvm/include/llvm/SandboxIR/Context.h

@@ -249,9 +249,9 @@ public:
   /// Convenience function for `getTracker().save()`
   void save() { IRTracker.save(); }
   /// Convenience function for `getTracker().revert()`
-  void revert() { IRTracker.revert(); }
+  void revert(bool RevertAll = false) { IRTracker.revert(RevertAll); }
   /// Convenience function for `getTracker().accept()`
-  void accept() { IRTracker.accept(); }
+  void accept(bool AcceptAll = false) { IRTracker.accept(AcceptAll); }
 
   LLVM_ABI sandboxir::Value *getValue(llvm::Value *V) const;
   const sandboxir::Value *getValue(const llvm::Value *V) const {

llvm/include/llvm/SandboxIR/Tracker.h

@@ -19,13 +19,18 @@
 // are saved in the order they are registered with the tracker and are stored in
 // the `Tracker::Changes` vector. All of this is done transparently to
 // the user.
+// Calling `Tracker::save()` a second time without having accepted/reverted the
+// state, creates a second nested checkpoint.
 //
 // Reverting changes
 // -----------------
-// Calling `Tracker::revert()` will restore the state saved when
+// Calling `Tracker::revert()` will restore the state saved when the last
 // `Tracker::save()` was called. Internally this goes through the
 // change objects in `Tracker::Changes` in reverse order, calling their
 // `IRChangeBase::revert()` function one by one.
+// In the context of a nested checkpoint, this will revert the state
+// until the last `Tracker::save()` checkpoint.
+// You can revert all changes with `Tracker::revert(/*RevertAll=*/true)`.
 //
 // Accepting changes
 // -----------------
@@ -34,6 +39,9 @@
 // This is the job of `Tracker::accept()`. Internally this will go
 // through the change objects in `Tracker::Changes` in order, calling
 // `IRChangeBase::accept()`.
+// In the context of a nested checkpoint, this will leave the state unchanged
+// and will remove the last checkpoint for the stack.
+// You can accept all changes with `Tracker::accept(/*AcceptAll=*/true)`.
 //
 //===----------------------------------------------------------------------===//
 
@@ -504,12 +512,16 @@ public:
   bool isTracking() const { return State == TrackerState::Record; }
   /// \Returns the current state of the tracker.
   TrackerState getState() const { return State; }
-  /// Turns on IR tracking.
+  /// Turns on IR tracking. If tracking is already enabled this creates a new
+  /// nested checkpoint.
   LLVM_ABI void save();
-  /// Stops tracking and accept changes.
-  LLVM_ABI void accept();
-  /// Stops tracking and reverts to saved state.
-  LLVM_ABI void revert();
+  /// Stops tracking and accept changes. If we have nested checkpoints, this
+  /// will remove the last checkpoint from the stack without modifying the
+  /// state.
+  LLVM_ABI void accept(bool AcceptAll = false);
+  /// Stops tracking and reverts to saved state. If we have nested checkpoints
+  /// this will revert the state to the last checkpoint.
+  LLVM_ABI void revert(bool RevertAll = false);
   /// \returns the number of nested (outstanding) checkpoints.
   unsigned nestingDepth() const { return Snapshots.size(); }
 

llvm/lib/SandboxIR/Tracker.cpp

@@ -338,28 +338,43 @@ void Tracker::save() {
 #endif
 }
 
-void Tracker::revert() {
+void Tracker::revert(bool RevertAll) {
   assert(State == TrackerState::Record && "Forgot to save()!");
   State = TrackerState::Reverting;
-  const unsigned ToRevert = Changes.size() - Snapshots.back();
+  unsigned UntilChangeIdx = RevertAll ? 0 : Snapshots.back();
+  const unsigned ToRevert = Changes.size() - UntilChangeIdx;
   unsigned CntReverts = 0;
   for (auto &Change : reverse(Changes)) {
     // Stop reverting if we reach the index of the last snapshot.
     if (CntReverts++ == ToRevert)
       break;
     Change->revert(*this);
+#if !defined(NDEBUG) && defined(EXPENSIVE_CHECKS)
+    // There may be multiple changes between snapshots, so use the snapshot
+    // checker only if this change has an associated snapshot.
+    unsigned ChangeIdx = Changes.size() - CntReverts;
+    bool ChangeHasSnapshot = ChangeIdx == Snapshots.back();
+    if (ChangeHasSnapshot) {
+      SnapshotChecker.back().expectNoDiff();
+      SnapshotChecker.pop_back();
+    }
+#endif
   }
   Changes.erase(Changes.end() - ToRevert, Changes.end());
-  Snapshots.pop_back();
-#if !defined(NDEBUG) && defined(EXPENSIVE_CHECKS)
-  SnapshotChecker.back().expectNoDiff();
-  SnapshotChecker.pop_back();
-#endif
+  if (RevertAll)
+    Snapshots.clear();
+  else
+    Snapshots.pop_back();
   State = Snapshots.empty() ? TrackerState::Disabled : TrackerState::Record;
 }
 
-void Tracker::accept() {
+void Tracker::accept(bool AcceptAll) {
   assert(State == TrackerState::Record && "Forgot to save()!");
+  if (!AcceptAll && Snapshots.size() > 1) {
+    // Just remove the last stacked checkpoint.
+    Snapshots.pop_back();
+    return;
+  }
   State = TrackerState::Disabled;
   for (auto &Change : Changes)
     Change->accept();

llvm/unittests/SandboxIR/TrackerTest.cpp

@@ -2053,6 +2053,61 @@ define i32 @foo(i32 %arg0, i32 %arg1) {
   EXPECT_EQ(Ctx.getTracker().nestingDepth(), 0u);
   EXPECT_EQ(Add1->getOperand(0), Arg0);
   Checker.expectNoDiff();
+
+  // Check revert(/*RevertAll=*/true)
+  Checker.save();
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 1u);
+  Add1->setOperand(0, Arg1);
+  Add1->setOperand(0, Arg1);
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 2u);
+  Add1->setOperand(0, Arg0);
+  Ctx.revert(/*RevertAll=*/true);
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 0u);
+  Checker.expectNoDiff();
+
+  // Check revert(/*RevertAll=*/false)
+  Checker.save();
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 1u);
+  Add1->setOperand(0, Arg1);
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 2u);
+  Add1->setOperand(0, Arg0);
+  Ctx.revert(/*RevertAll=*/false);
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 1u);
+  EXPECT_EQ(Add1->getOperand(0), Arg1);
+  Ctx.revert(/*RevertAll=*/false);
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 0u);
+  Checker.expectNoDiff();
+
+  // Check accept(/*AcceptAll=*/false)
+  auto *OrigOp = Add1->getOperand(0);
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 1u);
+  Add1->setOperand(0, Arg1);
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 2u);
+  Add1->setOperand(0, Arg0);
+  Ctx.accept(/*RevertAll=*/false);
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 1u);
+  EXPECT_EQ(Add1->getOperand(0), Arg0);
+  Ctx.accept(/*RevertAll=*/false);
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 0u);
+  EXPECT_EQ(Add1->getOperand(0), Arg0);
+  Add1->setOperand(0, OrigOp);
+
+  // Check accept(/*AcceptAll=*/true)
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 1u);
+  Add1->setOperand(0, Arg1);
+  Ctx.save();
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 2u);
+  Add1->setOperand(0, Arg0);
+  Ctx.accept(/*RevertAll=*/true);
+  EXPECT_EQ(Ctx.getTracker().nestingDepth(), 0u);
+  EXPECT_EQ(Add1->getOperand(0), Arg0);
 }
 
 #endif // NDEBUG