[MustExecute] Use the annotation style printer

As suggested in the original review (https://reviews.llvm.org/D44524), use an annotation style printer instead.

Note: The switch from -analyze to -disable-output in tests was driven by the fact that seems to be the idiomatic style used in annoation passes.  I tried to keep both working, but the old style pass API for printers really doesn't make this easy.  It invokes (runOnFunction, print(Module)) repeatedly.  I decided the extra state wasn't worth it given the old pass manager is going away soonish anyway.
llvm-svn: 328015

1 files changed, 54 insertions, 33 deletions

diff --git a/llvm/lib/Analysis/MustExecute.cpp b/llvm/lib/Analysis/MustExecute.cpp
index 3d98a61..f741b35 100644
--- a/llvm/lib/Analysis/MustExecute.cpp
+++ b/llvm/lib/Analysis/MustExecute.cpp
@@ -10,19 +10,19 @@
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/Passes.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormattedStream.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 using namespace llvm;
 
 namespace {
   struct MustExecutePrinter : public FunctionPass {
-    DenseMap<Value*, SmallVector<Loop*, 4> > MustExec;
-    SmallVector<Value *, 4> Ordering;
 
     static char ID; // Pass identification, replacement for typeid
     MustExecutePrinter() : FunctionPass(ID) {
@@ -34,11 +34,6 @@ namespace {
       AU.addRequired<LoopInfoWrapperPass>();
     }
     bool runOnFunction(Function &F) override;
-    void print(raw_ostream &OS, const Module * = nullptr) const override;
-    void releaseMemory() override {
-      MustExec.clear();
-      Ordering.clear();
-    }
   };
 }
 
@@ -54,7 +49,7 @@ FunctionPass *llvm::createMustExecutePrinter() {
   return new MustExecutePrinter();
 }
 
-bool isMustExecuteIn(Instruction &I, Loop *L, DominatorTree *DT) {
+bool isMustExecuteIn(const Instruction &I, Loop *L, DominatorTree *DT) {
   // TODO: move loop specific code to analysis
   //LoopSafetyInfo LSI;
   //computeLoopSafetyInfo(&LSI, L);
@@ -62,41 +57,67 @@ bool isMustExecuteIn(Instruction &I, Loop *L, DominatorTree *DT) {
   return isGuaranteedToExecuteForEveryIteration(&I, L);
 }
 
-bool MustExecutePrinter::runOnFunction(Function &F) {
-  auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
-  auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-  for (auto &I: instructions(F)) {
-    Loop *L = LI.getLoopFor(I.getParent());
-    while (L) {
-      if (isMustExecuteIn(I, L, &DT)) {
-        if (!MustExec.count(&I))
-          Ordering.push_back(&I);
-        MustExec[&I].push_back(L);
-      }
-      L = L->getParentLoop();
-    };
+/// \brief An assembly annotator class to print must execute information in
+/// comments.
+class MustExecuteAnnotatedWriter : public AssemblyAnnotationWriter {
+  DenseMap<const Value*, SmallVector<Loop*, 4> > MustExec;
+
+public:
+  MustExecuteAnnotatedWriter(const Function &F,
+                             DominatorTree &DT, LoopInfo &LI) {
+    for (auto &I: instructions(F)) {
+      Loop *L = LI.getLoopFor(I.getParent());
+      while (L) {
+        if (isMustExecuteIn(I, L, &DT)) {
+          MustExec[&I].push_back(L);
+        }
+        L = L->getParentLoop();
+      };
+    }
   }
-  return false;
-}
+  MustExecuteAnnotatedWriter(const Module &M,
+                             DominatorTree &DT, LoopInfo &LI) {
+    for (auto &F : M)
+    for (auto &I: instructions(F)) {
+      Loop *L = LI.getLoopFor(I.getParent());
+      while (L) {
+        if (isMustExecuteIn(I, L, &DT)) {
+          MustExec[&I].push_back(L);
+        }
+        L = L->getParentLoop();
+      };
+    }
+  }
+
+
+  void printInfoComment(const Value &V, formatted_raw_ostream &OS) override {  
+    if (!MustExec.count(&V))
+      return;
 
-void MustExecutePrinter::print(raw_ostream &OS, const Module *M) const {
-  OS << "The following are guaranteed to execute (for the respective loops):\n";
-  for (Value *V: Ordering) {
-    V->printAsOperand(OS);
-    auto NumLoops = MustExec.lookup(V).size();
+    const auto &Loops = MustExec.lookup(&V);
+    const auto NumLoops = Loops.size();
     if (NumLoops > 1)
-      OS << "\t(mustexec in " << NumLoops << " loops: ";
+      OS << " ; (mustexec in " << NumLoops << " loops: ";
     else
-      OS << "\t(mustexec in: ";
+      OS << " ; (mustexec in: ";
     
     bool first = true;
-    for (const Loop *L : MustExec.lookup(V)) {
+    for (const Loop *L : Loops) {
       if (!first)
         OS << ", ";
       first = false;
       OS << L->getHeader()->getName();
     }
-    OS << ")\n";
+    OS << ")";
   }
-  OS << "\n";
+};
+
+bool MustExecutePrinter::runOnFunction(Function &F) {
+  auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+  auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+
+  MustExecuteAnnotatedWriter Writer(F, DT, LI);
+  F.print(dbgs(), &Writer);
+  
+  return false;
 }