1 files changed, 128 insertions, 0 deletions
diff --git a/include/llvm/Analysis/AssumptionTracker.h b/include/llvm/Analysis/AssumptionTracker.h
new file mode 100644
index 0000000..5a050a8
--- /dev/null
+++ b/include/llvm/Analysis/AssumptionTracker.h
@@ -0,0 +1,128 @@
+//===- llvm/Analysis/AssumptionTracker.h - Track @llvm.assume ---*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a pass that keeps track of @llvm.assume intrinsics in
+// the functions of a module (allowing assumptions within any function to be
+// found cheaply by other parts of the optimizer).
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_ASSUMPTIONTRACKER_H
+#define LLVM_ANALYSIS_ASSUMPTIONTRACKER_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/ValueHandle.h"
+#include "llvm/Pass.h"
+#include <memory>
+
+namespace llvm {
+
+/// An immutable pass that tracks @llvm.assume intrinsics in a module.
+class AssumptionTracker : public ImmutablePass {
+  /// A callback value handle applied to function objects, which we use to
+  /// delete our cache of intrinsics for a function when it is deleted.
+  class FunctionCallbackVH : public CallbackVH {
+    AssumptionTracker *AT;
+    void deleted() override;
+
+    public:
+      typedef DenseMapInfo<Value *> DMI;
+
+      FunctionCallbackVH(Value *V, AssumptionTracker *AT = nullptr)
+        : CallbackVH(V), AT(AT) {}
+  };
+
+  /// A callback value handle applied to call instructions, which keeps
+  /// track of the call's parent function so that we can remove a
+  /// assumption intrinsic call from our cache when the instruction is
+  /// deleted.
+  class CallCallbackVH : public CallbackVH {
+    AssumptionTracker *AT;
+    void deleted() override;
+
+    // We store the function here because we need it to lookup the set
+    // containing this handle when the underlying CallInst is being deleted.
+    Function *F;
+
+    public:
+      typedef DenseMapInfo<Instruction *> DMI;
+
+      CallCallbackVH(Instruction *I, AssumptionTracker *AT = nullptr)
+        : CallbackVH(I), AT(AT), F(nullptr) {
+        if (I != DMI::getEmptyKey() && I != DMI::getTombstoneKey())
+          F = I->getParent()->getParent();
+      }
+
+      operator CallInst*() const {
+        Value *V = getValPtr();
+        if (V == DMI::getEmptyKey() || V == DMI::getTombstoneKey())
+          return reinterpret_cast<CallInst*>(V);
+
+        return cast<CallInst>(V);
+      }
+
+      CallInst *operator->() const { return cast<CallInst>(getValPtr()); }
+      CallInst &operator*() const { return *cast<CallInst>(getValPtr()); }
+  };
+
+  friend FunctionCallbackVH;
+  friend CallCallbackVH;
+
+  // FIXME: SmallSet might be better here, but it currently has no iterators.
+  typedef DenseSet<CallCallbackVH, CallCallbackVH::DMI> CallHandleSet;
+  typedef DenseMap<FunctionCallbackVH, std::unique_ptr<CallHandleSet>,
+                   FunctionCallbackVH::DMI> FunctionCallsMap;
+  FunctionCallsMap CachedAssumeCalls;
+
+  /// Scan the provided function for @llvm.assume intrinsic calls. Returns an
+  /// iterator to the set for this function in the CachedAssumeCalls map.
+  FunctionCallsMap::iterator scanFunction(Function *F);
+
+public:
+  /// Remove the cache of @llvm.assume intrinsics for the given function.
+  void forgetCachedAssumptions(Function *F);
+
+  /// Add an @llvm.assume intrinsic to the cache for its parent function.
+  void registerAssumption(CallInst *CI);
+
+  typedef CallHandleSet::iterator assumption_iterator;
+  typedef iterator_range<assumption_iterator> assumption_range;
+
+  inline assumption_range assumptions(Function *F) {
+    FunctionCallsMap::iterator I = CachedAssumeCalls.find_as(F);
+    if (I == CachedAssumeCalls.end()) {
+      I = scanFunction(F);
+    }
+
+    return assumption_range(I->second->begin(), I->second->end());
+  }
+
+  AssumptionTracker();
+  ~AssumptionTracker();
+
+  void releaseMemory() override {
+    CachedAssumeCalls.shrink_and_clear();
+  }
+
+  void verifyAnalysis() const override;
+  bool doFinalization(Module &) override {
+    verifyAnalysis();
+    return false;
+  }
+
+  static char ID; // Pass identification, replacement for typeid
+};
+
+} // end namespace llvm
+
+#endif