[Function builders] Use buildLimitedAvailability() for #available block

The use of "if #available" in function builders can subvert availability
checking if the function builder carries all type information for the
values within the "then" block outside of the "else" block. Tighten up
the model in two ways:

* Check whether the type coming out of an "if #available" references
any declarations that are not available in the outer context, to close
up the model.
* If the function builder provides a buildLimitedAvailability(_:)
operation, call that on the result of the "then" block in an "if
that it cannot leak out of the "if #available"; if it doesn't, the
check above will still fire.

Stage this in with a warning so function builders out there in the wild
can adapt. We'll upgrade the warning to an error later.

Fixes rdar://problem/65021017.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -5127,6 +5127,10 @@ NOTE(function_builder_infer_add_return, none,
 NOTE(function_builder_infer_pick_specific, none,
      "apply function builder %0 (inferred from %select{protocol|dynamic replacement of}1 %2)",
      (Type, unsigned, DeclName))
+WARNING(function_builder_missing_limited_availability, none,
+        "function builder %0 does not implement `buildLimitedAvailability`; "
+        "this code may crash on earlier versions of the OS",
+        (Type))
 
 //------------------------------------------------------------------------------
 // MARK: Tuple Shuffle Diagnostics

include/swift/AST/KnownIdentifiers.def

@@ -38,6 +38,7 @@ IDENTIFIER(buildEither)
 IDENTIFIER(buildExpression)
 IDENTIFIER(buildFinalResult)
 IDENTIFIER(buildIf)
+IDENTIFIER(buildLimitedAvailability)
 IDENTIFIER(buildOptional)
 IDENTIFIER(callAsFunction)
 IDENTIFIER(Change)

lib/Sema/BuilderTransform.cpp

@@ -19,6 +19,7 @@
 #include "MiscDiagnostics.h"
 #include "SolutionResult.h"
 #include "TypeChecker.h"
+#include "TypeCheckAvailability.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/NameLookup.h"
@@ -38,6 +39,24 @@ using namespace constraints;
 
 namespace {
 
+/// Find the first #available condition within the statement condition,
+/// or return NULL if there isn't one.
+const StmtConditionElement *findAvailabilityCondition(StmtCondition stmtCond) {
+  for (const auto &cond : stmtCond) {
+    switch (cond.getKind()) {
+    case StmtConditionElement::CK_Boolean:
+    case StmtConditionElement::CK_PatternBinding:
+      continue;
+
+    case StmtConditionElement::CK_Availability:
+      return &cond;
+      break;
+    }
+  }
+
+  return nullptr;
+}
+
 /// Visitor to classify the contents of the given closure.
 class BuilderClosureVisitor
     : private StmtVisitor<BuilderClosureVisitor, VarDecl *> {
@@ -498,10 +517,20 @@ class BuilderClosureVisitor
     if (!cs || !thenVar || (elseChainVar && !*elseChainVar))
       return nullptr;
 
+    // If there is a #available in the condition, the 'then' will need to
+    // be wrapped in a call to buildAvailabilityErasure(_:), if available.
+    Expr *thenVarRefExpr = buildVarRef(
+        thenVar, ifStmt->getThenStmt()->getEndLoc());
+    if (findAvailabilityCondition(ifStmt->getCond()) &&
+        builderSupports(ctx.Id_buildLimitedAvailability)) {
+      thenVarRefExpr = buildCallIfWanted(
+          ifStmt->getThenStmt()->getEndLoc(), ctx.Id_buildLimitedAvailability,
+          { thenVarRefExpr }, { Identifier() });
+    }
+
     // Prepare the `then` operand by wrapping it to produce a chain result.
     Expr *thenExpr = buildWrappedChainPayload(
-        buildVarRef(thenVar, ifStmt->getThenStmt()->getEndLoc()),
-        payloadIndex, numPayloads, isOptional);
+        thenVarRefExpr, payloadIndex, numPayloads, isOptional);
 
     // Prepare the `else operand:
     Expr *elseExpr;
@@ -1175,6 +1204,35 @@ class BuilderClosureRewriter
             capturedThen.first, {capturedThen.second.front()}));
     ifStmt->setThenStmt(newThen);
 
+    // Look for a #available condition. If there is one, we need to check
+    // that the resulting type of the "then" does refer to any types that
+    // are unavailable in the enclosing context.
+    //
+    // Note that this is for staging in support for
+    if (auto availabilityCond = findAvailabilityCondition(ifStmt->getCond())) {
+      SourceLoc loc = availabilityCond->getStartLoc();
+      Type thenBodyType = solution.simplifyType(
+          solution.getType(target.captured.second[0]));
+      thenBodyType.findIf([&](Type type) {
+        auto nominal = type->getAnyNominal();
+        if (!nominal)
+          return false;
+
+        if (auto reason = TypeChecker::checkDeclarationAvailability(
+                              nominal, loc, dc)) {
+          // Note that the problem is with the function builder, not the body.
+          // This is for staging only. We want to disable #available in
+          // function builders that don't support this operation.
+          ctx.Diags.diagnose(
+              loc, diag::function_builder_missing_limited_availability,
+              builderTransform.builderType);
+          return true;
+        }
+
+        return false;
+      });
+    }
+
     if (auto elseBraceStmt =
             dyn_cast_or_null<BraceStmt>(ifStmt->getElseStmt())) {
       // Translate the "else" branch when it's a stmt-brace.

test/Constraints/function_builder_availability.swift

@@ -16,7 +16,7 @@ struct TupleBuilder {
   static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
     return (t1, t2)
   }
-  
+
   static func buildBlock<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
       -> (T1, T2, T3) {
     return (t1, t2, t3)
@@ -55,14 +55,17 @@ func globalFuncAvailableOn10_9() -> Int { return 9 }
 func globalFuncAvailableOn10_51() -> Int { return 10 }
 
 @available(OSX, introduced: 10.52)
-func globalFuncAvailableOn10_52() -> Int { return 11 }
+struct Only10_52 { }
+
+@available(OSX, introduced: 10.52)
+func globalFuncAvailableOn10_52() -> Only10_52 { .init() }
 
 tuplify(true) { cond in
   globalFuncAvailableOn10_9()
   if #available(OSX 10.51, *) {
     globalFuncAvailableOn10_51()
     tuplify(false) { cond2 in
-      if cond, #available(OSX 10.52, *) {
+      if cond, #available(OSX 10.52, *) { // expected-warning{{function builder 'TupleBuilder' does not implement `buildLimitedAvailability`; this code may crash on earlier versions of the OS}}
         cond2
         globalFuncAvailableOn10_52()
       } else {
@@ -72,3 +75,55 @@ tuplify(true) { cond in
     }
   }
 }
+
+// Function builder that can perform type erasure for #available.
+@_functionBuilder
+struct TupleBuilderAvailability {
+  static func buildBlock<T1>(_ t1: T1) -> (T1) {
+    return (t1)
+  }
+
+  static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
+    return (t1, t2)
+  }
+
+  static func buildBlock<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
+      -> (T1, T2, T3) {
+    return (t1, t2, t3)
+  }
+
+  static func buildBlock<T1, T2, T3, T4>(_ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4)
+      -> (T1, T2, T3, T4) {
+    return (t1, t2, t3, t4)
+  }
+
+  static func buildBlock<T1, T2, T3, T4, T5>(
+    _ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4, _ t5: T5
+  ) -> (T1, T2, T3, T4, T5) {
+    return (t1, t2, t3, t4, t5)
+  }
+
+  static func buildDo<T>(_ value: T) -> T { return value }
+  static func buildIf<T>(_ value: T?) -> T? { return value }
+
+  static func buildEither<T,U>(first value: T) -> Either<T,U> {
+    return .first(value)
+  }
+  static func buildEither<T,U>(second value: U) -> Either<T,U> {
+    return .second(value)
+  }
+
+  static func buildLimitedAvailability<T>(_ value: T) -> Any {
+    return value
+  }
+}
+
+func tuplifyWithAvailabilityErasure<T>(_ cond: Bool, @TupleBuilderAvailability body: (Bool) -> T) {
+  print(body(cond))
+}
+
+tuplifyWithAvailabilityErasure(true) { cond in
+  if cond, #available(OSX 10.52, *) {
+    globalFuncAvailableOn10_52()
+  }
+}