[llvm] [SandboxIR] Pass registry (PR #108084)

[via llvm-commits]

https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/108084

>From 8c1e712f0a71f68ef93cd87e611c7162f3d95f82 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Mon, 9 Sep 2024 15:27:23 -0700
Subject: [PATCH] [SandboxIR] Pass registry

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.
---
 llvm/include/llvm/SandboxIR/PassManager.h | 29 +++++++++++++++++++++++
 llvm/lib/SandboxIR/Pass.cpp               |  1 +
 llvm/lib/SandboxIR/PassManager.cpp        |  6 +++++
 llvm/unittests/SandboxIR/PassTest.cpp     | 29 +++++++++++++++++++++++
 4 files changed, 65 insertions(+)

diff --git a/llvm/include/llvm/SandboxIR/PassManager.h b/llvm/include/llvm/SandboxIR/PassManager.h
index cb321fe699a562..5e250641f3b3f6 100644
--- a/llvm/include/llvm/SandboxIR/PassManager.h
+++ b/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
diff --git a/llvm/lib/SandboxIR/Pass.cpp b/llvm/lib/SandboxIR/Pass.cpp
index 64e1b609a9f49d..c6ec1aec48b19d 100644
--- a/llvm/lib/SandboxIR/Pass.cpp
+++ b/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;
diff --git a/llvm/lib/SandboxIR/PassManager.cpp b/llvm/lib/SandboxIR/PassManager.cpp
index d10f3926f7bcdb..2dd19e74734db8 100644
--- a/llvm/lib/SandboxIR/PassManager.cpp
+++ b/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
diff --git a/llvm/unittests/SandboxIR/PassTest.cpp b/llvm/unittests/SandboxIR/PassTest.cpp
index 8e080128b15b34..3517f0e32b1bbe 100644
--- a/llvm/unittests/SandboxIR/PassTest.cpp
+++ b/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
+}