[SandboxIR] Pass registry (#108084)

This patch implements a simple Pass Registry class, which takes
ownership of the passes registered with it and provides an interface to
get the pass pointer by its name.

Author: vporpo

llvm/include/llvm/SandboxIR/PassManager.h

@@ -18,6 +18,7 @@
 #ifndef LLVM_SANDBOXIR_PASSMANAGER_H
 #define LLVM_SANDBOXIR_PASSMANAGER_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/SandboxIR/Pass.h"
 #include "llvm/Support/Debug.h"
@@ -65,6 +66,34 @@ class FunctionPassManager final
   bool runOnFunction(Function &F) final;
 };
 
+/// Owns the passes and provides an API to get a pass by its name.
+class PassRegistry {
+  SmallVector<std::unique_ptr<Pass>, 8> Passes;
+  DenseMap<StringRef, Pass *> NameToPassMap;
+
+public:
+  PassRegistry() = default;
+  /// Registers \p PassPtr and takes ownership.
+  Pass &registerPass(std::unique_ptr<Pass> &&PassPtr) {
+    auto &PassRef = *PassPtr.get();
+    NameToPassMap[PassRef.getName()] = &PassRef;
+    Passes.push_back(std::move(PassPtr));
+    return PassRef;
+  }
+  /// \Returns the pass with name \p Name, or null if not registered.
+  Pass *getPassByName(StringRef Name) const {
+    auto It = NameToPassMap.find(Name);
+    return It != NameToPassMap.end() ? It->second : nullptr;
+  }
+#ifndef NDEBUG
+  void print(raw_ostream &OS) const {
+    for (const auto &PassPtr : Passes)
+      OS << PassPtr->getName() << "\n";
+  }
+  LLVM_DUMP_METHOD void dump() const;
+#endif
+};
+
 } // namespace llvm::sandboxir
 
 #endif // LLVM_SANDBOXIR_PASSMANAGER_H

llvm/lib/SandboxIR/Pass.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/SandboxIR/Pass.h"
+#include "llvm/SandboxIR/PassManager.h"
 #include "llvm/Support/Debug.h"
 
 using namespace llvm::sandboxir;

llvm/lib/SandboxIR/PassManager.cpp

@@ -20,3 +20,9 @@ bool FunctionPassManager::runOnFunction(Function &F) {
   // TODO: Check ChangeAll against hashes before/after.
   return Change;
 }
+#ifndef NDEBUG
+void PassRegistry::dump() const {
+  print(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG

llvm/unittests/SandboxIR/PassTest.cpp

@@ -133,3 +133,32 @@ define void @foo() {
   EXPECT_EQ(Buff, "test-fpm(test-pass1,test-pass2)");
 #endif // NDEBUG
 }
+
+TEST_F(PassTest, PassRegistry) {
+  class TestPass1 final : public FunctionPass {
+  public:
+    TestPass1() : FunctionPass("test-pass1") {}
+    bool runOnFunction(Function &F) final { return false; }
+  };
+  class TestPass2 final : public FunctionPass {
+  public:
+    TestPass2() : FunctionPass("test-pass2") {}
+    bool runOnFunction(Function &F) final { return false; }
+  };
+
+  PassRegistry Registry;
+  auto &TP1 = Registry.registerPass(std::make_unique<TestPass1>());
+  auto &TP2 = Registry.registerPass(std::make_unique<TestPass2>());
+
+  // Check getPassByName().
+  EXPECT_EQ(Registry.getPassByName("test-pass1"), &TP1);
+  EXPECT_EQ(Registry.getPassByName("test-pass2"), &TP2);
+
+#ifndef NDEBUG
+  // Check print().
+  std::string Buff;
+  llvm::raw_string_ostream SS(Buff);
+  Registry.print(SS);
+  EXPECT_EQ(Buff, "test-pass1\ntest-pass2\n");
+#endif // NDEBUG
+}