Merge pull request #69757 from tshortli/refactor-lazy-type-refinement-contexts

Sema: Expand TypeRefinementContexts lazily for unparsed function bodies

Author: tshortli

include/swift/AST/TypeCheckRequests.h

@@ -4541,23 +4541,26 @@ class InitAccessorReferencedVariablesRequest
   bool isCached() const { return true; }
 };
 
-/// Expand the children of the type refinement context for the given
-/// declaration.
+/// Expand the children of the given type refinement context.
 class ExpandChildTypeRefinementContextsRequest
     : public SimpleRequest<ExpandChildTypeRefinementContextsRequest,
-                           bool(Decl *, TypeRefinementContext *),
-                           RequestFlags::Cached> {
+                           std::vector<TypeRefinementContext *>(
+                               TypeRefinementContext *),
+                           RequestFlags::SeparatelyCached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  bool evaluate(Evaluator &evaluator, Decl *decl,
-                TypeRefinementContext *parentTRC) const;
+  std::vector<TypeRefinementContext *>
+  evaluate(Evaluator &evaluator, TypeRefinementContext *parentTRC) const;
 
 public:
+  // Separate caching.
   bool isCached() const { return true; }
+  llvm::Optional<std::vector<TypeRefinementContext *>> getCachedResult() const;
+  void cacheResult(std::vector<TypeRefinementContext *> children) const;
 };
 
 class SerializeAttrGenericSignatureRequest

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -513,8 +513,8 @@ SWIFT_REQUEST(TypeChecker, InitAccessorReferencedVariablesRequest,
                                   ArrayRef<Identifier>),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandChildTypeRefinementContextsRequest,
-              bool(Decl *, TypeRefinementContext *),
-              Cached, NoLocationInfo)
+              std::vector<TypeRefinementContext *>(TypeRefinementContext *),
+              SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, SerializeAttrGenericSignatureRequest,
               GenericSignature(Decl *, SpecializeAttr *),
               SeparatelyCached, NoLocationInfo)

include/swift/AST/TypeRefinementContext.h

@@ -97,6 +97,7 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
   };
 
 private:
+  friend class ExpandChildTypeRefinementContextsRequest;
 
   /// Represents the AST node that introduced a refinement context.
   class IntroNode {
@@ -179,6 +180,11 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
 
   std::vector<TypeRefinementContext *> Children;
 
+  struct {
+    /// Whether this node has child nodes that have not yet been expanded.
+    unsigned needsExpansion : 1;
+  } LazyInfo = {};
+
   TypeRefinementContext(ASTContext &Ctx, IntroNode Node,
                         TypeRefinementContext *Parent, SourceRange SrcRange,
                         const AvailabilityContext &Info,
@@ -240,6 +246,13 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
                          TypeRefinementContext *Parent,
                          const AvailabilityContext &Info);
 
+  Decl *getDeclOrNull() const {
+    auto IntroReason = getReason();
+    if (IntroReason == Reason::Decl || IntroReason == Reason::DeclImplicit)
+      return getIntroductionNode().getAsDecl();
+    return nullptr;
+  }
+
   /// Returns the reason this context was introduced.
   Reason getReason() const;
 
@@ -288,7 +301,13 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
   /// Returns the inner-most TypeRefinementContext descendant of this context
   /// for the given source location.
   TypeRefinementContext *findMostRefinedSubContext(SourceLoc Loc,
-                                                   SourceManager &SM);
+                                                   ASTContext &Ctx);
+
+  bool getNeedsExpansion() const { return LazyInfo.needsExpansion; }
+
+  void setNeedsExpansion(bool needsExpansion) {
+    LazyInfo.needsExpansion = needsExpansion;
+  }
 
   SWIFT_DEBUG_DUMPER(dump(SourceManager &SrcMgr));
   void dump(raw_ostream &OS, SourceManager &SrcMgr) const;
@@ -300,6 +319,10 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
 void simple_display(llvm::raw_ostream &out,
                     const TypeRefinementContext *trc);
 
+inline SourceLoc extractNearestSourceLoc(const TypeRefinementContext *TRC) {
+  return TRC->getIntroductionLoc();
+}
+
 } // end namespace swift
 
 #endif

lib/AST/TypeRefinementContext.cpp

@@ -57,7 +57,7 @@ TypeRefinementContext::createRoot(SourceFile *SF,
       auto charRange = Ctx.SourceMgr.getRangeForBuffer(*SF->getBufferID());
       range = SourceRange(charRange.getStart(), charRange.getEnd());
       parentContext = parentTRC->findMostRefinedSubContext(
-          parentExpansion.getStartLoc(), Ctx.SourceMgr);
+          parentExpansion.getStartLoc(), Ctx);
       availabilityContext = parentContext->getAvailabilityInfo();
     }
   }
@@ -186,29 +186,20 @@ static bool rangeContainsTokenLocWithGeneratedSource(
 
 TypeRefinementContext *
 TypeRefinementContext::findMostRefinedSubContext(SourceLoc Loc,
-                                                 SourceManager &SM) {
+                                                 ASTContext &Ctx) {
   assert(Loc.isValid());
-  
+
   if (SrcRange.isValid() &&
-      !rangeContainsTokenLocWithGeneratedSource(SM, SrcRange, Loc))
+      !rangeContainsTokenLocWithGeneratedSource(Ctx.SourceMgr, SrcRange, Loc))
     return nullptr;
 
-  // If this context is for a declaration, ensure that we've expanded the
-  // children of the declaration.
-  if (Node.getReason() == Reason::Decl ||
-      Node.getReason() == Reason::DeclImplicit) {
-    if (auto decl = Node.getAsDecl()) {
-      ASTContext &ctx = decl->getASTContext();
-      (void)evaluateOrDefault(
-          ctx.evaluator, ExpandChildTypeRefinementContextsRequest{decl, this},
-          false);
-    }
-  }
+  auto expandedChildren = evaluateOrDefault(
+      Ctx.evaluator, ExpandChildTypeRefinementContextsRequest{this}, {});
 
   // For the moment, we perform a linear search here, but we can and should
   // do something more efficient.
-  for (TypeRefinementContext *Child : Children) {
-    if (auto *Found = Child->findMostRefinedSubContext(Loc, SM)) {
+  for (TypeRefinementContext *Child : expandedChildren) {
+    if (auto *Found = Child->findMostRefinedSubContext(Loc, Ctx)) {
       return Found;
     }
   }
@@ -433,3 +424,18 @@ void swift::simple_display(
     llvm::raw_ostream &out, const TypeRefinementContext *trc) {
   out << "TRC @" << trc;
 }
+
+llvm::Optional<std::vector<TypeRefinementContext *>>
+ExpandChildTypeRefinementContextsRequest::getCachedResult() const {
+  auto *TRC = std::get<0>(getStorage());
+  if (TRC->getNeedsExpansion())
+    return llvm::None;
+  return TRC->Children;
+}
+
+void ExpandChildTypeRefinementContextsRequest::cacheResult(
+    std::vector<TypeRefinementContext *> children) const {
+  auto *TRC = std::get<0>(getStorage());
+  TRC->Children = children;
+  TRC->setNeedsExpansion(false);
+}