1 files changed, 62 insertions, 0 deletions
diff --git a/unittests/IR/LegacyPassManagerTest.cpp b/unittests/IR/LegacyPassManagerTest.cpp
index 0ff2ec717597..9f5f431a5585 100644
--- a/unittests/IR/LegacyPassManagerTest.cpp
+++ b/unittests/IR/LegacyPassManagerTest.cpp
@@ -26,6 +26,7 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/OptBisect.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
@@ -396,6 +397,67 @@ namespace llvm {
       delete M;
     }
 
+    // Skips or runs optional passes.
+    struct CustomOptPassGate : public OptPassGate {
+      bool Skip;
+      CustomOptPassGate(bool Skip) : Skip(Skip) { }
+      bool shouldRunPass(const Pass *P, const Module &U) { return !Skip; }
+    };
+
+    // Optional module pass.
+    struct ModuleOpt: public ModulePass {
+      char run = 0;
+      static char ID;
+      ModuleOpt() : ModulePass(ID) { }
+      bool runOnModule(Module &M) override {
+        if (!skipModule(M))
+          run++;
+        return false;
+      }
+    };
+    char ModuleOpt::ID=0;
+
+    TEST(PassManager, CustomOptPassGate) {
+      LLVMContext Context0;
+      LLVMContext Context1;
+      LLVMContext Context2;
+      CustomOptPassGate SkipOptionalPasses(true);
+      CustomOptPassGate RunOptionalPasses(false);
+
+      Module M0("custom-opt-bisect", Context0);
+      Module M1("custom-opt-bisect", Context1);
+      Module M2("custom-opt-bisect2", Context2);
+      struct ModuleOpt *mOpt0 = new ModuleOpt();
+      struct ModuleOpt *mOpt1 = new ModuleOpt();
+      struct ModuleOpt *mOpt2 = new ModuleOpt();
+
+      mOpt0->run = mOpt1->run = mOpt2->run = 0;
+
+      legacy::PassManager Passes0;
+      legacy::PassManager Passes1;
+      legacy::PassManager Passes2;
+
+      Passes0.add(mOpt0);
+      Passes1.add(mOpt1);
+      Passes2.add(mOpt2);
+
+      Context1.setOptPassGate(SkipOptionalPasses);
+      Context2.setOptPassGate(RunOptionalPasses);
+
+      Passes0.run(M0);
+      Passes1.run(M1);
+      Passes2.run(M2);
+
+      // By default optional passes are run.
+      EXPECT_EQ(1, mOpt0->run);
+
+      // The first context skips optional passes.
+      EXPECT_EQ(0, mOpt1->run);
+
+      // The second context runs optional passes.
+      EXPECT_EQ(1, mOpt2->run);
+    }
+
     Module *makeLLVMModule(LLVMContext &Context) {
       // Module Construction
       Module *mod = new Module("test-mem", Context);