Adjust BriefCommentRequest to only query swiftdoc if we have it

If we have both loaded a swiftdoc, and the decl we
have should have had its doc comment serialized into
it, we can check it without needing to fall back
to the swiftsourceinfo.

This requires a couple of refactorings:

- Factoring out the `shouldIncludeDecl` logic
into `getDocCommentSerializationTargetFor` for
determining whether a doc comment should end up
in the swiftdoc or not.
- Factoring out `CommentProviderFinder` for searching
for the doc providing comment decl for brief
comments, in order to allow us to avoid querying
the raw comment when searching for it. This has the
added bonus of meaning we no longer need to fall
back to parsing the raw comment for the brief
comment if the comment is provided by another decl
in the swiftdoc.

This diff is best viewed without whitespace.

Author: hamishknight

include/swift/AST/Comment.h

@@ -111,6 +111,25 @@ extractCommentParts(swift::markup::MarkupContext &MC,
 
 /// Extract brief comment from \p RC, and print it to \p OS .
 void printBriefComment(RawComment RC, llvm::raw_ostream &OS);
+
+/// Describes the intended serialization target for a doc comment.
+enum class DocCommentSerializationTarget : uint8_t {
+  /// The doc comment should not be serialized.
+  None = 0,
+
+  /// The doc comment should only be serialized in the 'swiftsourceinfo'.
+  SourceInfoOnly,
+
+  /// The doc comment should be serialized in both the 'swiftdoc' and
+  /// 'swiftsourceinfo'.
+  SwiftDocAndSourceInfo,
+};
+
+/// Retrieve the expected serialization target for a documentation comment
+/// attached to the given decl.
+DocCommentSerializationTarget
+getDocCommentSerializationTargetFor(const Decl *D);
+
 } // namespace swift
 
 #endif // LLVM_SWIFT_AST_COMMENT_H

include/swift/AST/FileUnit.h

@@ -167,6 +167,10 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
     return None;
   }
 
+  /// For a serialized AST file, returns \c true if an adjacent swiftdoc has been
+  /// loaded. Otherwise, returns \c false.
+  virtual bool hasLoadedSwiftDoc() const { return false; }
+
   virtual Optional<StringRef>
   getGroupNameForDecl(const Decl *D) const {
     return None;

include/swift/Serialization/SerializedModuleLoader.h

@@ -430,6 +430,8 @@ class SerializedASTFile final : public LoadedFile {
 
   Optional<CommentInfo> getCommentForDecl(const Decl *D) const override;
 
+  bool hasLoadedSwiftDoc() const override;
+
   Optional<StringRef> getGroupNameForDecl(const Decl *D) const override;
 
 

lib/AST/DocComment.cpp

@@ -392,85 +392,122 @@ DocComment *swift::getSingleDocComment(swift::markup::MarkupContext &MC,
 }
 
 namespace {
-const ValueDecl *findOverriddenDeclWithDocComment(const ValueDecl *VD) {
-  // Only applies to class member.
-  if (!VD->getDeclContext()->getSelfClassDecl())
-    return nullptr;
-
-  while (auto *baseDecl = VD->getOverriddenDecl()) {
-    if (!baseDecl->getRawComment().isEmpty())
-      return baseDecl;
-    VD = baseDecl;
+/// Helper class for finding the comment providing decl for either a brief or
+/// raw comment.
+template <typename Result>
+class CommentProviderFinder final {
+  using ResultWithDecl = std::pair<Result, const Decl *>;
+  using VisitFnTy = Optional<Result>(*)(const Decl *);
+
+  VisitFnTy VisitFn;
+
+public:
+  CommentProviderFinder(VisitFnTy visitFn) : VisitFn(visitFn) {}
+
+private:
+  Optional<ResultWithDecl> visit(const Decl *D) {
+    // Adapt the provided visitor function to also return the decl.
+    if (auto result = VisitFn(D))
+      return {{*result, D}};
+    return None;
   }
 
-  return nullptr;
-}
+  Optional<ResultWithDecl> findOverriddenDecl(const ValueDecl *VD) {
+    // Only applies to class member.
+    if (!VD->getDeclContext()->getSelfClassDecl())
+      return None;
 
-const ValueDecl *findDefaultProvidedDeclWithDocComment(const ValueDecl *VD) {
-  auto protocol = VD->getDeclContext()->getExtendedProtocolDecl();
-  // Only applies to protocol extension member.
-  if (!protocol)
-    return nullptr;
+    while (auto *baseDecl = VD->getOverriddenDecl()) {
+      if (auto result = visit(baseDecl))
+        return result;
 
-  ValueDecl *requirement = nullptr;
+      VD = baseDecl;
+    }
+    return None;
+  }
 
-  SmallVector<ValueDecl *, 2> members;
-  protocol->lookupQualified(const_cast<ProtocolDecl *>(protocol),
-                            DeclNameRef(VD->getName()),
-                            NLOptions::NL_ProtocolMembers,
-                            members);
+  Optional<ResultWithDecl> findDefaultProvidedDecl(const ValueDecl *VD) {
+    // Only applies to protocol extension member.
+    auto *protocol = VD->getDeclContext()->getExtendedProtocolDecl();
+    if (!protocol)
+      return None;
+
+    SmallVector<ValueDecl *, 2> members;
+    protocol->lookupQualified(const_cast<ProtocolDecl *>(protocol),
+                              DeclNameRef(VD->getName()),
+                              NLOptions::NL_ProtocolMembers, members);
+
+    Optional<ResultWithDecl> result;
+    for (auto *member : members) {
+      if (!isa<ProtocolDecl>(member->getDeclContext()) ||
+          !member->isProtocolRequirement())
+        continue;
 
-  for (auto *member : members) {
-    if (!isa<ProtocolDecl>(member->getDeclContext()) ||
-        !member->isProtocolRequirement() || member->getRawComment().isEmpty())
-      continue;
-    if (requirement)
-      // Found two or more decls with doc-comment.
-      return nullptr;
+      auto newResult = visit(member);
+      if (!newResult)
+        continue;
 
-    requirement = member;
+      if (result) {
+        // Found two or more decls with doc-comment.
+        return None;
+      }
+      result = newResult;
+    }
+    return result;
   }
-  return requirement;
-}
 
-const ValueDecl *findRequirementDeclWithDocComment(const ValueDecl *VD) {
-  std::queue<const ValueDecl *> requirements;
-  while (true) {
-    for (auto *req : VD->getSatisfiedProtocolRequirements()) {
-      if (!req->getRawComment().isEmpty())
-        return req;
-      else
+  Optional<ResultWithDecl> findRequirementDecl(const ValueDecl *VD) {
+    std::queue<const ValueDecl *> requirements;
+    while (true) {
+      for (auto *req : VD->getSatisfiedProtocolRequirements()) {
+        if (auto result = visit(req))
+          return result;
+
         requirements.push(req);
+      }
+      if (requirements.empty())
+        return None;
+
+      VD = requirements.front();
+      requirements.pop();
     }
-    if (requirements.empty())
-      return nullptr;
-    VD = requirements.front();
-    requirements.pop();
   }
-}
-} // namespace
 
-const Decl *swift::getDocCommentProvidingDecl(const Decl *D) {
-  if (!D->canHaveComment())
-    return nullptr;
+public:
+  Optional<ResultWithDecl> findCommentProvider(const Decl *D) {
+    if (auto result = visit(D))
+      return result;
 
-  if (!D->getRawComment().isEmpty())
-    return D;
+    auto *VD = dyn_cast<ValueDecl>(D);
+    if (!VD)
+      return None;
 
-  auto *VD = dyn_cast<ValueDecl>(D);
-  if (!VD)
-    return nullptr;
+    if (auto result = findOverriddenDecl(VD))
+      return result;
 
-  if (auto *overridden = findOverriddenDeclWithDocComment(VD))
-    return overridden;
+    if (auto result = findDefaultProvidedDecl(VD))
+      return result;
 
-  if (auto *requirement = findDefaultProvidedDeclWithDocComment(VD))
-    return requirement;
+    if (auto result = findRequirementDecl(VD))
+      return result;
 
-  if (auto *requirement = findRequirementDeclWithDocComment(VD))
-    return requirement;
+    return None;
+  }
+};
+} // end anonymous namespace
 
-  return nullptr;
+const Decl *swift::getDocCommentProvidingDecl(const Decl *D) {
+  // Search for the first decl we see with a non-empty raw comment.
+  auto finder = CommentProviderFinder<RawComment>(
+      [](const Decl *D) -> Optional<RawComment> {
+        auto comment = D->getRawComment();
+        if (comment.isEmpty())
+          return None;
+        return comment;
+      });
+
+  auto result = finder.findCommentProvider(D);
+  return result ? result->second : nullptr;
 }
 
 DocComment *swift::getCascadingDocComment(swift::markup::MarkupContext &MC,
@@ -499,31 +536,50 @@ DocComment *swift::getCascadingDocComment(swift::markup::MarkupContext &MC,
   return doc;
 }
 
-StringRef
-BriefCommentRequest::evaluate(Evaluator &evaluator, const Decl *D) const {
-  auto *DC = D->getDeclContext();
-  auto &ctx = DC->getASTContext();
+/// Retrieve the brief comment for a given decl \p D, without attempting to
+/// walk any requirements or overrides.
+static Optional<StringRef> getDirectBriefComment(const Decl *D) {
+  if (!D->canHaveComment())
+    return None;
 
-  // 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;
+  auto *ModuleDC = D->getDeclContext()->getModuleScopeContext();
+  auto &Ctx = ModuleDC->getASTContext();
+
+  // If we expect the comment to be in the swiftdoc, check for it if we loaded a
+  // swiftdoc. If missing from the swiftdoc, we know it will not be in the
+  // swiftsourceinfo either, so we can bail early.
+  if (auto *Unit = dyn_cast<FileUnit>(ModuleDC)) {
+    if (Unit->hasLoadedSwiftDoc()) {
+      auto target = getDocCommentSerializationTargetFor(D);
+      if (target == DocCommentSerializationTarget::SwiftDocAndSourceInfo) {
+        auto C = Unit->getCommentForDecl(D);
+        if (!C)
+          return None;
+
+        return C->Brief;
+      }
+    }
   }
 
-  // Otherwise, parse the brief from the raw comment itself. This will look into the
-  // swiftsourceinfo if needed.
+  // 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())
-    return StringRef();
+    return None;
 
   SmallString<256> BriefStr;
   llvm::raw_svector_ostream OS(BriefStr);
   printBriefComment(RC, OS);
-  return ctx.AllocateCopy(BriefStr.str());
+  return Ctx.AllocateCopy(BriefStr.str());
+}
+
+StringRef BriefCommentRequest::evaluate(Evaluator &evaluator,
+                                        const Decl *D) const {
+  // Perform a walk over the potential providers of the brief comment,
+  // retrieving the first one we come across.
+  CommentProviderFinder finder(getDirectBriefComment);
+  auto result = finder.findCommentProvider(D);
+  return result ? result->first : StringRef();
 }
 
 StringRef Decl::getBriefComment() const {

lib/AST/RawComment.cpp

@@ -236,3 +236,78 @@ CharSourceRange RawComment::getCharSourceRange() {
                 static_cast<const char *>(Start.getOpaquePointerValue());
   return CharSourceRange(Start, Length);
 }
+
+static bool hasDoubleUnderscore(const Decl *D) {
+  // Exclude decls with double-underscored names, either in arguments or
+  // base names.
+  static StringRef Prefix = "__";
+
+  // If it's a function or subscript with a parameter with leading
+  // double underscore, it's a private function or subscript.
+  if (isa<AbstractFunctionDecl>(D) || isa<SubscriptDecl>(D)) {
+    auto *params = getParameterList(cast<ValueDecl>(const_cast<Decl *>(D)));
+    if (params->hasInternalParameter(Prefix))
+      return true;
+  }
+
+  if (auto *VD = dyn_cast<ValueDecl>(D)) {
+    auto Name = VD->getBaseName();
+    if (!Name.isSpecial() && Name.getIdentifier().str().startswith(Prefix)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+static DocCommentSerializationTarget
+getDocCommentSerializationTargetImpl(const Decl *D) {
+  if (auto *ED = dyn_cast<ExtensionDecl>(D)) {
+    auto *extended = ED->getExtendedNominal();
+    if (!extended)
+      return DocCommentSerializationTarget::None;
+
+    return getDocCommentSerializationTargetFor(extended);
+  }
+  auto *VD = dyn_cast<ValueDecl>(D);
+  if (!VD)
+    return DocCommentSerializationTarget::None;
+
+  // The use of getEffectiveAccess is unusual here; we want to take the
+  // testability state into account and emit documentation if and only if they
+  // are visible to clients (which means public ordinarily, but public+internal
+  // when testing enabled).
+  switch (VD->getEffectiveAccess()) {
+  case AccessLevel::Private:
+  case AccessLevel::FilePrivate:
+  case AccessLevel::Internal:
+    // There's no point serializing anything internal or below, as they are not
+    // accessible outside their defining module.
+    return DocCommentSerializationTarget::None;
+  case AccessLevel::Package:
+    // Package doc comments can be referenced outside their module, but only
+    // locally, so can't be included in swiftdoc.
+    return DocCommentSerializationTarget::SourceInfoOnly;
+  case AccessLevel::Public:
+  case AccessLevel::Open:
+    return DocCommentSerializationTarget::SwiftDocAndSourceInfo;
+  }
+  llvm_unreachable("Unhandled case in switch!");
+}
+
+DocCommentSerializationTarget
+swift::getDocCommentSerializationTargetFor(const Decl *D) {
+  auto Limit = DocCommentSerializationTarget::SwiftDocAndSourceInfo;
+
+  // We can't include SPI decls in swiftdoc.
+  if (D->isSPI())
+    Limit = DocCommentSerializationTarget::SourceInfoOnly;
+
+  // .swiftdoc doesn't include comments for double underscored symbols, but
+  // for .swiftsourceinfo, having the source location for these symbols isn't
+  // a concern because these symbols are in .swiftinterface anyway.
+  if (hasDoubleUnderscore(D))
+    Limit = DocCommentSerializationTarget::SourceInfoOnly;
+
+  auto Result = getDocCommentSerializationTargetImpl(D);
+  return std::min(Result, Limit);
+}

lib/Serialization/ModuleFile.cpp

@@ -1065,6 +1065,10 @@ Optional<CommentInfo> ModuleFile::getCommentForDecl(const Decl *D) const {
   return getCommentForDeclByUSR(USRBuffer.str());
 }
 
+bool ModuleFile::hasLoadedSwiftDoc() const {
+  return Core->DeclCommentTable != nullptr;
+}
+
 void ModuleFile::collectSerializedSearchPath(
     llvm::function_ref<void(StringRef)> callback) const {
   for (auto path: Core->SearchPaths) {

lib/Serialization/ModuleFile.h

@@ -883,6 +883,7 @@ class ModuleFile
   Optional<unsigned> getSourceOrderForDecl(const Decl *D) const;
   void collectAllGroups(SmallVectorImpl<StringRef> &Names) const;
   Optional<CommentInfo> getCommentForDecl(const Decl *D) const;
+  bool hasLoadedSwiftDoc() const;
   Optional<CommentInfo> getCommentForDeclByUSR(StringRef USR) const;
   Optional<StringRef> getGroupNameByUSR(StringRef USR) const;
   Optional<ExternalSourceLocs::RawLocs>