AST/Sema: Refactor lazy TypeRefinementContext expansion.

When building out the TypeRefinementContext tree, if we encounter an AST node
that we cannot fully build out the children for, instead add a leaf node to the
tree and mark it as needing expansion. Then skip walking the children of the
AST node. When querying the tree, if when a node needing expansion is
encountered, expand the children of the node on-demand. The previous
implementation of lazy tree expansion did not actually skip the children of AST
nodes requiring expansion, so the resulting TypeRefinementContext tree could
contain malformed nodes with duplicate descendants.

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);
+}

lib/Sema/TypeCheckAvailability.cpp

@@ -472,6 +472,17 @@ class TypeRefinementContextBuilder : private ASTWalker {
   PreWalkAction walkToDeclPre(Decl *D) override {
     PrettyStackTraceDecl trace(stackTraceAction(), D);
 
+    // Implicit decls don't have source locations so they cannot have a TRC.
+    if (D->isImplicit())
+      return Action::Continue();
+
+    // The AST of this decl may not be ready to traverse yet if it hasn't been
+    // full typechecked. If that's the case, we leave a placeholder node in the
+    // tree to indicate that the subtree should be expanded lazily when it
+    // needs to be traversed.
+    if (buildLazyContextForDecl(D))
+      return Action::SkipChildren();
+
     // Adds in a TRC that covers the entire declaration.
     if (auto DeclTRC = getNewContextForSignatureOfDecl(D)) {
       pushContext(DeclTRC, D);
@@ -497,6 +508,44 @@ class TypeRefinementContextBuilder : private ASTWalker {
     return Action::Continue();
   }
 
+  /// Constructs a placeholder TRC node that should be expanded later if the AST
+  /// associated with the given declaration is not ready to be traversed yet.
+  /// Returns true if a node was created.
+  bool buildLazyContextForDecl(Decl *D) {
+    if (!isa<PatternBindingDecl>(D))
+      return false;
+
+    // Check whether the current TRC is already a lazy placeholder. If it is,
+    // we should try to expand it rather than creating a new placeholder.
+    auto currentTRC = getCurrentTRC();
+    if (currentTRC->getNeedsExpansion() && currentTRC->getDeclOrNull() == D)
+      return false;
+
+    // Pattern binding declarations may have attached property wrappers that
+    // get expanded from macros attached to the parent declaration. We must
+    // not eagerly expand the attached property wrappers to avoid a request
+    // cycle.
+    if (auto *pattern = dyn_cast<PatternBindingDecl>(D)) {
+      if (auto firstVar = pattern->getAnchoringVarDecl(0)) {
+        // If there's no initial value, or the init is exposed to clients, then
+        // we don't need to create any implicit TRCs for the init bodies.
+        if (!firstVar->hasInitialValue() || firstVar->isInitExposedToClients())
+          return false;
+
+        // FIXME: We could narrow this further by detecting whether there are
+        // any macro expansions required to visit the CustomAttrs of the var.
+      }
+    }
+
+    // If we've made it this far then we've identified a declaration that
+    // requires lazy expansion later.
+    auto lazyTRC = TypeRefinementContext::createForDeclImplicit(
+        Context, D, currentTRC, currentTRC->getAvailabilityInfo(),
+        refinementSourceRangeForDecl(D));
+    lazyTRC->setNeedsExpansion(true);
+    return true;
+  }
+
   /// Returns a new context to be introduced for the declaration, or nullptr
   /// if no new context should be introduced.
   TypeRefinementContext *getNewContextForSignatureOfDecl(Decl *D) {
@@ -519,10 +568,6 @@ class TypeRefinementContextBuilder : private ASTWalker {
         return nullptr;
     }
 
-    // Ignore implicit declarations (mainly skips over `DeferStmt` functions).
-    if (D->isImplicit())
-      return nullptr;
-
     // Declarations with an explicit availability attribute always get a TRC.
     if (hasActiveAvailableAttribute(D, Context)) {
       AvailabilityContext DeclaredAvailability =
@@ -539,17 +584,11 @@ class TypeRefinementContextBuilder : private ASTWalker {
     // internal property in a public struct can be effectively less available
     // than the containing struct decl because the internal property will only
     // be accessed by code running at the deployment target or later.
-    //
-    // For declarations that have their child construction delayed, always
-    // create this implicit declaration context. It will be used to trigger
-    // lazy creation of the child TRCs.
     AvailabilityContext CurrentAvailability =
         getCurrentTRC()->getAvailabilityInfo();
     AvailabilityContext EffectiveAvailability =
         getEffectiveAvailabilityForDeclSignature(D, CurrentAvailability);
-    if ((isa<PatternBindingDecl>(D) &&
-         refinementSourceRangeForDecl(D).isValid()) ||
-        CurrentAvailability.isSupersetOf(EffectiveAvailability))
+    if (CurrentAvailability.isSupersetOf(EffectiveAvailability))
       return TypeRefinementContext::createForDeclImplicit(
           Context, D, getCurrentTRC(), EffectiveAvailability,
           refinementSourceRangeForDecl(D));
@@ -732,6 +771,11 @@ class TypeRefinementContextBuilder : private ASTWalker {
       }
       return;
     }
+
+    if (auto pattern = dyn_cast<PatternBindingDecl>(D)) {
+      buildContextsForPatternBindingDecl(pattern);
+      return;
+    }
   }
 
   PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
@@ -1208,7 +1252,7 @@ void TypeChecker::buildTypeRefinementContextHierarchyDelayed(SourceFile &SF, Abs
 
   // Build the refinement context for the function body.
   ASTContext &Context = SF.getASTContext();
-  auto LocalTRC = RootTRC->findMostRefinedSubContext(AFD->getLoc(), Context.SourceMgr);
+  auto LocalTRC = RootTRC->findMostRefinedSubContext(AFD->getLoc(), Context);
   TypeRefinementContextBuilder Builder(LocalTRC, Context);
   Builder.build(AFD);
 }
@@ -1224,23 +1268,16 @@ TypeChecker::getOrBuildTypeRefinementContext(SourceFile *SF) {
   return TRC;
 }
 
-bool ExpandChildTypeRefinementContextsRequest::evaluate(
-    Evaluator &evaluator, Decl *decl, TypeRefinementContext *parentTRC
-) const {
-  // If the parent TRC is already contained by the deployment target, there's
-  // nothing more to do.
-  ASTContext &ctx = decl->getASTContext();
-  if (computeContainedByDeploymentTarget(parentTRC, ctx))
-    return false;
-
-  // Pattern binding declarations can have children corresponding to property
-  // wrappers and the initial values provided in each pattern binding entry.
-  if (auto *binding = dyn_cast<PatternBindingDecl>(decl)) {
+std::vector<TypeRefinementContext *>
+ExpandChildTypeRefinementContextsRequest::evaluate(
+    Evaluator &evaluator, TypeRefinementContext *parentTRC) const {
+  assert(parentTRC->getNeedsExpansion());
+  if (auto decl = parentTRC->getDeclOrNull()) {
+    ASTContext &ctx = decl->getASTContext();
     TypeRefinementContextBuilder builder(parentTRC, ctx);
-    builder.buildContextsForPatternBindingDecl(binding);
+    builder.build(decl);
   }
-
-  return false;
+  return parentTRC->Children;
 }
 
 AvailabilityContext
@@ -1293,7 +1330,7 @@ TypeChecker::overApproximateAvailabilityAtLocation(SourceLoc loc,
   if (SF && loc.isValid()) {
     TypeRefinementContext *rootTRC = getOrBuildTypeRefinementContext(SF);
     TypeRefinementContext *TRC =
-        rootTRC->findMostRefinedSubContext(loc, Context.SourceMgr);
+        rootTRC->findMostRefinedSubContext(loc, Context);
     OverApproximateContext.constrainWith(TRC->getAvailabilityInfo());
     if (MostRefined) {
       *MostRefined = TRC;