[llvm] [SandboxIR] IR Tracker (PR #99238)

[Arthur Eubanks via llvm-commits]

================
@@ -0,0 +1,181 @@
+//===- SandboxIRTracker.h ---------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is the component of SandboxIR that tracks all changes made to its
+// state, such that we can revert the state when needed.
+//
+// Tracking changes
+// ----------------
+// The user needs to call `SandboxIRTracker::save()` to enable tracking changes
+// made to SandboxIR. From that point on, any change made to SandboxIR, will
+// automatically create a change tracking object and register it with the
+// tracker. IR-change objects are subclasses of `IRChangeBase` and get
+// registered with the `SandboxIRTracker::track()` function. The change objects
+// are saved in the order they are registered with the tracker and are stored in
+// the `SandboxIRTracker::Changes` vector. All of this is done transparently to
+// the user.
+//
+// Reverting changes
+// -----------------
+// Calling `SandboxIRTracker::revert()` will restore the state saved when
+// `SandboxIRTracker::save()` was called. Internally this goes through the
+// change objects in `SandboxIRTracker::Changes` in reverse order, calling their
+// `IRChangeBase::revert()` function one by one.
+//
+// Accepting changes
+// -----------------
+// The user needs to either revert or accept changes before the tracker object
+// is destroyed, or else the tracker destructor will cause a crash.
+// This is the job of `SandboxIRTracker::accept()`. Internally this will go
+// through the change objects in `SandboxIRTracker::Changes` in order, calling
+// `IRChangeBase::accept()`.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
+#define LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+#include "llvm/SandboxIR/Use.h"
+#include "llvm/Support/Debug.h"
+#include <memory>
+#include <regex>
+
+namespace llvm::sandboxir {
+
+class BasicBlock;
+
+/// Each IR change type has an ID.
+enum class TrackID {
+  UseSet,
+};
+
+#ifndef NDEBUG
+static const char *trackIDToStr(TrackID ID) {
+  switch (ID) {
+  case TrackID::UseSet:
+    return "UseSet";
+  }
+  llvm_unreachable("Unimplemented ID");
+}
+#endif // NDEBUG
+
+class SandboxIRTracker;
+
+/// The base class for IR Change classes.
+class IRChangeBase {
+protected:
+#ifndef NDEBUG
+  unsigned Idx = 0;
+#endif
+  const TrackID ID;
+  SandboxIRTracker &Parent;
+
+public:
+  IRChangeBase(TrackID ID, SandboxIRTracker &Parent);
+  TrackID getTrackID() const { return ID; }
+  /// This runs when changes get reverted.
+  virtual void revert() = 0;
+  /// This runs when changes get accepted.
+  virtual void accept() = 0;
+  virtual ~IRChangeBase() = default;
+#ifndef NDEBUG
+  void dumpCommon(raw_ostream &OS) const {
+    OS << Idx << ". " << trackIDToStr(ID);
+  }
+  virtual void dump(raw_ostream &OS) const = 0;
+  LLVM_DUMP_METHOD virtual void dump() const = 0;
+#endif
+};
+
+/// Change the source Value of a sandboxir::Use.
+class UseSet : public IRChangeBase {
+  Use U;
+  Value *OrigV = nullptr;
+
+public:
+  UseSet(const Use &U, SandboxIRTracker &Tracker)
+      : IRChangeBase(TrackID::UseSet, Tracker), U(U), OrigV(U.get()) {}
+  // For isa<> etc.
+  static bool classof(const IRChangeBase *Other) {
+    return Other->getTrackID() == TrackID::UseSet;
+  }
+  void revert() final { U.set(OrigV); }
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final { dumpCommon(OS); }
+  LLVM_DUMP_METHOD void dump() const final;
+  friend raw_ostream &operator<<(raw_ostream &OS, const UseSet &C) {
+    C.dump(OS);
+    return OS;
+  }
+#endif
+};
+
+/// The tracker collects all the change objects and implements the main API for
+/// saving / reverting / accepting.
+class SandboxIRTracker {
+public:
+  enum class TrackerState {
+    Disabled, ///> Tracking is disabled
+    Record,   ///> Tracking changes
+    Revert,   ///> Undoing changes
+    Accept,   ///> Accepting changes
+  };
+
+private:
+  /// The list of changes that are being tracked.
+  SmallVector<std::unique_ptr<IRChangeBase>> Changes;
----------------
aeubanks wrote:

I wonder if this could be a `BumpPtrAllocator` with a separate vector of `IRChangeBase *` pointing to all the values in it. Easy to apply in (reverse) order and call destructors of them all. Would help to not have so many dynamic allocations. But we can optimize later if this contributes to compile time.

https://github.com/llvm/llvm-project/pull/99238