1 files changed, 57 insertions, 0 deletions
diff --git a/llvm/unittests/IR/PassManagerTest.cpp b/llvm/unittests/IR/PassManagerTest.cpp
index c2ac863260a..47fb8b70f1b 100644
--- a/llvm/unittests/IR/PassManagerTest.cpp
+++ b/llvm/unittests/IR/PassManagerTest.cpp
@@ -331,4 +331,61 @@ TEST_F(PassManagerTest, Basic) {
 
   EXPECT_EQ(1, ModuleAnalysisRuns);
 }
+
+// A customized pass manager that passes extra arguments through the
+// infrastructure.
+typedef AnalysisManager<Function, int> CustomizedAnalysisManager;
+typedef PassManager<Function, CustomizedAnalysisManager, int, int &>
+    CustomizedPassManager;
+
+class CustomizedAnalysis : public AnalysisInfoMixin<CustomizedAnalysis> {
+public:
+  struct Result {
+    Result(int I) : I(I) {}
+    int I;
+  };
+
+  Result run(Function &F, CustomizedAnalysisManager &AM, int I) {
+    return Result(I);
+  }
+
+private:
+  friend AnalysisInfoMixin<CustomizedAnalysis>;
+  static char PassID;
+};
+
+char CustomizedAnalysis::PassID;
+
+struct CustomizedPass : PassInfoMixin<CustomizedPass> {
+  std::function<void(CustomizedAnalysis::Result &, int &)> Callback;
+
+  template <typename CallbackT>
+  CustomizedPass(CallbackT Callback) : Callback(Callback) {}
+
+  PreservedAnalyses run(Function &F, CustomizedAnalysisManager &AM, int I,
+                        int &O) {
+    Callback(AM.getResult<CustomizedAnalysis>(F, I), O);
+    return PreservedAnalyses::none();
+  }
+};
+
+TEST_F(PassManagerTest, CustomizedPassManagerArgs) {
+  CustomizedAnalysisManager AM;
+  AM.registerPass([&] { return CustomizedAnalysis(); });
+
+  CustomizedPassManager PM;
+
+  // Add an instance of the customized pass that just accumulates the input
+  // after it is round-tripped through the analysis.
+  int Result = 0;
+  PM.addPass(
+      CustomizedPass([](CustomizedAnalysis::Result &R, int &O) { O += R.I; }));
+
+  // Run this over every function with the input of 42.
+  for (Function &F : *M)
+    PM.run(F, AM, 42, Result);
+
+  // And ensure that we accumulated the correct result.
+  EXPECT_EQ(42 * (int)M->size(), Result);
+}
 }