1 files changed, 174 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index fa5dbc23014..f4e2ed4df98 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -16,6 +16,7 @@
 #include "llvm/Transforms/IPO/Attributor.h"
 
 #include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -31,6 +32,9 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Local.h"
+
 #include <cassert>
 
 using namespace llvm;
@@ -1411,6 +1415,173 @@ ChangeStatus AANoAliasReturned::updateImpl(Attributor &A) {
   return ChangeStatus::UNCHANGED;
 }
 
+/// -------------------AAIsDead Function Attribute-----------------------
+
+struct AAIsDeadFunction : AAIsDead, BooleanState {
+
+  AAIsDeadFunction(Function &F, InformationCache &InfoCache)
+      : AAIsDead(F, InfoCache) {}
+
+  /// See AbstractAttribute::getState()
+  /// {
+  AbstractState &getState() override { return *this; }
+  const AbstractState &getState() const override { return *this; }
+  /// }
+
+  /// See AbstractAttribute::getManifestPosition().
+  ManifestPosition getManifestPosition() const override { return MP_FUNCTION; }
+
+  void initialize(Attributor &A) override {
+    Function &F = getAnchorScope();
+
+    ToBeExploredPaths.insert(&(F.getEntryBlock().front()));
+    AssumedLiveBlocks.insert(&(F.getEntryBlock()));
+    for (size_t i = 0; i < ToBeExploredPaths.size(); ++i)
+      explorePath(A, ToBeExploredPaths[i]);
+  }
+
+  /// Explores new instructions starting from \p I. If instruction is dead, stop
+  /// and return true if it discovered a new instruction.
+  bool explorePath(Attributor &A, Instruction *I);
+
+  const std::string getAsStr() const override {
+    return "LiveBBs(" + std::to_string(AssumedLiveBlocks.size()) + "/" +
+           std::to_string(getAnchorScope().size()) + ")";
+  }
+
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+    assert(getState().isValidState() &&
+           "Attempted to manifest an invalid state!");
+
+    ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
+
+    for (Instruction *I : NoReturnCalls) {
+      BasicBlock *BB = I->getParent();
+
+      /// Invoke is replaced with a call and unreachable is placed after it.
+      if (auto *II = dyn_cast<InvokeInst>(I)) {
+        changeToCall(II);
+        changeToUnreachable(BB->getTerminator(), /* UseLLVMTrap */ false);
+        LLVM_DEBUG(dbgs() << "[AAIsDead] Replaced invoke with call inst\n");
+        continue;
+      }
+
+      SplitBlock(BB, I->getNextNode());
+      changeToUnreachable(BB->getTerminator(), /* UseLLVMTrap */ false);
+      HasChanged = ChangeStatus::CHANGED;
+    }
+
+    return HasChanged;
+  }
+
+  /// See AbstractAttribute::updateImpl(...).
+  ChangeStatus updateImpl(Attributor &A) override;
+
+  /// See AAIsDead::isAssumedDead().
+  bool isAssumedDead(BasicBlock *BB) const override {
+    if (!getAssumed())
+      return false;
+    return !AssumedLiveBlocks.count(BB);
+  }
+
+  /// See AAIsDead::isKnownDead().
+  bool isKnownDead(BasicBlock *BB) const override {
+    if (!getKnown())
+      return false;
+    return !AssumedLiveBlocks.count(BB);
+  }
+
+  /// Collection of to be explored paths.
+  SmallSetVector<Instruction *, 8> ToBeExploredPaths;
+
+  /// Collection of all assumed live BasicBlocks.
+  DenseSet<BasicBlock *> AssumedLiveBlocks;
+
+  /// Collection of calls with noreturn attribute, assumed or knwon.
+  SmallSetVector<Instruction *, 4> NoReturnCalls;
+};
+
+bool AAIsDeadFunction::explorePath(Attributor &A, Instruction *I) {
+  BasicBlock *BB = I->getParent();
+
+  while (I) {
+    ImmutableCallSite ICS(I);
+
+    if (ICS) {
+      auto *NoReturnAA = A.getAAFor<AANoReturn>(*this, *I);
+
+      if (NoReturnAA && NoReturnAA->isAssumedNoReturn()) {
+        if (!NoReturnCalls.insert(I))
+          // If I is already in the NoReturnCalls set, then it stayed noreturn
+          // and we didn't discover any new instructions.
+          return false;
+
+        // Discovered new noreturn call, return true to indicate that I is not
+        // noreturn anymore and should be deleted from NoReturnCalls.
+        return true;
+      }
+
+      if (ICS.hasFnAttr(Attribute::NoReturn)) {
+        if(!NoReturnCalls.insert(I))
+          return false;
+
+        return true;
+      }
+    }
+
+    I = I->getNextNode();
+  }
+
+  // get new paths (reachable blocks).
+  for (BasicBlock *SuccBB : successors(BB)) {
+    Instruction *Inst = &(SuccBB->front());
+    AssumedLiveBlocks.insert(SuccBB);
+    ToBeExploredPaths.insert(Inst);
+  }
+
+  return true;
+}
+
+ChangeStatus AAIsDeadFunction::updateImpl(Attributor &A) {
+  Function &F = getAnchorScope();
+
+  // Temporary collection to iterate over existing noreturn instructions. This
+  // will alow easier modification of NoReturnCalls collection
+  SmallVector<Instruction *, 8> NoReturnChanged;
+  ChangeStatus Status = ChangeStatus::UNCHANGED;
+
+  for (Instruction *I : NoReturnCalls)
+    NoReturnChanged.push_back(I);
+
+  for (Instruction *I : NoReturnChanged) {
+    size_t Size = ToBeExploredPaths.size();
+
+    // Still noreturn.
+    if (!explorePath(A, I))
+      continue;
+
+    NoReturnCalls.remove(I);
+
+    // No new paths.
+    if (Size == ToBeExploredPaths.size())
+      continue;
+
+    // At least one new path.
+    Status = ChangeStatus::CHANGED;
+
+    // explore new paths.
+    while (Size != ToBeExploredPaths.size())
+      explorePath(A, ToBeExploredPaths[Size++]);
+  }
+
+  LLVM_DEBUG(dbgs() << "[AAIsDead] AssumedLiveBlocks: "
+                    << AssumedLiveBlocks.size()
+                    << "Total number of blocks: " << F.size() << "\n");
+
+  return Status;
+}
+
 /// ----------------------------------------------------------------------------
 ///                               Attributor
 /// ----------------------------------------------------------------------------
@@ -1627,6 +1798,9 @@ void Attributor::identifyDefaultAbstractAttributes(
   // Every function might be "will-return".
   registerAA(*new AAWillReturnFunction(F, InfoCache));
 
+  // Check for dead BasicBlocks in every function.
+  registerAA(*new AAIsDeadFunction(F, InfoCache));
+
   // Walk all instructions to find more attribute opportunities and also
   // interesting instructions that might be queried by abstract attributes
   // during their initialization or update.