Naive Dependency Replay

Finish off private intransitive dependencies with an implementation of
dependency replay.

For the sake of illustration, imagine a chain of requests

A -> B -> C -> ...

Supposing each request is never cached, then every invocation of the
compiler with the same inputs will always kick off the exact same set of
requests. For the purposes of dependency tracking, that also means every
single lookup request will run without issue, and all dependencies will
be accurately reported. But we live in a world with cached requests.
Suppose request B* is cached. The first time we encounter that request,
its evaluation order looks identical:

A -> B* -> C -> ...

If we are in a mode that compiles single primaries, this is not
a problem because every request graph will look like this.
But if we are in a mode where we are compiling multiple primaries, then
subsequent request graphs will *actually* hit the cache and never
execute request C or any of its dependent computations!

A -> B*

Supposing C was a lookup request, that means the name(s) looked up
downstream of B* will *never* be recorded in the referenced name tracker
which can lead to miscompilation. Note that this is not a problem
inherent to the design of the request evaluator - caches in the compiler
have *always* hidden dependent lookups. In fact, the request evaluator
provides us our first opportunity to resolve this correctness bug!

Author: CodaFi

include/swift/AST/AnyRequest.h

@@ -66,6 +66,9 @@ struct AnyRequestVTable {
     static SourceLoc getNearestLoc(const void *ptr) {
       return static_cast<const Request *>(ptr)->getNearestLoc();
     }
+    static bool isCached(const void *ptr) {
+      return static_cast<const Request *>(ptr)->isCached();
+    }
   };
 
   const uint64_t typeID;
@@ -78,8 +81,29 @@ struct AnyRequestVTable {
   const std::function<void(const void *, DiagnosticEngine &)> diagnoseCycle;
   const std::function<void(const void *, DiagnosticEngine &)> noteCycleStep;
   const std::function<SourceLoc(const void *)> getNearestLoc;
+  const std::function<bool(const void *)> isCached;
 
-  template <typename Request>
+  template <typename Request,
+            typename std::enable_if<Request::isEverCached>::type * = nullptr>
+  static const AnyRequestVTable *get() {
+    static const AnyRequestVTable vtable = {
+        TypeID<Request>::value,
+        sizeof(Request),
+        &Impl<Request>::copy,
+        &Impl<Request>::getHash,
+        &Impl<Request>::deleter,
+        &Impl<Request>::isEqual,
+        &Impl<Request>::simpleDisplay,
+        &Impl<Request>::diagnoseCycle,
+        &Impl<Request>::noteCycleStep,
+        &Impl<Request>::getNearestLoc,
+        &Impl<Request>::isCached,
+    };
+    return &vtable;
+  }
+
+  template <typename Request,
+            typename std::enable_if<!Request::isEverCached>::type * = nullptr>
   static const AnyRequestVTable *get() {
     static const AnyRequestVTable vtable = {
         TypeID<Request>::value,
@@ -92,6 +116,7 @@ struct AnyRequestVTable {
         &Impl<Request>::diagnoseCycle,
         &Impl<Request>::noteCycleStep,
         &Impl<Request>::getNearestLoc,
+        [](auto){ return false; },
     };
     return &vtable;
   }
@@ -194,6 +219,10 @@ class AnyRequestBase {
     return getVTable()->getNearestLoc(getRawStorage());
   }
 
+  bool isCached() const {
+    return getVTable()->isCached(getRawStorage());
+  }
+
   /// Compare two instances for equality.
   friend bool operator==(const AnyRequestBase<Derived> &lhs,
                          const AnyRequestBase<Derived> &rhs) {

include/swift/AST/Evaluator.h

@@ -431,19 +431,22 @@ class Evaluator {
   }
 
 private:
-  // Report the result of evaluating a request that is not a dependency sink -
-  // which is to say do nothing.
-  template <typename Request,
-            typename std::enable_if<!Request::isDependencySink>::type * = nullptr>
+  // Report the result of evaluating a request that is not a dependency sink.
+  template <typename Request, typename std::enable_if<
+                                  !Request::isDependencySink>::type * = nullptr>
   void reportEvaluatedResult(const Request &r,
-                             const typename Request::OutputType &o) {}
+                             const typename Request::OutputType &o) {
+    collector.replay(ActiveRequest(r));
+  }
 
   // Report the result of evaluating a request that is a dependency sink.
   template <typename Request,
             typename std::enable_if<Request::isDependencySink>::type * = nullptr>
   void reportEvaluatedResult(const Request &r,
                              const typename Request::OutputType &o) {
-    r.writeDependencySink(collector, o);
+    return collector.record(activeRequests, [&r, &o](auto &c) {
+      return r.writeDependencySink(c, o);
+    });
   }
 
 public:

include/swift/AST/EvaluatorDependencies.h

@@ -18,6 +18,7 @@
 #ifndef SWIFT_AST_EVALUATOR_DEPENDENCIES_H
 #define SWIFT_AST_EVALUATOR_DEPENDENCIES_H
 
+#include "swift/AST/AnyRequest.h"
 #include "swift/AST/AttrKind.h"
 #include "swift/AST/SourceFile.h"
 #include "llvm/ADT/PointerIntPair.h"
@@ -113,6 +114,70 @@ struct DependencyCollector {
   /// A stack of dependency sources in the order they were evaluated.
   llvm::SmallVector<evaluator::DependencySource, 8> dependencySources;
 
+  struct Reference {
+  public:
+    enum class Kind {
+      Empty,
+      Tombstone,
+      UsedMember,
+      PotentialMember,
+      TopLevel,
+      Dynamic,
+    } kind;
+
+    NominalTypeDecl *subject;
+    DeclBaseName name;
+
+  private:
+    Reference(Kind kind, NominalTypeDecl *subject, DeclBaseName name)
+        : kind(kind), subject(subject), name(name) {}
+
+  public:
+    static Reference empty() {
+      return {Kind::Empty, llvm::DenseMapInfo<NominalTypeDecl *>::getEmptyKey(),
+              llvm::DenseMapInfo<DeclBaseName>::getEmptyKey()};
+    }
+
+    static Reference tombstone() {
+      return {Kind::Empty,
+              llvm::DenseMapInfo<NominalTypeDecl *>::getTombstoneKey(),
+              llvm::DenseMapInfo<DeclBaseName>::getTombstoneKey()};
+    }
+
+  public:
+    static Reference usedMember(NominalTypeDecl *subject, DeclBaseName name) {
+      return {Kind::UsedMember, subject, name};
+    }
+
+    static Reference potentialMember(NominalTypeDecl *subject) {
+      return {Kind::PotentialMember, subject, DeclBaseName()};
+    }
+
+    static Reference topLevel(DeclBaseName name) {
+      return {Kind::TopLevel, nullptr, name};
+    }
+
+    static Reference dynamic(DeclBaseName name) {
+      return {Kind::Dynamic, nullptr, name};
+    }
+
+  public:
+    struct Info {
+      static inline Reference getEmptyKey() { return Reference::empty(); }
+      static inline Reference getTombstoneKey() {
+        return Reference::tombstone();
+      }
+      static inline unsigned getHashValue(const Reference &Val) {
+        return llvm::hash_combine(Val.kind, Val.subject,
+                                  Val.name.getAsOpaquePointer());
+      }
+      static bool isEqual(const Reference &LHS, const Reference &RHS) {
+        return LHS.kind == RHS.kind && LHS.subject == RHS.subject &&
+               LHS.name == RHS.name;
+      }
+    };
+  };
+
 public:
   enum class Mode {
     // Enables the current "status quo" behavior of the dependency collector.
@@ -128,8 +193,21 @@ struct DependencyCollector {
     ExperimentalPrivateDependencies,
   };
   Mode mode;
+  llvm::DenseMap<AnyRequest, llvm::DenseSet<Reference, Reference::Info>>
+      requestReferences;
+  llvm::DenseSet<Reference, Reference::Info> scratch;
+  bool isRecording;
+
+  explicit DependencyCollector(Mode mode)
+      : mode{mode}, requestReferences{}, scratch{}, isRecording{false} {};
 
-  explicit DependencyCollector(Mode mode) : mode{mode} {};
+private:
+  void realizeOrRecord(const Reference &ref);
+
+public:
+  void replay(const swift::ActiveRequest &req);
+  void record(const llvm::SetVector<swift::ActiveRequest> &stack,
+              llvm::function_ref<void(DependencyCollector &)> rec);
 
 public:
   /// Registers a named reference from the current dependency scope to a member

lib/AST/Evaluator.cpp

@@ -21,6 +21,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/SaveAndRestore.h"
 
 using namespace swift;
 
@@ -378,24 +379,93 @@ void Evaluator::dumpDependenciesGraphviz() const {
   printDependenciesGraphviz(llvm::dbgs());
 }
 
+void evaluator::DependencyCollector::realizeOrRecord(const Reference &ref) {
+  if (isRecording) {
+    scratch.insert(std::move(ref));
+    return;
+  }
+
+  auto *tracker = getActiveDependencyTracker();
+  if (!tracker) {
+    return;
+  }
+
+  switch (ref.kind) {
+  case Reference::Kind::Empty:
+  case Reference::Kind::Tombstone:
+    llvm_unreachable("cannot record empty dependency");
+  case Reference::Kind::UsedMember:
+    tracker->addUsedMember({ref.subject, ref.name}, isActiveSourceCascading());
+    break;
+  case Reference::Kind::PotentialMember:
+    tracker->addUsedMember({ref.subject, Identifier()},
+                           isActiveSourceCascading());
+    break;
+  case Reference::Kind::TopLevel:
+    tracker->addTopLevelName(ref.name, isActiveSourceCascading());
+    break;
+  case Reference::Kind::Dynamic:
+    tracker->addDynamicLookupName(ref.name, isActiveSourceCascading());
+    break;
+  }
+}
+
 void evaluator::DependencyCollector::addUsedMember(NominalTypeDecl *subject,
                                                    DeclBaseName name) {
-  if (auto *tracker = getActiveDependencyTracker())
-    tracker->addUsedMember({subject, name}, isActiveSourceCascading());
+  return realizeOrRecord(Reference::usedMember(subject, name));
 }
 
 void evaluator::DependencyCollector::addPotentialMember(
     NominalTypeDecl *subject) {
-  if (auto *tracker = getActiveDependencyTracker())
-    tracker->addUsedMember({subject, Identifier()}, isActiveSourceCascading());
+  return realizeOrRecord(Reference::potentialMember(subject));
 }
 
 void evaluator::DependencyCollector::addTopLevelName(DeclBaseName name) {
-  if (auto *tracker = getActiveDependencyTracker())
-    tracker->addTopLevelName(name, isActiveSourceCascading());
+  return realizeOrRecord(Reference::topLevel(name));
 }
 
 void evaluator::DependencyCollector::addDynamicLookupName(DeclBaseName name) {
-  if (auto *tracker = getActiveDependencyTracker())
-    tracker->addDynamicLookupName(name, isActiveSourceCascading());
+  return realizeOrRecord(Reference::dynamic(name));
+}
+
+void evaluator::DependencyCollector::record(
+    const llvm::SetVector<swift::ActiveRequest> &stack,
+    llvm::function_ref<void(DependencyCollector &)> rec) {
+  assert(!isRecording && "Probably not a good idea to allow nested recording");
+  if (mode == Mode::StatusQuo) {
+    return rec(*this);
+  }
+
+  llvm::SaveAndRestore<bool> restore(isRecording, true);
+  scratch = {};
+  rec(*this);
+  for (const auto &request : stack) {
+    if (!request.isCached()) {
+      continue;
+    }
+
+    auto entry = requestReferences.find_as(request);
+    if (entry == requestReferences.end()) {
+      requestReferences.insert({AnyRequest(request), scratch});
+      continue;
+    }
+
+    entry->second.insert(scratch.begin(), scratch.end());
+  }
+}
+
+void evaluator::DependencyCollector::replay(const swift::ActiveRequest &req) {
+  if (mode == Mode::StatusQuo) {
+    return;
+  }
+
+  assert(!isRecording && "Probably not a good idea to allow nested recording");
+  auto entry = requestReferences.find_as(req);
+  if (entry == requestReferences.end()) {
+    return;
+  }
+
+  for (const auto &ref : entry->second) {
+    realizeOrRecord(ref);
+  }
 }