AST: Per-request dependency maps

Just as with the result cache, instead of a single DenseMap with
type-erased AnyRequest keys, we can use per-request maps for a
nice performance improvement.

Author: slavapestov

include/swift/AST/Evaluator.h

@@ -323,6 +323,7 @@ class Evaluator {
            typename std::enable_if<!Request::hasExternalCache>::type* = nullptr>
   void clearCachedOutput(const Request &request) {
     cache.erase<Request>(request);
+    recorder.clearRequest<Request>(request);
   }
 
   /// Clear the cache stored within this evaluator.

include/swift/AST/EvaluatorDependencies.h

@@ -20,6 +20,7 @@
 
 #include "swift/AST/AnyRequest.h"
 #include "swift/AST/DependencyCollector.h"
+#include "swift/AST/RequestCache.h"
 #include "swift/Basic/NullablePtr.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
@@ -63,8 +64,7 @@ struct DependencyRecorder {
   /// downstream request as well. Note that uncached requests don't appear as
   /// keys in this map; their references are charged to the innermost cached
   /// active request.
-  llvm::DenseMap<AnyRequest, std::vector<DependencyCollector::Reference>>
-      requestReferences;
+  RequestReferences requestReferences;
 
   /// Stack of references from each cached active request. When evaluating a
   /// dependency sink request, we update the innermost set of references.
@@ -103,6 +103,10 @@ struct DependencyRecorder {
   void handleDependencySourceRequest(const Request &req,
                                      SourceFile *source);
 
+  /// Clear the recorded dependencies of a request, if any.
+  template<typename Request>
+  void clearRequest(const Request &req);
+
 private:
   /// Add an entry to the innermost set on the activeRequestReferences stack.
   /// Called from the DependencyCollector.
@@ -159,7 +163,7 @@ void evaluator::DependencyRecorder::endRequest(const Request &req) {
 
   // Finally, record the dependencies so we can replay them
   // later when the request is re-evaluated.
-  requestReferences.insert({AnyRequest(req), std::move(vec)});
+  requestReferences.insert<Request>(std::move(req), std::move(vec));
 }
 
 template<typename Request>
@@ -169,8 +173,8 @@ void evaluator::DependencyRecorder::replayCachedRequest(const Request &req) {
   if (activeRequestReferences.empty())
     return;
 
-  auto found = requestReferences.find_as(req);
-  if (found == requestReferences.end())
+  auto found = requestReferences.find_as<Request>(req);
+  if (found == requestReferences.end<Request>())
     return;
 
   activeRequestReferences.back().insert(found->second.begin(),
@@ -181,13 +185,19 @@ template<typename Request>
 void evaluator::DependencyRecorder::handleDependencySourceRequest(
     const Request &req,
     SourceFile *source) {
-  auto found = requestReferences.find_as(req);
-  if (found != requestReferences.end()) {
+  auto found = requestReferences.find_as<Request>(req);
+  if (found != requestReferences.end<Request>()) {
     fileReferences[source].insert(found->second.begin(),
                                   found->second.end());
   }
 }
 
+template<typename Request>
+void evaluator::DependencyRecorder::clearRequest(
+    const Request &req) {
+  requestReferences.erase(req);
+}
+
 } // end namespace evaluator
 
 } // end namespace swift

include/swift/AST/RequestCache.h

@@ -11,10 +11,11 @@
 //===----------------------------------------------------------------------===//
 //
 // This file defines data structures to efficiently support the request
-// evaluator's request caching.
+// evaluator's per-request caching and dependency tracking maps.
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/AST/DependencyCollector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 
@@ -279,6 +280,127 @@ class RequestCache {
   }
 };
 
+/// Type-erased wrapper for caching dependencies from a single type of request.
+class PerRequestReferences {
+  void *Storage;
+  std::function<void(void *)> Deleter;
+
+  PerRequestReferences(void *storage, std::function<void(void *)> deleter)
+      : Storage(storage), Deleter(deleter) {}
+
+public:
+  PerRequestReferences() : Storage(nullptr), Deleter([](void *) {}) {}
+  PerRequestReferences(PerRequestReferences &&other)
+      : Storage(other.Storage), Deleter(std::move(other.Deleter)) {
+    other.Storage = nullptr;
+  }
+
+  PerRequestReferences &operator=(PerRequestReferences &&other) {
+    if (&other != this) {
+      this->~PerRequestReferences();
+      new (this) PerRequestReferences(std::move(other));
+    }
+    return *this;
+  }
+
+  PerRequestReferences(const PerRequestReferences &) = delete;
+  PerRequestReferences &operator=(const PerRequestReferences &) = delete;
+
+  template <typename Request>
+  static PerRequestReferences makeEmpty() {
+    using Map =
+        llvm::DenseMap<RequestKey<Request>,
+                       std::vector<DependencyCollector::Reference>>;
+    return PerRequestReferences(new Map(),
+                                [](void *ptr) { delete static_cast<Map *>(ptr); });
+  }
+
+  template <typename Request>
+  llvm::DenseMap<RequestKey<Request>,
+                 std::vector<DependencyCollector::Reference>> *
+  get() const {
+    using Map =
+        llvm::DenseMap<RequestKey<Request>,
+                       std::vector<DependencyCollector::Reference>>;
+    assert(Storage);
+    return static_cast<Map *>(Storage);
+  }
+
+  bool isNull() const { return !Storage; }
+  ~PerRequestReferences() {
+    if (Storage)
+      Deleter(Storage);
+  }
+};
+
+/// Data structure for caching dependencies from requests. Sharded by the
+/// type ID zone and request kind, with a PerRequestReferences for each
+/// request kind.
+///
+/// Conceptually equivalent to DenseMap<AnyRequest, vector<Reference>>, but
+/// without type erasure overhead for keys.
+class RequestReferences {
+
+#define SWIFT_TYPEID_ZONE(Name, Id)                                            \
+  std::vector<PerRequestReferences> Name##ZoneRefs;                            \
+                                                                               \
+  template <                                                                   \
+      typename Request, typename ZoneTypes = TypeIDZoneTypes<Zone::Name>,      \
+      typename std::enable_if<TypeID<Request>::zone == Zone::Name>::type * =   \
+          nullptr>                                                             \
+  llvm::DenseMap<RequestKey<Request>,                                          \
+                 std::vector<DependencyCollector::Reference>> *                \
+  getRefs() {                                                                  \
+    auto &refs = Name##ZoneRefs;                                               \
+    if (refs.empty()) {                                                        \
+      refs.resize(ZoneTypes::Count);                                           \
+    }                                                                          \
+    auto idx = TypeID<Request>::localID;                                       \
+    if (refs[idx].isNull()) {                                                  \
+      refs[idx] = PerRequestReferences::makeEmpty<Request>();                  \
+    }                                                                          \
+    return refs[idx].template get<Request>();                                  \
+  }
+#include "swift/Basic/TypeIDZones.def"
+#undef SWIFT_TYPEID_ZONE
+
+public:
+  template <typename Request>
+  typename llvm::DenseMap<RequestKey<Request>,
+                          std::vector<DependencyCollector::Reference>>::const_iterator
+  find_as(const Request &req) {
+    auto *refs = getRefs<Request>();
+    return refs->find_as(req);
+  }
+
+  template <typename Request>
+  typename llvm::DenseMap<RequestKey<Request>,
+                          std::vector<DependencyCollector::Reference>>::const_iterator
+  end() {
+    auto *refs = getRefs<Request>();
+    return refs->end();
+  }
+
+  template <typename Request>
+  void insert(Request req, std::vector<DependencyCollector::Reference> val) {
+    auto *refs = getRefs<Request>();
+    refs->insert({RequestKey<Request>(std::move(req)),
+                  std::move(val)});
+  }
+
+  template <typename Request>
+  void erase(Request req) {
+    auto *refs = getRefs<Request>();
+    refs->erase(RequestKey<Request>(std::move(req)));
+  }
+
+  void clear() {
+#define SWIFT_TYPEID_ZONE(Name, Id) Name##ZoneRefs.clear();
+#include "swift/Basic/TypeIDZones.def"
+#undef SWIFT_TYPEID_ZONE
+  }
+};
+
 } // end namespace evaluator
 
 } // end namespace swift