From 55140fbbcd6527a3fe3636dda2078fe3652558ec Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks@google.com>
Date: Thu, 16 Mar 2023 09:57:35 -0700
Subject: [PATCH] [StandardInstrumentations] Check that module analyses are
 properly invalidated

Followup to D146003/D146160

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D146238
---
 llvm/lib/Passes/StandardInstrumentations.cpp | 33 ++++++++++++++++++++
 llvm/unittests/IR/PassManagerTest.cpp        | 45 +++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index af97517..f1ff58d 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -1067,6 +1067,23 @@ struct PreservedFunctionHashAnalysis
 
 AnalysisKey PreservedFunctionHashAnalysis::Key;
 
+struct PreservedModuleHashAnalysis
+    : public AnalysisInfoMixin<PreservedModuleHashAnalysis> {
+  static AnalysisKey Key;
+
+  struct ModuleHash {
+    uint64_t Hash;
+  };
+
+  using Result = ModuleHash;
+
+  Result run(Module &F, ModuleAnalysisManager &FAM) {
+    return Result{StructuralHash(F)};
+  }
+};
+
+AnalysisKey PreservedModuleHashAnalysis::Key;
+
 bool PreservedCFGCheckerInstrumentation::CFG::invalidate(
     Function &F, const PreservedAnalyses &PA,
     FunctionAnalysisManager::Invalidator &) {
@@ -1106,6 +1123,7 @@ void PreservedCFGCheckerInstrumentation::registerCallbacks(
     if (!Registered) {
       FAM.registerPass([&] { return PreservedCFGCheckerAnalysis(); });
       FAM.registerPass([&] { return PreservedFunctionHashAnalysis(); });
+      MAM.registerPass([&] { return PreservedModuleHashAnalysis(); });
       Registered = true;
     }
 
@@ -1114,6 +1132,11 @@ void PreservedCFGCheckerInstrumentation::registerCallbacks(
       FAM.getResult<PreservedCFGCheckerAnalysis>(*F);
       FAM.getResult<PreservedFunctionHashAnalysis>(*F);
     }
+
+    if (auto *MaybeM = any_cast<const Module *>(&IR)) {
+      Module &M = **const_cast<Module **>(MaybeM);
+      MAM.getResult<PreservedModuleHashAnalysis>(M);
+    }
   });
 
   PIC.registerAfterPassInvalidatedCallback(
@@ -1169,6 +1192,16 @@ void PreservedCFGCheckerInstrumentation::registerCallbacks(
         CheckCFG(P, F->getName(), *GraphBefore,
                  CFG(F, /* TrackBBLifetime */ false));
     }
+    if (auto *MaybeM = any_cast<const Module *>(&IR)) {
+      Module &M = **const_cast<Module **>(MaybeM);
+      if (auto *HashBefore =
+              MAM.getCachedResult<PreservedModuleHashAnalysis>(M)) {
+        if (HashBefore->Hash != StructuralHash(M)) {
+          report_fatal_error(formatv(
+              "Module changed by {0} without invalidating analyses", P));
+        }
+      }
+    }
   });
 }
 
diff --git a/llvm/unittests/IR/PassManagerTest.cpp b/llvm/unittests/IR/PassManagerTest.cpp
index dda1d0c..a648716 100644
--- a/llvm/unittests/IR/PassManagerTest.cpp
+++ b/llvm/unittests/IR/PassManagerTest.cpp
@@ -985,6 +985,7 @@ TEST_F(PassManagerTest, FunctionPassMissedFunctionAnalysisInvalidation) {
   MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
   MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
 
   FunctionPassManager FPM;
   FPM.addPass(WrongFunctionPass());
@@ -998,7 +999,10 @@ struct WrongModulePass : PassInfoMixin<WrongModulePass> {
     for (Function &F : M)
       F.getEntryBlock().begin()->eraseFromParent();
 
-    return PreservedAnalyses::all();
+    PreservedAnalyses PA;
+    PA.preserveSet<AllAnalysesOn<Function>>();
+    PA.preserve<FunctionAnalysisManagerModuleProxy>();
+    return PA;
   }
   static StringRef name() { return "WrongModulePass"; }
 };
@@ -1018,6 +1022,7 @@ TEST_F(PassManagerTest, ModulePassMissedFunctionAnalysisInvalidation) {
   MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
   MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
 
   ModulePassManager MPM;
   MPM.addPass(WrongModulePass());
@@ -1027,5 +1032,43 @@ TEST_F(PassManagerTest, ModulePassMissedFunctionAnalysisInvalidation) {
       "Function @foo changed by WrongModulePass without invalidating analyses");
 }
 
+struct WrongModulePass2 : PassInfoMixin<WrongModulePass2> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    for (Function &F : M)
+      F.getEntryBlock().begin()->eraseFromParent();
+
+    PreservedAnalyses PA;
+    PA.preserveSet<AllAnalysesOn<Module>>();
+    PA.abandon<FunctionAnalysisManagerModuleProxy>();
+    return PA;
+  }
+  static StringRef name() { return "WrongModulePass2"; }
+};
+
+TEST_F(PassManagerTest, ModulePassMissedModuleAnalysisInvalidation) {
+  LLVMContext Context;
+  auto M = parseIR(Context, "define void @foo() {\n"
+                            "  %a = add i32 0, 0\n"
+                            "  ret void\n"
+                            "}\n");
+
+  FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
+  PassInstrumentationCallbacks PIC;
+  StandardInstrumentations SI(M->getContext(), /*DebugLogging*/ false);
+  SI.registerCallbacks(PIC, &MAM);
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+  FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
+
+  ModulePassManager MPM;
+  MPM.addPass(WrongModulePass2());
+
+  EXPECT_DEATH(
+      MPM.run(*M, MAM),
+      "Module changed by WrongModulePass2 without invalidating analyses");
+}
+
 #endif
 }
-- 
2.7.4