From c55b74f7fe9b0b0bc911ab9b3cbc2656adbdc862 Mon Sep 17 00:00:00 2001
From: Marcel Ullrich <ullrich@cs.uni-saarland.de>
Date: Wed, 14 Sep 2022 11:53:18 +0200
Subject: [PATCH] pre prep phase

---
 dialects/tool/tool.cpp          |  2 +-
 thorin/pass/optimize.cpp        |  4 +++-
 thorin/pass/pipelinebuilder.cpp | 24 ++++++++++++++++++++++++
 thorin/pass/pipelinebuilder.h   |  6 ++++++
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/dialects/tool/tool.cpp b/dialects/tool/tool.cpp
index ba75465187..e600ff31e4 100644
--- a/dialects/tool/tool.cpp
+++ b/dialects/tool/tool.cpp
@@ -14,7 +14,7 @@ using namespace thorin;
 extern "C" THORIN_EXPORT thorin::DialectInfo thorin_get_dialect_info() {
     return {"tool",
             [](thorin::PipelineBuilder& builder) {
-                builder.extend_codegen_prep_phase([](PassMan& man) { man.add<thorin::tool::SetFilter>(); });
+                builder.extend_opt_prep_phase2([](PassMan& man) { man.add<thorin::tool::SetFilter>(); });
             },
             nullptr, [](Normalizers& normalizers) { tool::register_normalizers(normalizers); }};
 }
diff --git a/thorin/pass/optimize.cpp b/thorin/pass/optimize.cpp
index e189a48bf6..686e7faae8 100644
--- a/thorin/pass/optimize.cpp
+++ b/thorin/pass/optimize.cpp
@@ -11,13 +11,15 @@ namespace thorin {
 
 void optimize(World& world, PipelineBuilder& builder) {
     Pipeline pipe(world);
+    pipe.add<PassManPhase>(builder.opt_prep_phase1(world));
     pipe.add<Scalerize>();
     pipe.add<EtaRed>();
     pipe.add<TailRecElim>();
+    pipe.add<PassManPhase>(builder.opt_prep_phase2(world));
     pipe.add<PassManPhase>(builder.opt_phase(world));
     pipe.add<LamSpec>();
     pipe.add<PassManPhase>(builder.codegen_prep_phase(world));
-    pipe.add<PassManPhase>(builder.opt_phase2(world));
+    // pipe.add<PassManPhase>(builder.opt_phase2(world));
     pipe.run();
 }
 
diff --git a/thorin/pass/pipelinebuilder.cpp b/thorin/pass/pipelinebuilder.cpp
index 3180bcfa0b..e57482f903 100644
--- a/thorin/pass/pipelinebuilder.cpp
+++ b/thorin/pass/pipelinebuilder.cpp
@@ -23,6 +23,14 @@ void PipelineBuilder::extend_codegen_prep_phase(std::function<void(PassMan&)> ex
     codegen_prep_phase_extensions_.push_back(extension);
 }
 
+void PipelineBuilder::extend_opt_prep_phase1(std::function<void(PassMan&)> extension) {
+    opt_prep_phase1_extensions_.push_back(extension);
+}
+
+void PipelineBuilder::extend_opt_prep_phase2(std::function<void(PassMan&)> extension) {
+    opt_prep_phase2_extensions_.push_back(extension);
+}
+
 std::unique_ptr<PassMan> PipelineBuilder::opt_phase2(World& world) {
     auto man = std::make_unique<PassMan>(world);
 
@@ -61,4 +69,20 @@ std::unique_ptr<PassMan> PipelineBuilder::codegen_prep_phase(World& world) {
     return man;
 }
 
+std::unique_ptr<PassMan> PipelineBuilder::opt_prep_phase1(World& world) {
+    auto man = std::make_unique<PassMan>(world);
+
+    for (const auto& ext : opt_prep_phase1_extensions_) ext(*man);
+
+    return man;
+}
+
+std::unique_ptr<PassMan> PipelineBuilder::opt_prep_phase2(World& world) {
+    auto man = std::make_unique<PassMan>(world);
+
+    for (const auto& ext : opt_prep_phase2_extensions_) ext(*man);
+
+    return man;
+}
+
 } // namespace thorin
diff --git a/thorin/pass/pipelinebuilder.h b/thorin/pass/pipelinebuilder.h
index 8e470b5275..54c5f964b8 100644
--- a/thorin/pass/pipelinebuilder.h
+++ b/thorin/pass/pipelinebuilder.h
@@ -13,14 +13,20 @@ class PipelineBuilder {
 
     void extend_opt_phase(std::function<void(PassMan&)>);
     void extend_codegen_prep_phase(std::function<void(PassMan&)>);
+    void extend_opt_prep_phase1(std::function<void(PassMan&)>);
+    void extend_opt_prep_phase2(std::function<void(PassMan&)>);
 
     std::unique_ptr<PassMan> opt_phase(World& world);
     std::unique_ptr<PassMan> opt_phase2(World& world);
     std::unique_ptr<PassMan> codegen_prep_phase(World& world);
+    std::unique_ptr<PassMan> opt_prep_phase1(World& world);
+    std::unique_ptr<PassMan> opt_prep_phase2(World& world);
 
 private:
     std::vector<std::function<void(PassMan&)>> opt_phase_extensions_;
     std::vector<std::function<void(PassMan&)>> codegen_prep_phase_extensions_;
+    std::vector<std::function<void(PassMan&)>> opt_prep_phase1_extensions_;
+    std::vector<std::function<void(PassMan&)>> opt_prep_phase2_extensions_;
 };
 
 } // namespace thorin