From bb999d44a78b5e901ca96a914648b6f4ec00ff54 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas@google.com>
Date: Tue, 10 Sep 2024 14:04:40 -0700
Subject: [PATCH] [SandboxIR][PassRegistry] Parse pipeline string

This patch implements a simple version of the pipeline parsing function.
It currently only handles a single FPM and adds function passes to it.
---
 llvm/include/llvm/SandboxIR/PassManager.h |  4 +++
 llvm/lib/SandboxIR/PassManager.cpp        | 32 +++++++++++++++++++++++
 llvm/unittests/SandboxIR/PassTest.cpp     | 31 ++++++++++++++++++++++
 3 files changed, 67 insertions(+)

diff --git a/llvm/include/llvm/SandboxIR/PassManager.h b/llvm/include/llvm/SandboxIR/PassManager.h
index 5e250641f3b3f..2cc669a966e0b 100644
--- a/llvm/include/llvm/SandboxIR/PassManager.h
+++ b/llvm/include/llvm/SandboxIR/PassManager.h
@@ -72,6 +72,7 @@ class PassRegistry {
   DenseMap<StringRef, Pass *> NameToPassMap;
 
 public:
+  static constexpr const char PassDelimToken = ',';
   PassRegistry() = default;
   /// Registers \p PassPtr and takes ownership.
   Pass &registerPass(std::unique_ptr<Pass> &&PassPtr) {
@@ -85,6 +86,9 @@ class PassRegistry {
     auto It = NameToPassMap.find(Name);
     return It != NameToPassMap.end() ? It->second : nullptr;
   }
+  /// Creates a pass pipeline and returns the first pass manager.
+  FunctionPassManager &parseAndCreatePassPipeline(StringRef Pipeline);
+
 #ifndef NDEBUG
   void print(raw_ostream &OS) const {
     for (const auto &PassPtr : Passes)
diff --git a/llvm/lib/SandboxIR/PassManager.cpp b/llvm/lib/SandboxIR/PassManager.cpp
index 2dd19e74734db..4abd39b28e87a 100644
--- a/llvm/lib/SandboxIR/PassManager.cpp
+++ b/llvm/lib/SandboxIR/PassManager.cpp
@@ -20,6 +20,38 @@ bool FunctionPassManager::runOnFunction(Function &F) {
   // TODO: Check ChangeAll against hashes before/after.
   return Change;
 }
+
+FunctionPassManager &
+PassRegistry::parseAndCreatePassPipeline(StringRef Pipeline) {
+  static constexpr const char EndToken = '\0';
+  // Add EndToken to the end to ease parsing.
+  std::string PipelineStr = std::string(Pipeline) + EndToken;
+  int FlagBeginIdx = 0;
+  // Start with a FunctionPassManager.
+  auto &InitialPM = static_cast<FunctionPassManager &>(
+      registerPass(std::make_unique<FunctionPassManager>("init-fpm")));
+
+  for (auto [Idx, C] : enumerate(PipelineStr)) {
+    // Keep moving Idx until we find the end of the pass name.
+    bool FoundDelim = C == EndToken || C == PassDelimToken;
+    if (!FoundDelim)
+      continue;
+    unsigned Sz = Idx - FlagBeginIdx;
+    std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
+    FlagBeginIdx = Idx + 1;
+
+    // Get the pass that corresponds to PassName and add it to the pass manager.
+    auto *Pass = getPassByName(PassName);
+    if (Pass == nullptr) {
+      errs() << "Pass '" << PassName << "' not registered!\n";
+      exit(1);
+    }
+    // TODO: This is safe for now, but would require proper upcasting once we
+    // add more Pass sub-classes.
+    InitialPM.addPass(static_cast<FunctionPass *>(Pass));
+  }
+  return InitialPM;
+}
 #ifndef NDEBUG
 void PassRegistry::dump() const {
   print(dbgs());
diff --git a/llvm/unittests/SandboxIR/PassTest.cpp b/llvm/unittests/SandboxIR/PassTest.cpp
index 3517f0e32b1bb..ed226d5765586 100644
--- a/llvm/unittests/SandboxIR/PassTest.cpp
+++ b/llvm/unittests/SandboxIR/PassTest.cpp
@@ -162,3 +162,34 @@ TEST_F(PassTest, PassRegistry) {
   EXPECT_EQ(Buff, "test-pass1\ntest-pass2\n");
 #endif // NDEBUG
 }
+
+TEST_F(PassTest, ParsePassPipeline) {
+  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;
+  Registry.registerPass(std::make_unique<TestPass1>());
+  Registry.registerPass(std::make_unique<TestPass2>());
+
+  auto &FPM =
+      Registry.parseAndCreatePassPipeline("test-pass1,test-pass2,test-pass1");
+#ifndef NDEBUG
+  std::string Buff;
+  llvm::raw_string_ostream SS(Buff);
+  FPM.print(SS);
+  EXPECT_EQ(Buff, "init-fpm(test-pass1,test-pass2,test-pass1)");
+#endif // NDEBUG
+
+  EXPECT_DEATH(Registry.parseAndCreatePassPipeline("bad-pass-name"),
+               ".*not registered.*");
+  EXPECT_DEATH(Registry.parseAndCreatePassPipeline(""), ".*not registered.*");
+  EXPECT_DEATH(Registry.parseAndCreatePassPipeline(","), ".*not registered.*");
+}