make ASTContext RawComment cache aware of serialization

rdar://76162972

Author: QuietMisdreavus

include/swift/AST/ASTContext.h

@@ -1170,8 +1170,8 @@ class ASTContext final {
 
 private:
   friend Decl;
-  Optional<RawComment> getRawComment(const Decl *D);
-  void setRawComment(const Decl *D, RawComment RC);
+  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);

lib/AST/ASTContext.cpp

@@ -285,7 +285,7 @@ struct ASTContext::Implementation {
   ClangModuleLoader *TheDWARFModuleLoader = nullptr;
 
   /// Map from Swift declarations to raw comments.
-  llvm::DenseMap<const Decl *, RawComment> RawComments;
+  llvm::DenseMap<const Decl *, std::pair<RawComment, bool>> RawComments;
 
   /// Map from Swift declarations to brief comments.
   llvm::DenseMap<const Decl *, StringRef> BriefComments;
@@ -2020,16 +2020,16 @@ ModuleDecl *ASTContext::getStdlibModule(bool loadIfAbsent) {
   return TheStdlibModule;
 }
 
-Optional<RawComment> ASTContext::getRawComment(const Decl *D) {
+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) {
-  getImpl().RawComments[D] = RC;
+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) {

lib/AST/RawComment.cpp

@@ -137,13 +137,16 @@ RawComment Decl::getRawComment(bool SerializedOK) const {
 
   // Check the cache in ASTContext.
   auto &Context = getASTContext();
-  if (Optional<RawComment> RC = Context.getRawComment(this))
-    return RC.getValue();
+  if (Optional<std::pair<RawComment, bool>> RC = Context.getRawComment(this)) {
+    auto P = RC.getValue();
+    if (!SerializedOK || P.second)
+      return P.first;
+  }
 
   // Check the declaration itself.
   if (auto *Attr = getAttrs().getAttribute<RawDocCommentAttr>()) {
     RawComment Result = toRawComment(Context, Attr->getCommentRange());
-    Context.setRawComment(this, Result);
+    Context.setRawComment(this, Result, true);
     return Result;
   }
 
@@ -172,15 +175,15 @@ RawComment Decl::getRawComment(bool SerializedOK) const {
           auto RC = RawComment(Context.AllocateCopy(llvm::makeArrayRef(SRCs)));
 
           if (!RC.isEmpty()) {
-            Context.setRawComment(this, RC);
+            Context.setRawComment(this, RC, true);
             return RC;
           }
         }
       }
     }
 
     if (Optional<CommentInfo> C = Unit->getCommentForDecl(this)) {
-      Context.setRawComment(this, C->Raw);
+      Context.setRawComment(this, C->Raw, false);
       return C->Raw;
     }
 

test/SymbolGraph/Symbols/DocComment/BatchMode.swift

@@ -0,0 +1,8 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -module-name BatchMode -emit-module-path %t/BatchMode.swiftmodule -emit-symbol-graph -emit-symbol-graph-dir %t -O -enable-batch-mode -incremental
+// RUN: %FileCheck %s --input-file %t/BatchMode.symbols.json
+
+/// This is some func.
+public func someFunc() {}
+
+// CHECK: range