Requestify raw and brief comment computation

Turn these ASTContext cached computations into
request evaluator cached computations.

Author: hamishknight

include/swift/AST/ASTContext.h

@@ -1522,12 +1522,6 @@ class ASTContext final {
   Optional<ExternalSourceLocs *> getExternalSourceLocs(const Decl *D);
   void setExternalSourceLocs(const Decl *D, ExternalSourceLocs *Locs);
 
-  Optional<std::pair<RawComment, bool>> getRawComment(const Decl *D);
-  void setRawComment(const Decl *D, RawComment RC, bool FromSerialized);
-
-  Optional<StringRef> getBriefComment(const Decl *D);
-  void setBriefComment(const Decl *D, StringRef Comment);
-
   friend TypeBase;
   friend ArchetypeType;
   friend OpaqueTypeDecl;

include/swift/AST/Decl.h

@@ -747,6 +747,7 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   friend class IterableDeclContext;
   friend class MemberLookupTable;
   friend class DeclDeserializer;
+  friend class RawCommentRequest;
 
 private:
   llvm::PointerUnion<DeclContext *, ASTContext *> Context;

include/swift/AST/TypeCheckRequests.h

@@ -4188,6 +4188,44 @@ class LocalDiscriminatorsRequest
   bool isCached() const { return true; }
 };
 
+/// Retrieve the raw documentation comment of a declaration.
+class RawCommentRequest
+    : public SimpleRequest<RawCommentRequest,
+                           RawComment(const Decl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  RawComment evaluate(Evaluator &evaluator, const Decl *D) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+};
+
+/// Retrieve the brief portion of a declaration's document comment.
+class BriefCommentRequest
+    : public SimpleRequest<BriefCommentRequest,
+                           StringRef(const Decl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  StringRef evaluate(Evaluator &evaluator, const Decl *D) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+};
+
 /// Checks that all of a class's \c \@objcImplementation extensions provide
 /// complete and correct implementations for their corresponding interfaces.
 /// This is done on all of a class's implementations at once to improve diagnostics.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -469,6 +469,12 @@ SWIFT_REQUEST(TypeChecker, GetRuntimeDiscoverableAttributes,
 SWIFT_REQUEST(TypeChecker, LocalDiscriminatorsRequest,
               unsigned(DeclContext *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, RawCommentRequest,
+              RawComment(const Decl *),
+              Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, BriefCommentRequest,
+              StringRef(const Decl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, IsNonUserModuleRequest,
               bool(ModuleDecl *),
               Cached, NoLocationInfo)

lib/AST/ASTContext.cpp

@@ -314,12 +314,6 @@ struct ASTContext::Implementation {
   /// actual \c SourceLocs that require opening their external buffer.
   llvm::DenseMap<const Decl *, ExternalSourceLocs *> ExternalSourceLocs;
 
-  /// Map from Swift declarations to raw comments.
-  llvm::DenseMap<const Decl *, std::pair<RawComment, bool>> RawComments;
-
-  /// Map from Swift declarations to brief comments.
-  llvm::DenseMap<const Decl *, StringRef> BriefComments;
-
   /// Map from declarations to foreign error conventions.
   /// This applies to both actual imported functions and to @objc functions.
   llvm::DenseMap<const AbstractFunctionDecl *,
@@ -2596,30 +2590,6 @@ void ASTContext::setExternalSourceLocs(const Decl *D,
   getImpl().ExternalSourceLocs[D] = Locs;
 }
 
-Optional<std::pair<RawComment, bool>> ASTContext::getRawComment(const Decl *D) {
-  auto Known = getImpl().RawComments.find(D);
-  if (Known == getImpl().RawComments.end())
-    return None;
-
-  return Known->second;
-}
-
-void ASTContext::setRawComment(const Decl *D, RawComment RC, bool FromSerialized) {
-  getImpl().RawComments[D] = std::make_pair(RC, FromSerialized);
-}
-
-Optional<StringRef> ASTContext::getBriefComment(const Decl *D) {
-  auto Known = getImpl().BriefComments.find(D);
-  if (Known == getImpl().BriefComments.end())
-    return None;
-
-  return Known->second;
-}
-
-void ASTContext::setBriefComment(const Decl *D, StringRef Comment) {
-  getImpl().BriefComments[D] = Comment;
-}
-
 NormalProtocolConformance *
 ASTContext::getConformance(Type conformingType,
                            ProtocolDecl *protocol,
@@ -2922,8 +2892,6 @@ size_t ASTContext::getTotalMemory() const {
     getImpl().Allocator.getTotalMemory() +
     getImpl().Cleanups.capacity() +
     llvm::capacity_in_bytes(getImpl().ModuleLoaders) +
-    llvm::capacity_in_bytes(getImpl().RawComments) +
-    llvm::capacity_in_bytes(getImpl().BriefComments) +
     llvm::capacity_in_bytes(getImpl().ModuleTypes) +
     llvm::capacity_in_bytes(getImpl().GenericParamTypes) +
     // getImpl().GenericFunctionTypes ?

lib/AST/DocComment.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/Comment.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/FileUnit.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/RawComment.h"
@@ -498,39 +499,37 @@ DocComment *swift::getCascadingDocComment(swift::markup::MarkupContext &MC,
   return doc;
 }
 
-StringRef Decl::getBriefComment() const {
-  if (!this->canHaveComment())
-    return StringRef();
-
-  // Check the cache in ASTContext.
-  auto &Context = getASTContext();
-  if (Optional<StringRef> Comment = Context.getBriefComment(this))
-    return Comment.value();
+StringRef
+BriefCommentRequest::evaluate(Evaluator &evaluator, const Decl *D) const {
+  auto *DC = D->getDeclContext();
+  auto &ctx = DC->getASTContext();
 
-  // Check if the serialized module may have the brief comment available.
-  if (auto *Unit =
-          dyn_cast<FileUnit>(this->getDeclContext()->getModuleScopeContext())) {
-    if (Optional<CommentInfo> C = Unit->getCommentForDecl(this)) {
-      Context.setBriefComment(this, C->Brief);
+  // Check if the brief comment is available in the swiftdoc.
+  if (auto *Unit = dyn_cast<FileUnit>(DC->getModuleScopeContext())) {
+    if (auto C = Unit->getCommentForDecl(D))
       return C->Brief;
-    }
   }
 
-  // Otherwise, parse the brief from the raw comment itself.
-  auto RC = getRawComment();
-
-  StringRef Result;
-  if (RC.isEmpty())
-    if (auto *docD = getDocCommentProvidingDecl(this))
+  // Otherwise, parse the brief from the raw comment itself. This will look into the
+  // swiftsourceinfo if needed.
+  auto RC = D->getRawComment();
+  if (RC.isEmpty()) {
+    if (auto *docD = getDocCommentProvidingDecl(D))
       RC = docD->getRawComment();
-  if (!RC.isEmpty()) {
-    SmallString<256> BriefStr;
-    llvm::raw_svector_ostream OS(BriefStr);
-    printBriefComment(RC, OS);
-    Result = Context.AllocateCopy(BriefStr.str());
   }
+  if (RC.isEmpty())
+    return StringRef();
+
+  SmallString<256> BriefStr;
+  llvm::raw_svector_ostream OS(BriefStr);
+  printBriefComment(RC, OS);
+  return ctx.AllocateCopy(BriefStr.str());
+}
+
+StringRef Decl::getBriefComment() const {
+  if (!this->canHaveComment())
+    return StringRef();
 
-  // Cache it.
-  Context.setBriefComment(this, Result);
-  return Result;
+  auto &eval = getASTContext().evaluator;
+  return evaluateOrDefault(eval, BriefCommentRequest{this}, StringRef());
 }

lib/AST/RawComment.cpp

@@ -23,6 +23,7 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/PrimitiveParsing.h"
@@ -130,37 +131,28 @@ static RawComment toRawComment(ASTContext &Context, CharSourceRange Range) {
   return Result;
 }
 
-RawComment Decl::getRawComment() const {
-  if (!this->canHaveComment())
-    return RawComment();
-
-  // Check the cache in ASTContext.
-  auto &Context = getASTContext();
-  if (Optional<std::pair<RawComment, bool>> RC = Context.getRawComment(this)) {
-    return RC.value().first;
-  }
+RawComment RawCommentRequest::evaluate(Evaluator &eval, const Decl *D) const {
+  auto *DC = D->getDeclContext();
+  auto &ctx = DC->getASTContext();
 
   // Check the declaration itself.
-  if (auto *Attr = getAttrs().getAttribute<RawDocCommentAttr>()) {
-    RawComment Result = toRawComment(Context, Attr->getCommentRange());
-    Context.setRawComment(this, Result, true);
-    return Result;
-  }
+  if (auto *Attr = D->getAttrs().getAttribute<RawDocCommentAttr>())
+    return toRawComment(ctx, Attr->getCommentRange());
 
-  if (!getDeclContext())
-    return RawComment();
-  auto *Unit = dyn_cast<FileUnit>(getDeclContext()->getModuleScopeContext());
+  auto *Unit = dyn_cast<FileUnit>(DC->getModuleScopeContext());
   if (!Unit)
     return RawComment();
 
   switch (Unit->getKind()) {
   case FileUnitKind::SerializedAST: {
-    auto *CachedLocs = getSerializedLocs();
+    // First check to see if we have the comment location available in the
+    // swiftsourceinfo, allowing us to grab it from the original file.
+    auto *CachedLocs = D->getSerializedLocs();
     if (!CachedLocs->DocRanges.empty()) {
       SmallVector<SingleRawComment, 4> SRCs;
       for (const auto &Range : CachedLocs->DocRanges) {
         if (Range.isValid()) {
-          SRCs.push_back({Range, Context.SourceMgr});
+          SRCs.push_back({Range, ctx.SourceMgr});
         } else {
           // if we've run into an invalid range, don't bother trying to load
           // any of the other comments
@@ -169,17 +161,13 @@ RawComment Decl::getRawComment() const {
         }
       }
 
-      if (!SRCs.empty()) {
-        auto RC = RawComment(Context.AllocateCopy(llvm::makeArrayRef(SRCs)));
-        Context.setRawComment(this, RC, true);
-        return RC;
-      }
+      if (!SRCs.empty())
+        return RawComment(ctx.AllocateCopy(llvm::makeArrayRef(SRCs)));
     }
 
-    if (Optional<CommentInfo> C = Unit->getCommentForDecl(this)) {
-      Context.setRawComment(this, C->Raw, false);
+    // Otherwise check to see if we have a comment available in the swiftdoc.
+    if (auto C = Unit->getCommentForDecl(D))
       return C->Raw;
-    }
 
     return RawComment();
   }
@@ -193,6 +181,14 @@ RawComment Decl::getRawComment() const {
   llvm_unreachable("invalid file kind");
 }
 
+RawComment Decl::getRawComment() const {
+  if (!this->canHaveComment())
+    return RawComment();
+
+  auto &eval = getASTContext().evaluator;
+  return evaluateOrDefault(eval, RawCommentRequest{this}, RawComment());
+}
+
 Optional<StringRef> Decl::getGroupName() const {
   if (hasClangNode())
     return None;