llvm/unittests/MIR/PassBuilderCallbacksTest.cpp - third_party/llvm-project - Git at Google

 //===- unittests/MIR/PassBuilderCallbacksTest.cpp - PB Callback Tests -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//

 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Testing/Support/Error.h"
 #include <functional>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <llvm/ADT/Any.h>
 #include <llvm/AsmParser/Parser.h>
 #include <llvm/CodeGen/MIRParser/MIRParser.h>
 #include <llvm/CodeGen/MachineFunction.h>
 #include <llvm/CodeGen/MachineModuleInfo.h>
 #include <llvm/CodeGen/MachinePassManager.h>
 #include <llvm/IR/LLVMContext.h>
 #include <llvm/IR/PassInstrumentation.h>
 #include <llvm/IR/PassManager.h>
 #include <llvm/Passes/PassBuilder.h>
 #include <llvm/Support/Regex.h>
 #include <llvm/Support/SourceMgr.h>
 #include <llvm/Support/TargetSelect.h>

 using namespace llvm;

 namespace {
 using testing::_;
 using testing::AnyNumber;
 using testing::DoAll;
 using testing::Not;
 using testing::Return;
 using testing::WithArgs;

 StringRef MIRString = R"MIR(
 --- |
   define void @test() {
     ret void
   }
 ...
 ---
 name:            test
 body:             |
   bb.0 (%ir-block.0):
     RET64
 ...
 )MIR";

 /// Helper for HasName matcher that returns getName both for IRUnit and
 /// for IRUnit pointer wrapper into llvm::Any (wrapped by PassInstrumentation).
 template <typename IRUnitT> std::string getName(const IRUnitT &IR) {
   return std::string(IR.getName());
 }

 template <> std::string getName(const StringRef &name) {
   return std::string(name);
 }

 template <> std::string getName(const Any &WrappedIR) {
   if (const auto *const *M = llvm::any_cast<const Module *>(&WrappedIR))
     return (*M)->getName().str();
   if (const auto *const *F = llvm::any_cast<const Function *>(&WrappedIR))
     return (*F)->getName().str();
   if (const auto *const *MF =
           llvm::any_cast<const MachineFunction *>(&WrappedIR))
     return (*MF)->getName().str();
   return "<UNKNOWN>";
 }
 /// Define a custom matcher for objects which support a 'getName' method.
 ///
 /// LLVM often has IR objects or analysis objects which expose a name
 /// and in tests it is convenient to match these by name for readability.
 /// Usually, this name is either a StringRef or a plain std::string. This
 /// matcher supports any type exposing a getName() method of this form whose
 /// return value is compatible with an std::ostream. For StringRef, this uses
 /// the shift operator defined above.
 ///
 /// It should be used as:
 ///
 ///   HasName("my_function")
 ///
 /// No namespace or other qualification is required.
 MATCHER_P(HasName, Name, "") {
   *result_listener << "has name '" << getName(arg) << "'";
   return Name == getName(arg);
 }

 MATCHER_P(HasNameRegex, Name, "") {
   *result_listener << "has name '" << getName(arg) << "'";
   llvm::Regex r(Name);
   return r.match(getName(arg));
 }

 struct MockPassInstrumentationCallbacks {
   PassInstrumentationCallbacks Callbacks;

   MockPassInstrumentationCallbacks() {
     ON_CALL(*this, runBeforePass(_, _)).WillByDefault(Return(true));
   }
   MOCK_METHOD2(runBeforePass, bool(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runBeforeSkippedPass, void(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runBeforeNonSkippedPass, void(StringRef PassID, llvm::Any));
   MOCK_METHOD3(runAfterPass,
                void(StringRef PassID, llvm::Any, const PreservedAnalyses &PA));
   MOCK_METHOD2(runAfterPassInvalidated,
                void(StringRef PassID, const PreservedAnalyses &PA));
   MOCK_METHOD2(runBeforeAnalysis, void(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runAfterAnalysis, void(StringRef PassID, llvm::Any));

   void registerPassInstrumentation() {
     Callbacks.registerShouldRunOptionalPassCallback(
         [this](StringRef P, llvm::Any IR) {
           return this->runBeforePass(P, IR);
         });
     Callbacks.registerBeforeSkippedPassCallback(
         [this](StringRef P, llvm::Any IR) {
           this->runBeforeSkippedPass(P, IR);
         });
     Callbacks.registerBeforeNonSkippedPassCallback(
         [this](StringRef P, llvm::Any IR) {
           this->runBeforeNonSkippedPass(P, IR);
         });
     Callbacks.registerAfterPassCallback(
         [this](StringRef P, llvm::Any IR, const PreservedAnalyses &PA) {
           this->runAfterPass(P, IR, PA);
         });
     Callbacks.registerAfterPassInvalidatedCallback(
         [this](StringRef P, const PreservedAnalyses &PA) {
           this->runAfterPassInvalidated(P, PA);
         });
     Callbacks.registerBeforeAnalysisCallback([this](StringRef P, llvm::Any IR) {
       return this->runBeforeAnalysis(P, IR);
     });
     Callbacks.registerAfterAnalysisCallback(
         [this](StringRef P, llvm::Any IR) { this->runAfterAnalysis(P, IR); });
   }

   void ignoreNonMockPassInstrumentation(StringRef IRName) {
     // Generic EXPECT_CALLs are needed to match instrumentation on unimportant
     // parts of a pipeline that we do not care about (e.g. various passes added
     // by default by PassBuilder - Verifier pass etc).
     // Make sure to avoid ignoring Mock passes/analysis, we definitely want
     // to check these explicitly.
     EXPECT_CALL(*this,
                 runBeforePass(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(
         *this, runBeforeSkippedPass(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this, runBeforeNonSkippedPass(Not(HasNameRegex("Mock")),
                                                HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runAfterPass(Not(HasNameRegex("Mock")), HasName(IRName), _))
         .Times(AnyNumber());
     EXPECT_CALL(*this, runBeforeAnalysis(HasNameRegex("MachineModuleAnalysis"),
                                          HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runBeforeAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this, runAfterAnalysis(HasNameRegex("MachineModuleAnalysis"),
                                         HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runAfterAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
   }
 };

 template <typename DerivedT> class MockAnalysisHandleBase {
 public:
   class Analysis : public AnalysisInfoMixin<Analysis> {
     friend AnalysisInfoMixin<Analysis>;
     friend MockAnalysisHandleBase;
     static AnalysisKey Key;

     DerivedT *Handle;

     Analysis(DerivedT &Handle) : Handle(&Handle) {
       static_assert(std::is_base_of<MockAnalysisHandleBase, DerivedT>::value,
                     "Must pass the derived type to this template!");
     }

   public:
     class Result {
       friend MockAnalysisHandleBase;

       DerivedT *Handle;

       Result(DerivedT &Handle) : Handle(&Handle) {}

     public:
       // Forward invalidation events to the mock handle.
       bool invalidate(MachineFunction &IR, const PreservedAnalyses &PA,
                       MachineFunctionAnalysisManager::Invalidator &Inv) {
         return Handle->invalidate(IR, PA, Inv);
       }
     };

     Result run(MachineFunction &IR, MachineFunctionAnalysisManager::Base &AM) {
       return Handle->run(IR, AM);
     }
   };

   Analysis getAnalysis() { return Analysis(static_cast<DerivedT &>(*this)); }
   typename Analysis::Result getResult() {
     return typename Analysis::Result(static_cast<DerivedT &>(*this));
   }
   static StringRef getName() { return llvm::getTypeName<DerivedT>(); }

 protected:
   // FIXME: MSVC seems unable to handle a lambda argument to Invoke from within
   // the template, so we use a boring static function.
   static bool
   invalidateCallback(MachineFunction &IR, const PreservedAnalyses &PA,
                      MachineFunctionAnalysisManager::Invalidator &Inv) {
     auto PAC = PA.template getChecker<Analysis>();
     return !PAC.preserved() &&
            !PAC.template preservedSet<AllAnalysesOn<MachineFunction>>();
   }

   /// Derived classes should call this in their constructor to set up default
   /// mock actions. (We can't do this in our constructor because this has to
   /// run after the DerivedT is constructed.)
   void setDefaults() {
     ON_CALL(static_cast<DerivedT &>(*this), run(_, _))
         .WillByDefault(Return(this->getResult()));
     ON_CALL(static_cast<DerivedT &>(*this), invalidate(_, _, _))
         .WillByDefault(&invalidateCallback);
   }
 };

 template <typename DerivedT> class MockPassHandleBase {
 public:
   class Pass : public MachinePassInfoMixin<Pass> {
     friend MockPassHandleBase;

     DerivedT *Handle;

     Pass(DerivedT &Handle) : Handle(&Handle) {
       static_assert(std::is_base_of<MockPassHandleBase, DerivedT>::value,
                     "Must pass the derived type to this template!");
     }

   public:
     static MachinePassKey Key;
     PreservedAnalyses run(MachineFunction &IR,
                           MachineFunctionAnalysisManager::Base &AM) {
       return Handle->run(IR, AM);
     }
   };

   static StringRef getName() { return llvm::getTypeName<DerivedT>(); }

   Pass getPass() { return Pass(static_cast<DerivedT &>(*this)); }

 protected:
   /// Derived classes should call this in their constructor to set up default
   /// mock actions. (We can't do this in our constructor because this has to
   /// run after the DerivedT is constructed.)
   void setDefaults() {
     ON_CALL(static_cast<DerivedT &>(*this), run(_, _))
         .WillByDefault(Return(PreservedAnalyses::all()));
   }
 };

 struct MockAnalysisHandle : public MockAnalysisHandleBase<MockAnalysisHandle> {
   MOCK_METHOD2(run, Analysis::Result(MachineFunction &,
                                      MachineFunctionAnalysisManager::Base &));

   MOCK_METHOD3(invalidate, bool(MachineFunction &, const PreservedAnalyses &,
                                 MachineFunctionAnalysisManager::Invalidator &));

   MockAnalysisHandle() { setDefaults(); }
 };

 template <typename DerivedT>
 MachinePassKey MockPassHandleBase<DerivedT>::Pass::Key;

 template <typename DerivedT>
 AnalysisKey MockAnalysisHandleBase<DerivedT>::Analysis::Key;

 class MockPassHandle : public MockPassHandleBase<MockPassHandle> {
 public:
   MOCK_METHOD2(run, PreservedAnalyses(MachineFunction &,
                                       MachineFunctionAnalysisManager::Base &));

   MockPassHandle() { setDefaults(); }
 };

 class MachineFunctionCallbacksTest : public testing::Test {
 protected:
   static void SetUpTestCase() {
     InitializeAllTargetInfos();
     InitializeAllTargets();
     InitializeAllTargetMCs();
   }

   std::unique_ptr<LLVMTargetMachine> TM;

   LLVMContext Context;
   std::unique_ptr<Module> M;
   std::unique_ptr<MIRParser> MIR;

   MockPassInstrumentationCallbacks CallbacksHandle;

   PassBuilder PB;
   ModulePassManager PM;
   MachineFunctionPassManager MFPM;
   FunctionAnalysisManager FAM;
   ModuleAnalysisManager AM;
   MachineFunctionAnalysisManager MFAM;

   MockPassHandle PassHandle;
   MockAnalysisHandle AnalysisHandle;

   std::unique_ptr<Module> parseMIR(const TargetMachine &TM, StringRef MIRCode,
                                    MachineModuleInfo &MMI) {
     SMDiagnostic Diagnostic;
     std::unique_ptr<MemoryBuffer> MBuffer = MemoryBuffer::getMemBuffer(MIRCode);
     MIR = createMIRParser(std::move(MBuffer), Context);
     if (!MIR)
       return nullptr;

     std::unique_ptr<Module> Mod = MIR->parseIRModule();
     if (!Mod)
       return nullptr;

     Mod->setDataLayout(TM.createDataLayout());

     if (MIR->parseMachineFunctions(*Mod, MMI)) {
       M.reset();
       return nullptr;
     }
     return Mod;
   }

   static PreservedAnalyses
   getAnalysisResult(MachineFunction &U,
                     MachineFunctionAnalysisManager::Base &AM) {
     auto &MFAM = static_cast<MachineFunctionAnalysisManager &>(AM);
     MFAM.getResult<MockAnalysisHandle::Analysis>(U);
     return PreservedAnalyses::all();
   }

   void SetUp() override {
     std::string Error;
     auto TripleName = "x86_64-pc-linux-gnu";
     auto *T = TargetRegistry::lookupTarget(TripleName, Error);
     if (!T)
       GTEST_SKIP();
     TM = std::unique_ptr<LLVMTargetMachine>(
         static_cast<LLVMTargetMachine *>(T->createTargetMachine(
             TripleName, "", "", TargetOptions(), std::nullopt)));
     if (!TM)
       GTEST_SKIP();
     MachineModuleInfo MMI(TM.get());
     M = parseMIR(*TM, MIRString, MMI);
     AM.registerPass([&] { return MachineModuleAnalysis(TM.get()); });
   }

   MachineFunctionCallbacksTest()
       : CallbacksHandle(), PB(nullptr, PipelineTuningOptions(), std::nullopt,
                               &CallbacksHandle.Callbacks),
         PM(), FAM(), AM(), MFAM(FAM, AM) {

     EXPECT_TRUE(&CallbacksHandle.Callbacks ==
                 PB.getPassInstrumentationCallbacks());

     /// Register a callback for analysis registration.
     ///
     /// The callback is a function taking a reference to an AnalyisManager
     /// object. When called, the callee gets to register its own analyses with
     /// this PassBuilder instance.
     PB.registerAnalysisRegistrationCallback(
         [this](MachineFunctionAnalysisManager &AM) {
           // Register our mock analysis
           AM.registerPass([this] { return AnalysisHandle.getAnalysis(); });
         });

     /// Register a callback for pipeline parsing.
     ///
     /// During parsing of a textual pipeline, the PassBuilder will call these
     /// callbacks for each encountered pass name that it does not know. This
     /// includes both simple pass names as well as names of sub-pipelines. In
     /// the latter case, the InnerPipeline is not empty.
     PB.registerPipelineParsingCallback(
         [this](StringRef Name, MachineFunctionPassManager &PM) {
           if (parseAnalysisUtilityPasses<MockAnalysisHandle::Analysis>(
                   "test-analysis", Name, PM))
             return true;

           /// Parse the name of our pass mock handle
           if (Name == "test-transform") {
             MFPM.addPass(PassHandle.getPass());
             return true;
           }
           return false;
         });

     /// Register builtin analyses and cross-register the analysis proxies
     PB.registerModuleAnalyses(AM);
     PB.registerFunctionAnalyses(FAM);
     PB.registerMachineFunctionAnalyses(MFAM);
   }
 };

 TEST_F(MachineFunctionCallbacksTest, Passes) {
   EXPECT_CALL(AnalysisHandle, run(HasName("test"), _));
   EXPECT_CALL(PassHandle, run(HasName("test"), _)).WillOnce(&getAnalysisResult);

   StringRef PipelineText = "test-transform";
   ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
       << "Pipeline was: " << PipelineText;
   ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
 }

 TEST_F(MachineFunctionCallbacksTest, InstrumentedPasses) {
   CallbacksHandle.registerPassInstrumentation();
   // Non-mock instrumentation not specifically mentioned below can be ignored.
   CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
   CallbacksHandle.ignoreNonMockPassInstrumentation("test");
   CallbacksHandle.ignoreNonMockPassInstrumentation("");

   // PassInstrumentation calls should happen in-sequence, in the same order
   // as passes/analyses are scheduled.
   ::testing::Sequence PISequence;
   EXPECT_CALL(CallbacksHandle,
               runBeforePass(HasNameRegex("MockPassHandle"), HasName("test")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runBeforeNonSkippedPass(HasNameRegex("MockPassHandle"), HasName("test")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), HasName("test"), _))
       .InSequence(PISequence);

   EXPECT_CALL(AnalysisHandle, run(HasName("test"), _));
   EXPECT_CALL(PassHandle, run(HasName("test"), _)).WillOnce(&getAnalysisResult);

   StringRef PipelineText = "test-transform";
   ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
       << "Pipeline was: " << PipelineText;
   ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
 }

 TEST_F(MachineFunctionCallbacksTest, InstrumentedSkippedPasses) {
   CallbacksHandle.registerPassInstrumentation();
   // Non-mock instrumentation run here can safely be ignored.
   CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
   CallbacksHandle.ignoreNonMockPassInstrumentation("test");
   CallbacksHandle.ignoreNonMockPassInstrumentation("");

   // Skip the pass by returning false.
   EXPECT_CALL(CallbacksHandle,
               runBeforePass(HasNameRegex("MockPassHandle"), HasName("test")))
       .WillOnce(Return(false));

   EXPECT_CALL(
       CallbacksHandle,
       runBeforeSkippedPass(HasNameRegex("MockPassHandle"), HasName("test")))
       .Times(1);

   EXPECT_CALL(AnalysisHandle, run(HasName("test"), _)).Times(0);
   EXPECT_CALL(PassHandle, run(HasName("test"), _)).Times(0);

   // As the pass is skipped there is no afterPass, beforeAnalysis/afterAnalysis
   // as well.
   EXPECT_CALL(CallbacksHandle,
               runBeforeNonSkippedPass(HasNameRegex("MockPassHandle"), _))
       .Times(0);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), _, _))
       .Times(0);
   EXPECT_CALL(CallbacksHandle,
               runAfterPassInvalidated(HasNameRegex("MockPassHandle"), _))
       .Times(0);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), _, _))
       .Times(0);
   EXPECT_CALL(CallbacksHandle,
               runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
       .Times(0);
   EXPECT_CALL(CallbacksHandle,
               runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
       .Times(0);

   StringRef PipelineText = "test-transform";
   ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
       << "Pipeline was: " << PipelineText;
   ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
 }

 } // end anonymous namespace
	//===- unittests/MIR/PassBuilderCallbacksTest.cpp - PB Callback Tests -----===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//

	#include "llvm/MC/TargetRegistry.h"
	#include "llvm/Target/TargetMachine.h"
	#include "llvm/Testing/Support/Error.h"
	#include <functional>
	#include <gmock/gmock.h>
	#include <gtest/gtest.h>
	#include <llvm/ADT/Any.h>
	#include <llvm/AsmParser/Parser.h>
	#include <llvm/CodeGen/MIRParser/MIRParser.h>
	#include <llvm/CodeGen/MachineFunction.h>
	#include <llvm/CodeGen/MachineModuleInfo.h>
	#include <llvm/CodeGen/MachinePassManager.h>
	#include <llvm/IR/LLVMContext.h>
	#include <llvm/IR/PassInstrumentation.h>
	#include <llvm/IR/PassManager.h>
	#include <llvm/Passes/PassBuilder.h>
	#include <llvm/Support/Regex.h>
	#include <llvm/Support/SourceMgr.h>
	#include <llvm/Support/TargetSelect.h>

	using namespace llvm;

	namespace {
	using testing::_;
	using testing::AnyNumber;
	using testing::DoAll;
	using testing::Not;
	using testing::Return;
	using testing::WithArgs;

	StringRef MIRString = R"MIR(
	--- \|
	define void @test() {
	ret void
	}
	...
	---
	name: test
	body: \|
	bb.0 (%ir-block.0):
	RET64
	...
	)MIR";

	/// Helper for HasName matcher that returns getName both for IRUnit and
	/// for IRUnit pointer wrapper into llvm::Any (wrapped by PassInstrumentation).
	template <typename IRUnitT> std::string getName(const IRUnitT &IR) {
	return std::string(IR.getName());
	}

	template <> std::string getName(const StringRef &name) {
	return std::string(name);
	}

	template <> std::string getName(const Any &WrappedIR) {
	if (const auto const M = llvm::any_cast<const Module *>(&WrappedIR))
	return (*M)->getName().str();
	if (const auto const F = llvm::any_cast<const Function *>(&WrappedIR))
	return (*F)->getName().str();
	if (const auto const MF =
	llvm::any_cast<const MachineFunction *>(&WrappedIR))
	return (*MF)->getName().str();
	return "<UNKNOWN>";
	}
	/// Define a custom matcher for objects which support a 'getName' method.
	///
	/// LLVM often has IR objects or analysis objects which expose a name
	/// and in tests it is convenient to match these by name for readability.
	/// Usually, this name is either a StringRef or a plain std::string. This
	/// matcher supports any type exposing a getName() method of this form whose
	/// return value is compatible with an std::ostream. For StringRef, this uses
	/// the shift operator defined above.
	///
	/// It should be used as:
	///
	/// HasName("my_function")
	///
	/// No namespace or other qualification is required.
	MATCHER_P(HasName, Name, "") {
	*result_listener << "has name '" << getName(arg) << "'";
	return Name == getName(arg);
	}

	MATCHER_P(HasNameRegex, Name, "") {
	*result_listener << "has name '" << getName(arg) << "'";
	llvm::Regex r(Name);
	return r.match(getName(arg));
	}

	struct MockPassInstrumentationCallbacks {
	PassInstrumentationCallbacks Callbacks;

	MockPassInstrumentationCallbacks() {
	ON_CALL(*this, runBeforePass(_, _)).WillByDefault(Return(true));
	}
	MOCK_METHOD2(runBeforePass, bool(StringRef PassID, llvm::Any));
	MOCK_METHOD2(runBeforeSkippedPass, void(StringRef PassID, llvm::Any));
	MOCK_METHOD2(runBeforeNonSkippedPass, void(StringRef PassID, llvm::Any));
	MOCK_METHOD3(runAfterPass,
	void(StringRef PassID, llvm::Any, const PreservedAnalyses &PA));
	MOCK_METHOD2(runAfterPassInvalidated,
	void(StringRef PassID, const PreservedAnalyses &PA));
	MOCK_METHOD2(runBeforeAnalysis, void(StringRef PassID, llvm::Any));
	MOCK_METHOD2(runAfterAnalysis, void(StringRef PassID, llvm::Any));

	void registerPassInstrumentation() {
	Callbacks.registerShouldRunOptionalPassCallback(
	[this](StringRef P, llvm::Any IR) {
	return this->runBeforePass(P, IR);
	});
	Callbacks.registerBeforeSkippedPassCallback(
	[this](StringRef P, llvm::Any IR) {
	this->runBeforeSkippedPass(P, IR);
	});
	Callbacks.registerBeforeNonSkippedPassCallback(
	[this](StringRef P, llvm::Any IR) {
	this->runBeforeNonSkippedPass(P, IR);
	});
	Callbacks.registerAfterPassCallback(
	[this](StringRef P, llvm::Any IR, const PreservedAnalyses &PA) {
	this->runAfterPass(P, IR, PA);
	});
	Callbacks.registerAfterPassInvalidatedCallback(
	[this](StringRef P, const PreservedAnalyses &PA) {
	this->runAfterPassInvalidated(P, PA);
	});
	Callbacks.registerBeforeAnalysisCallback([this](StringRef P, llvm::Any IR) {
	return this->runBeforeAnalysis(P, IR);
	});
	Callbacks.registerAfterAnalysisCallback(
	[this](StringRef P, llvm::Any IR) { this->runAfterAnalysis(P, IR); });
	}

	void ignoreNonMockPassInstrumentation(StringRef IRName) {
	// Generic EXPECT_CALLs are needed to match instrumentation on unimportant
	// parts of a pipeline that we do not care about (e.g. various passes added
	// by default by PassBuilder - Verifier pass etc).
	// Make sure to avoid ignoring Mock passes/analysis, we definitely want
	// to check these explicitly.
	EXPECT_CALL(*this,
	runBeforePass(Not(HasNameRegex("Mock")), HasName(IRName)))
	.Times(AnyNumber());
	EXPECT_CALL(
	*this, runBeforeSkippedPass(Not(HasNameRegex("Mock")), HasName(IRName)))
	.Times(AnyNumber());
	EXPECT_CALL(*this, runBeforeNonSkippedPass(Not(HasNameRegex("Mock")),
	HasName(IRName)))
	.Times(AnyNumber());
	EXPECT_CALL(*this,
	runAfterPass(Not(HasNameRegex("Mock")), HasName(IRName), _))
	.Times(AnyNumber());
	EXPECT_CALL(*this, runBeforeAnalysis(HasNameRegex("MachineModuleAnalysis"),
	HasName(IRName)))
	.Times(AnyNumber());
	EXPECT_CALL(*this,
	runBeforeAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
	.Times(AnyNumber());
	EXPECT_CALL(*this, runAfterAnalysis(HasNameRegex("MachineModuleAnalysis"),
	HasName(IRName)))
	.Times(AnyNumber());
	EXPECT_CALL(*this,
	runAfterAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
	.Times(AnyNumber());
	}
	};

	template <typename DerivedT> class MockAnalysisHandleBase {
	public:
	class Analysis : public AnalysisInfoMixin<Analysis> {
	friend AnalysisInfoMixin<Analysis>;
	friend MockAnalysisHandleBase;
	static AnalysisKey Key;

	DerivedT *Handle;

	Analysis(DerivedT &Handle) : Handle(&Handle) {
	static_assert(std::is_base_of<MockAnalysisHandleBase, DerivedT>::value,
	"Must pass the derived type to this template!");
	}

	public:
	class Result {
	friend MockAnalysisHandleBase;

	DerivedT *Handle;

	Result(DerivedT &Handle) : Handle(&Handle) {}

	public:
	// Forward invalidation events to the mock handle.
	bool invalidate(MachineFunction &IR, const PreservedAnalyses &PA,
	MachineFunctionAnalysisManager::Invalidator &Inv) {
	return Handle->invalidate(IR, PA, Inv);
	}
	};

	Result run(MachineFunction &IR, MachineFunctionAnalysisManager::Base &AM) {
	return Handle->run(IR, AM);
	}
	};

	Analysis getAnalysis() { return Analysis(static_cast<DerivedT &>(*this)); }
	typename Analysis::Result getResult() {
	return typename Analysis::Result(static_cast<DerivedT &>(*this));
	}
	static StringRef getName() { return llvm::getTypeName<DerivedT>(); }

	protected:
	// FIXME: MSVC seems unable to handle a lambda argument to Invoke from within
	// the template, so we use a boring static function.
	static bool
	invalidateCallback(MachineFunction &IR, const PreservedAnalyses &PA,
	MachineFunctionAnalysisManager::Invalidator &Inv) {
	auto PAC = PA.template getChecker<Analysis>();
	return !PAC.preserved() &&
	!PAC.template preservedSet<AllAnalysesOn<MachineFunction>>();
	}

	/// Derived classes should call this in their constructor to set up default
	/// mock actions. (We can't do this in our constructor because this has to
	/// run after the DerivedT is constructed.)
	void setDefaults() {
	ON_CALL(static_cast<DerivedT &>(*this), run(_, _))
	.WillByDefault(Return(this->getResult()));
	ON_CALL(static_cast<DerivedT &>(*this), invalidate(_, _, _))
	.WillByDefault(&invalidateCallback);
	}
	};

	template <typename DerivedT> class MockPassHandleBase {
	public:
	class Pass : public MachinePassInfoMixin<Pass> {
	friend MockPassHandleBase;

	DerivedT *Handle;

	Pass(DerivedT &Handle) : Handle(&Handle) {
	static_assert(std::is_base_of<MockPassHandleBase, DerivedT>::value,
	"Must pass the derived type to this template!");
	}

	public:
	static MachinePassKey Key;
	PreservedAnalyses run(MachineFunction &IR,
	MachineFunctionAnalysisManager::Base &AM) {
	return Handle->run(IR, AM);
	}
	};

	static StringRef getName() { return llvm::getTypeName<DerivedT>(); }

	Pass getPass() { return Pass(static_cast<DerivedT &>(*this)); }

	protected:
	/// Derived classes should call this in their constructor to set up default
	/// mock actions. (We can't do this in our constructor because this has to
	/// run after the DerivedT is constructed.)
	void setDefaults() {
	ON_CALL(static_cast<DerivedT &>(*this), run(_, _))
	.WillByDefault(Return(PreservedAnalyses::all()));
	}
	};

	struct MockAnalysisHandle : public MockAnalysisHandleBase<MockAnalysisHandle> {
	MOCK_METHOD2(run, Analysis::Result(MachineFunction &,
	MachineFunctionAnalysisManager::Base &));

	MOCK_METHOD3(invalidate, bool(MachineFunction &, const PreservedAnalyses &,
	MachineFunctionAnalysisManager::Invalidator &));

	MockAnalysisHandle() { setDefaults(); }
	};

	template <typename DerivedT>
	MachinePassKey MockPassHandleBase<DerivedT>::Pass::Key;

	template <typename DerivedT>
	AnalysisKey MockAnalysisHandleBase<DerivedT>::Analysis::Key;

	class MockPassHandle : public MockPassHandleBase<MockPassHandle> {
	public:
	MOCK_METHOD2(run, PreservedAnalyses(MachineFunction &,
	MachineFunctionAnalysisManager::Base &));

	MockPassHandle() { setDefaults(); }
	};

	class MachineFunctionCallbacksTest : public testing::Test {
	protected:
	static void SetUpTestCase() {
	InitializeAllTargetInfos();
	InitializeAllTargets();
	InitializeAllTargetMCs();
	}

	std::unique_ptr<LLVMTargetMachine> TM;

	LLVMContext Context;
	std::unique_ptr<Module> M;
	std::unique_ptr<MIRParser> MIR;

	MockPassInstrumentationCallbacks CallbacksHandle;

	PassBuilder PB;
	ModulePassManager PM;
	MachineFunctionPassManager MFPM;
	FunctionAnalysisManager FAM;
	ModuleAnalysisManager AM;
	MachineFunctionAnalysisManager MFAM;

	MockPassHandle PassHandle;
	MockAnalysisHandle AnalysisHandle;

	std::unique_ptr<Module> parseMIR(const TargetMachine &TM, StringRef MIRCode,
	MachineModuleInfo &MMI) {
	SMDiagnostic Diagnostic;
	std::unique_ptr<MemoryBuffer> MBuffer = MemoryBuffer::getMemBuffer(MIRCode);
	MIR = createMIRParser(std::move(MBuffer), Context);
	if (!MIR)
	return nullptr;

	std::unique_ptr<Module> Mod = MIR->parseIRModule();
	if (!Mod)
	return nullptr;

	Mod->setDataLayout(TM.createDataLayout());

	if (MIR->parseMachineFunctions(*Mod, MMI)) {
	M.reset();
	return nullptr;
	}
	return Mod;
	}

	static PreservedAnalyses
	getAnalysisResult(MachineFunction &U,
	MachineFunctionAnalysisManager::Base &AM) {
	auto &MFAM = static_cast<MachineFunctionAnalysisManager &>(AM);
	MFAM.getResult<MockAnalysisHandle::Analysis>(U);
	return PreservedAnalyses::all();
	}

	void SetUp() override {
	std::string Error;
	auto TripleName = "x86_64-pc-linux-gnu";
	auto *T = TargetRegistry::lookupTarget(TripleName, Error);
	if (!T)
	GTEST_SKIP();
	TM = std::unique_ptr<LLVMTargetMachine>(
	static_cast<LLVMTargetMachine *>(T->createTargetMachine(
	TripleName, "", "", TargetOptions(), std::nullopt)));
	if (!TM)
	GTEST_SKIP();
	MachineModuleInfo MMI(TM.get());
	M = parseMIR(*TM, MIRString, MMI);
	AM.registerPass([&] { return MachineModuleAnalysis(TM.get()); });
	}

	MachineFunctionCallbacksTest()
	: CallbacksHandle(), PB(nullptr, PipelineTuningOptions(), std::nullopt,
	&CallbacksHandle.Callbacks),
	PM(), FAM(), AM(), MFAM(FAM, AM) {

	EXPECT_TRUE(&CallbacksHandle.Callbacks ==
	PB.getPassInstrumentationCallbacks());

	/// Register a callback for analysis registration.
	///
	/// The callback is a function taking a reference to an AnalyisManager
	/// object. When called, the callee gets to register its own analyses with
	/// this PassBuilder instance.
	PB.registerAnalysisRegistrationCallback(
	[this](MachineFunctionAnalysisManager &AM) {
	// Register our mock analysis
	AM.registerPass([this] { return AnalysisHandle.getAnalysis(); });
	});

	/// Register a callback for pipeline parsing.
	///
	/// During parsing of a textual pipeline, the PassBuilder will call these
	/// callbacks for each encountered pass name that it does not know. This
	/// includes both simple pass names as well as names of sub-pipelines. In
	/// the latter case, the InnerPipeline is not empty.
	PB.registerPipelineParsingCallback(
	[this](StringRef Name, MachineFunctionPassManager &PM) {
	if (parseAnalysisUtilityPasses<MockAnalysisHandle::Analysis>(
	"test-analysis", Name, PM))
	return true;

	/// Parse the name of our pass mock handle
	if (Name == "test-transform") {
	MFPM.addPass(PassHandle.getPass());
	return true;
	}
	return false;
	});

	/// Register builtin analyses and cross-register the analysis proxies
	PB.registerModuleAnalyses(AM);
	PB.registerFunctionAnalyses(FAM);
	PB.registerMachineFunctionAnalyses(MFAM);
	}
	};

	TEST_F(MachineFunctionCallbacksTest, Passes) {
	EXPECT_CALL(AnalysisHandle, run(HasName("test"), _));
	EXPECT_CALL(PassHandle, run(HasName("test"), _)).WillOnce(&getAnalysisResult);

	StringRef PipelineText = "test-transform";
	ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
	<< "Pipeline was: " << PipelineText;
	ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
	}

	TEST_F(MachineFunctionCallbacksTest, InstrumentedPasses) {
	CallbacksHandle.registerPassInstrumentation();
	// Non-mock instrumentation not specifically mentioned below can be ignored.
	CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
	CallbacksHandle.ignoreNonMockPassInstrumentation("test");
	CallbacksHandle.ignoreNonMockPassInstrumentation("");

	// PassInstrumentation calls should happen in-sequence, in the same order
	// as passes/analyses are scheduled.
	::testing::Sequence PISequence;
	EXPECT_CALL(CallbacksHandle,
	runBeforePass(HasNameRegex("MockPassHandle"), HasName("test")))
	.InSequence(PISequence);
	EXPECT_CALL(
	CallbacksHandle,
	runBeforeNonSkippedPass(HasNameRegex("MockPassHandle"), HasName("test")))
	.InSequence(PISequence);
	EXPECT_CALL(CallbacksHandle,
	runAfterPass(HasNameRegex("MockPassHandle"), HasName("test"), _))
	.InSequence(PISequence);

	EXPECT_CALL(AnalysisHandle, run(HasName("test"), _));
	EXPECT_CALL(PassHandle, run(HasName("test"), _)).WillOnce(&getAnalysisResult);

	StringRef PipelineText = "test-transform";
	ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
	<< "Pipeline was: " << PipelineText;
	ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
	}

	TEST_F(MachineFunctionCallbacksTest, InstrumentedSkippedPasses) {
	CallbacksHandle.registerPassInstrumentation();
	// Non-mock instrumentation run here can safely be ignored.
	CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
	CallbacksHandle.ignoreNonMockPassInstrumentation("test");
	CallbacksHandle.ignoreNonMockPassInstrumentation("");

	// Skip the pass by returning false.
	EXPECT_CALL(CallbacksHandle,
	runBeforePass(HasNameRegex("MockPassHandle"), HasName("test")))
	.WillOnce(Return(false));

	EXPECT_CALL(
	CallbacksHandle,
	runBeforeSkippedPass(HasNameRegex("MockPassHandle"), HasName("test")))
	.Times(1);

	EXPECT_CALL(AnalysisHandle, run(HasName("test"), _)).Times(0);
	EXPECT_CALL(PassHandle, run(HasName("test"), _)).Times(0);

	// As the pass is skipped there is no afterPass, beforeAnalysis/afterAnalysis
	// as well.
	EXPECT_CALL(CallbacksHandle,
	runBeforeNonSkippedPass(HasNameRegex("MockPassHandle"), _))
	.Times(0);
	EXPECT_CALL(CallbacksHandle,
	runAfterPass(HasNameRegex("MockPassHandle"), _, _))
	.Times(0);
	EXPECT_CALL(CallbacksHandle,
	runAfterPassInvalidated(HasNameRegex("MockPassHandle"), _))
	.Times(0);
	EXPECT_CALL(CallbacksHandle,
	runAfterPass(HasNameRegex("MockPassHandle"), _, _))
	.Times(0);
	EXPECT_CALL(CallbacksHandle,
	runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
	.Times(0);
	EXPECT_CALL(CallbacksHandle,
	runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
	.Times(0);

	StringRef PipelineText = "test-transform";
	ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
	<< "Pipeline was: " << PipelineText;
	ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
	}

	} // end anonymous namespace