Merge pull request swiftlang#35206 from slavapestov/per-request-caches

Per-request caches and dependency maps

Author: slavapestov

include/swift/AST/Evaluator.h

@@ -20,6 +20,7 @@
 
 #include "swift/AST/AnyRequest.h"
 #include "swift/AST/EvaluatorDependencies.h"
+#include "swift/AST/RequestCache.h"
 #include "swift/Basic/AnyValue.h"
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/LangOptions.h"
@@ -211,7 +212,7 @@ class Evaluator {
   llvm::SetVector<ActiveRequest> activeRequests;
 
   /// A cache that stores the results of requests.
-  llvm::DenseMap<AnyRequest, AnyValue> cache;
+  evaluator::RequestCache cache;
 
   /// Track the dependencies of each request.
   ///
@@ -314,14 +315,15 @@ class Evaluator {
            typename std::enable_if<!Request::hasExternalCache>::type* = nullptr>
   void cacheOutput(const Request &request,
                    typename Request::OutputType &&output) {
-    cache.insert({AnyRequest(request), std::move(output)});
+    cache.insert<Request>(request, std::move(output));
   }
 
   /// Do not introduce new callers of this function.
   template<typename Request,
            typename std::enable_if<!Request::hasExternalCache>::type* = nullptr>
   void clearCachedOutput(const Request &request) {
-    cache.erase(AnyRequest(request));
+    cache.erase<Request>(request);
+    recorder.clearRequest<Request>(request);
   }
 
   /// Clear the cache stored within this evaluator.
@@ -372,21 +374,21 @@ class Evaluator {
     FrontendStatsTracer statsTracer = make_tracer(stats, request);
     if (stats) reportEvaluatedRequest(*stats, request);
 
-    recorder.beginRequest(activeReq);
+    recorder.beginRequest<Request>();
 
-    auto &&r = getRequestFunction<Request>()(request, *this);
+    auto &&result = getRequestFunction<Request>()(request, *this);
 
-    recorder.endRequest(activeReq);
+    recorder.endRequest<Request>(request);
 
     handleDependencySourceRequest<Request>(request);
-    handleDependencySinkRequest<Request>(request, r);
+    handleDependencySinkRequest<Request>(request, result);
 
     // Make sure we remove this from the set of active requests once we're
     // done.
     assert(activeRequests.back() == activeReq);
     activeRequests.pop_back();
 
-    return std::move(r);
+    return std::move(result);
   }
 
   /// Get the result of a request, consulting an external cache
@@ -398,7 +400,7 @@ class Evaluator {
   getResultCached(const Request &request) {
     // If there is a cached result, return it.
     if (auto cached = request.getCachedResult()) {
-      recorder.replayCachedRequest(ActiveRequest(request));
+      recorder.replayCachedRequest(request);
       handleDependencySinkRequest<Request>(request, *cached);
       return *cached;
     }
@@ -424,12 +426,12 @@ class Evaluator {
   llvm::Expected<typename Request::OutputType>
   getResultCached(const Request &request) {
     // If we already have an entry for this request in the cache, return it.
-    auto known = cache.find_as(request);
-    if (known != cache.end()) {
-      auto r = known->second.template castTo<typename Request::OutputType>();
-      recorder.replayCachedRequest(ActiveRequest(request));
-      handleDependencySinkRequest<Request>(request, r);
-      return r;
+    auto known = cache.find_as<Request>(request);
+    if (known != cache.end<Request>()) {
+      auto result = known->second;
+      recorder.replayCachedRequest(request);
+      handleDependencySinkRequest<Request>(request, result);
+      return result;
     }
 
     // Compute the result.
@@ -438,7 +440,7 @@ class Evaluator {
       return result;
 
     // Cache the result.
-    cache.insert({AnyRequest(request), *result});
+    cache.insert<Request>(request, *result);
     return result;
   }
 
@@ -465,7 +467,7 @@ class Evaluator {
   void handleDependencySourceRequest(const Request &r) {
     auto source = r.readDependencySource(recorder);
     if (!source.isNull() && source.get()->isPrimary()) {
-      recorder.handleDependencySourceRequest(ActiveRequest(r), source.get());
+      recorder.handleDependencySourceRequest(r, source.get());
     }
   }
 

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,15 +64,14 @@ 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.
   /// Upon completion of a request, we union the completed request's references
   /// with the next innermost active request.
-  std::vector<llvm::DenseSet<DependencyCollector::Reference,
-                             DependencyCollector::Reference::Info>>
+  std::vector<llvm::SmallDenseSet<DependencyCollector::Reference, 2,
+                                  DependencyCollector::Reference::Info>>
       activeRequestReferences;
 
 #ifndef NDEBUG
@@ -83,22 +83,30 @@ struct DependencyRecorder {
 
 public:
   /// Push a new empty set onto the activeRequestReferences stack.
-  void beginRequest(const swift::ActiveRequest &req);
+  template<typename Request>
+  void beginRequest();
 
   /// Pop the activeRequestReferences stack, and insert recorded references
   /// into the requestReferences map, as well as the next innermost entry in
   /// activeRequestReferences.
-  void endRequest(const swift::ActiveRequest &req);
+  template<typename Request>
+  void endRequest(const Request &req);
 
   /// When replaying a request whose value has already been cached, we need
   /// to update the innermost set in the activeRequestReferences stack.
-  void replayCachedRequest(const swift::ActiveRequest &req);
+  template<typename Request>
+  void replayCachedRequest(const Request &req);
 
   /// Upon completion of a dependency source request, we update the
   /// fileReferences map.
-  void handleDependencySourceRequest(const swift::ActiveRequest &req,
+  template<typename Request>
+  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.
@@ -120,6 +128,76 @@ struct DependencyRecorder {
   void enumerateReferencesInFile(const SourceFile *SF,
                                  ReferenceEnumerator f) const ;
 };
+
+template<typename Request>
+void evaluator::DependencyRecorder::beginRequest() {
+  if (!Request::isEverCached && !Request::isDependencySource)
+    return;
+
+  activeRequestReferences.push_back({});
+}
+
+template<typename Request>
+void evaluator::DependencyRecorder::endRequest(const Request &req) {
+  if (!Request::isEverCached && !Request::isDependencySource)
+    return;
+
+  // Grab all the dependencies we've recorded so far, and pop
+  // the stack.
+  auto recorded = std::move(activeRequestReferences.back());
+  activeRequestReferences.pop_back();
+
+  // If we didn't record anything, there is nothing to do.
+  if (recorded.empty())
+    return;
+
+  // Convert the set of dependencies into a vector.
+  std::vector<DependencyCollector::Reference>
+      vec(recorded.begin(), recorded.end());
+
+  // The recorded dependencies bubble up to the parent request.
+  if (!activeRequestReferences.empty()) {
+    activeRequestReferences.back().insert(vec.begin(),
+                                          vec.end());
+  }
+
+  // Finally, record the dependencies so we can replay them
+  // later when the request is re-evaluated.
+  requestReferences.insert<Request>(std::move(req), std::move(vec));
+}
+
+template<typename Request>
+void evaluator::DependencyRecorder::replayCachedRequest(const Request &req) {
+  assert(req.isCached());
+
+  if (activeRequestReferences.empty())
+    return;
+
+  auto found = requestReferences.find_as<Request>(req);
+  if (found == requestReferences.end<Request>())
+    return;
+
+  activeRequestReferences.back().insert(found->second.begin(),
+                                        found->second.end());
+}
+
+template<typename Request>
+void evaluator::DependencyRecorder::handleDependencySourceRequest(
+    const Request &req,
+    SourceFile *source) {
+  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