1 files changed, 142 insertions, 0 deletions
diff --git a/include/llvm/Analysis/AssumptionCache.h b/include/llvm/Analysis/AssumptionCache.h
new file mode 100644
index 000000000000..b129e6796328
--- /dev/null
+++ b/include/llvm/Analysis/AssumptionCache.h
@@ -0,0 +1,142 @@
+//===- llvm/Analysis/AssumptionCache.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_ASSUMPTIONCACHE_H
+#define LLVM_ANALYSIS_ASSUMPTIONCACHE_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.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 {
+
+/// \brief A cache of @llvm.assume calls within a function.
+///
+/// This cache provides fast lookup of assumptions within a function by caching
+/// them and amortizing the cost of scanning for them across all queries. The
+/// cache is also conservatively self-updating so that it will never return
+/// incorrect results about a function even as the function is being mutated.
+/// However, flushing the cache and rebuilding it (or explicitly updating it)
+/// may allow it to discover new assumptions.
+class AssumptionCache {
+  /// \brief The function for which this cache is handling assumptions.
+  ///
+  /// We track this to lazily populate our assumptions.
+  Function &F;
+
+  /// \brief Vector of weak value handles to calls of the @llvm.assume
+  /// intrinsic.
+  SmallVector<WeakVH, 4> AssumeHandles;
+
+  /// \brief Flag tracking whether we have scanned the function yet.
+  ///
+  /// We want to be as lazy about this as possible, and so we scan the function
+  /// at the last moment.
+  bool Scanned;
+
+  /// \brief Scan the function for assumptions and add them to the cache.
+  void scanFunction();
+
+public:
+  /// \brief Construct an AssumptionCache from a function by scanning all of
+  /// its instructions.
+  AssumptionCache(Function &F) : F(F), Scanned(false) {}
+
+  /// \brief Add an @llvm.assume intrinsic to this function's cache.
+  ///
+  /// The call passed in must be an instruction within this fuction and must
+  /// not already be in the cache.
+  void registerAssumption(CallInst *CI);
+
+  /// \brief Clear the cache of @llvm.assume intrinsics for a function.
+  ///
+  /// It will be re-scanned the next time it is requested.
+  void clear() {
+    AssumeHandles.clear();
+    Scanned = false;
+  }
+
+  /// \brief Access the list of assumption handles currently tracked for this
+  /// fuction.
+  ///
+  /// Note that these produce weak handles that may be null. The caller must
+  /// handle that case.
+  /// FIXME: We should replace this with pointee_iterator<filter_iterator<...>>
+  /// when we can write that to filter out the null values. Then caller code
+  /// will become simpler.
+  MutableArrayRef<WeakVH> assumptions() {
+    if (!Scanned)
+      scanFunction();
+    return AssumeHandles;
+  }
+};
+
+/// \brief An immutable pass that tracks lazily created \c AssumptionCache
+/// objects.
+///
+/// This is essentially a workaround for the legacy pass manager's weaknesses
+/// which associates each assumption cache with Function and clears it if the
+/// function is deleted. The nature of the AssumptionCache is that it is not
+/// invalidated by any changes to the function body and so this is sufficient
+/// to be conservatively correct.
+class AssumptionCacheTracker : 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 {
+    AssumptionCacheTracker *ACT;
+    void deleted() override;
+
+  public:
+    typedef DenseMapInfo<Value *> DMI;
+
+    FunctionCallbackVH(Value *V, AssumptionCacheTracker *ACT = nullptr)
+        : CallbackVH(V), ACT(ACT) {}
+  };
+
+  friend FunctionCallbackVH;
+
+  typedef DenseMap<FunctionCallbackVH, std::unique_ptr<AssumptionCache>,
+                   FunctionCallbackVH::DMI> FunctionCallsMap;
+  FunctionCallsMap AssumptionCaches;
+
+public:
+  /// \brief Get the cached assumptions for a function.
+  ///
+  /// If no assumptions are cached, this will scan the function. Otherwise, the
+  /// existing cache will be returned.
+  AssumptionCache &getAssumptionCache(Function &F);
+
+  AssumptionCacheTracker();
+  ~AssumptionCacheTracker();
+
+  void releaseMemory() override { AssumptionCaches.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