Sema: Fix availability checking in function default argument closure expressions.

When building the TypeRefinementContext tree, delay pushing a decl's body context until the body Stmt is actually encountered by the `ASTWalker`. This ensures that default argument expressions do not automatically inherit the availability context of the function body.

Resolves rdar://95148645

Author: tshortli

lib/Sema/TypeCheckAvailability.cpp

@@ -356,21 +356,34 @@ namespace {
 /// A class to walk the AST to build the type refinement context hierarchy.
 class TypeRefinementContextBuilder : private ASTWalker {
 
+  ASTContext &Context;
+
+  /// Represents an entry in a stack of active type refinement contexts. The
+  /// stack is used to facilitate building the TRC's tree structure. A new TRC
+  /// is pushed onto this stack before visiting children whenever the current
+  /// AST node requires a new context and the TRC is then popped
+  /// post-visitation.
   struct ContextInfo {
     TypeRefinementContext *TRC;
 
-    /// The node whose end marks the end of the refinement context.
-    /// If the builder sees this node in a post-visitor, it will pop
-    /// the context from the stack. This node can be null (ParentTy()),
+    /// The AST node. This node can be null (ParentTy()),
     /// indicating that custom logic elsewhere will handle removing
     /// the context when needed.
     ParentTy ScopeNode;
 
     bool ContainedByDeploymentTarget;
   };
-
   std::vector<ContextInfo> ContextStack;
-  ASTContext &Context;
+
+  /// Represents an entry in a stack of pending decl body type refinement
+  /// contexts. TRCs in this stack should be pushed onto \p ContextStack when
+  /// \p BodyStmt is encountered.
+  struct DeclBodyContextInfo {
+    TypeRefinementContext *TRC;
+    Decl *Decl;
+    Stmt *BodyStmt;
+  };
+  std::vector<DeclBodyContextInfo> DeclBodyContextStack;
 
   /// A mapping from abstract storage declarations with accessors to
   /// to the type refinement contexts for those declarations. We refer to
@@ -411,6 +424,15 @@ class TypeRefinementContextBuilder : private ASTWalker {
     ContextStack.push_back(Info);
   }
 
+  void pushDeclBodyContext(TypeRefinementContext *TRC, Decl *D, Stmt *S) {
+    DeclBodyContextInfo Info;
+    Info.TRC = TRC;
+    Info.Decl = D;
+    Info.BodyStmt = S;
+
+    DeclBodyContextStack.push_back(Info);
+  }
+
   const char *stackTraceAction() const {
     return "building type refinement context for";
   }
@@ -474,6 +496,12 @@ class TypeRefinementContextBuilder : private ASTWalker {
     while (ContextStack.back().ScopeNode.getAsDecl() == D) {
       ContextStack.pop_back();
     }
+
+    while (!DeclBodyContextStack.empty() &&
+           DeclBodyContextStack.back().Decl == D) {
+      DeclBodyContextStack.pop_back();
+    }
+
     return true;
   }
 
@@ -689,18 +717,22 @@ class TypeRefinementContextBuilder : private ASTWalker {
 
     // Top level code always uses the deployment target.
     if (auto tlcd = dyn_cast<TopLevelCodeDecl>(D)) {
-      auto *topLevelTRC = createContext(tlcd, tlcd->getSourceRange());
-      pushContext(topLevelTRC, D);
+      if (auto bodyStmt = tlcd->getBody()) {
+        pushDeclBodyContext(createContext(tlcd, tlcd->getSourceRange()), tlcd,
+                            bodyStmt);
+      }
       return;
     }
 
     // Function bodies use the deployment target if they are within the module's
     // resilience domain.
     if (auto afd = dyn_cast<AbstractFunctionDecl>(D)) {
-      if (!afd->isImplicit() && afd->getBodySourceRange().isValid() &&
+      if (!afd->isImplicit() &&
           afd->getResilienceExpansion() != ResilienceExpansion::Minimal) {
-        auto *functionBodyTRC = createContext(afd, afd->getBodySourceRange());
-        pushContext(functionBodyTRC, D);
+        if (auto body = afd->getBody(/*canSynthesize*/ false)) {
+          pushDeclBodyContext(createContext(afd, afd->getBodySourceRange()),
+                              afd, body);
+        }
       }
       return;
     }
@@ -748,6 +780,10 @@ class TypeRefinementContextBuilder : private ASTWalker {
   std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
     PrettyStackTraceStmt trace(Context, stackTraceAction(), S);
 
+    if (consumeDeclBodyContextIfNecessary(S)) {
+      return std::make_pair(true, S);
+    }
+
     if (auto *IS = dyn_cast<IfStmt>(S)) {
       buildIfStmtRefinementContext(IS);
       return std::make_pair(false, S);
@@ -778,6 +814,22 @@ class TypeRefinementContextBuilder : private ASTWalker {
     return S;
   }
 
+  /// Consumes the top TRC from \p DeclBodyContextStack and pushes it onto the
+  /// \p Context stack if the given \p Stmt is the matching body statement.
+  /// Returns \p true if a context was pushed.
+  bool consumeDeclBodyContextIfNecessary(Stmt *S) {
+    if (DeclBodyContextStack.empty())
+      return false;
+
+    auto Info = DeclBodyContextStack.back();
+    if (S != Info.BodyStmt)
+      return false;
+
+    pushContext(Info.TRC, Info.BodyStmt);
+    DeclBodyContextStack.pop_back();
+    return true;
+  }
+
   /// Builds the type refinement hierarchy for the IfStmt if the guard
   /// introduces a new refinement context for the Then branch.
   /// There is no need for the caller to explicitly traverse the children

test/attr/attr_inlinable_available.swift

@@ -655,7 +655,7 @@ public func backDeployedToInliningTarget(
 
 // MARK: - Default arguments
 
-// Default arguments act like @inlinable.
+// Default arguments act like @inlinable when in a public function.
 
 public func defaultArgsUseNoAvailable( // expected-note 3 {{add @available attribute}}
   _: Any = NoAvailable.self,
@@ -666,6 +666,15 @@ public func defaultArgsUseNoAvailable( // expected-note 3 {{add @available attri
   _: Any = AfterDeploymentTarget.self // expected-error {{'AfterDeploymentTarget' is only available in macOS 11 or newer}}
 ) {}
 
+func defaultArgsUseInternal( // expected-note {{add @available attribute}}
+  _: Any = NoAvailable.self,
+  _: Any = BeforeInliningTarget.self,
+  _: Any = AtInliningTarget.self,
+  _: Any = BetweenTargets.self,
+  _: Any = AtDeploymentTarget.self,
+  _: Any = AfterDeploymentTarget.self // expected-error {{'AfterDeploymentTarget' is only available in macOS 11 or newer}}
+) {}
+
 @available(macOS, unavailable)
 public func defaultArgsUseUnavailable(
   _: Any = NoAvailable.self,
@@ -687,6 +696,46 @@ public func spiDefaultArgsUseNoAvailable( // expected-note 1 {{add @available at
   _: Any = AfterDeploymentTarget.self // expected-error {{'AfterDeploymentTarget' is only available in macOS 11 or newer}}
 ) {}
 
+// Verify that complex default argument expressions are checked appropriately.
+public func defaultArgsClosureExprNoAvailable( // expected-note 3 {{add @available attribute}}
+  _: Int = {
+    _ = NoAvailable.self
+    _ = BeforeInliningTarget.self
+    _ = AtInliningTarget.self
+    _ = BetweenTargets.self // expected-error {{'BetweenTargets' is only available in macOS 10.14.5 or newer; clients of 'Test' may have a lower deployment target}} expected-note {{add 'if #available' version check}}
+    _ = AtDeploymentTarget.self // expected-error {{'AtDeploymentTarget' is only available in macOS 10.15 or newer; clients of 'Test' may have a lower deployment target}} expected-note {{add 'if #available' version check}}
+    _ = AfterDeploymentTarget.self // expected-error {{'AfterDeploymentTarget' is only available in macOS 11 or newer}} expected-note {{add 'if #available' version check}}
+    if #available(macOS 10.14.5, *) {
+      _ = BetweenTargets.self
+    }
+    if #available(macOS 10.15, *) {
+      _ = AtDeploymentTarget.self
+    }
+    if #available(macOS 11, *) {
+      _ = AfterDeploymentTarget.self
+    }
+    return 42
+  }()
+) {}
+
+func defaultArgsClosureExprInternal( // expected-note {{add @available attribute}}
+  _: Int = {
+    _ = NoAvailable.self
+    _ = BeforeInliningTarget.self
+    _ = AtInliningTarget.self
+    _ = BetweenTargets.self
+    _ = AtDeploymentTarget.self
+    _ = AfterDeploymentTarget.self // expected-error {{'AfterDeploymentTarget' is only available in macOS 11 or newer}} expected-note {{add 'if #available' version check}}
+    if #available(macOS 11, *) {
+      _ = AfterDeploymentTarget.self
+    }
+    return 42
+  }()
+) {}
+
+
+// MARK: - Properties
+
 @propertyWrapper
 public struct PropertyWrapper<T> {
   public var wrappedValue: T