[SE-0458] withoutActuallyEscaping is unsafe for @convention(block)

The implementation of `withoutActuallyEscaping` for `@convention(block)`
functions cannot verify at runtime that the function did not actually
escape. Diagnose this as unsafe code under strict memory safety checking.

Fixes rdar://139994149.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -8204,6 +8204,8 @@ NOTE(note_use_of_unsafe_conformance_is_unsafe,none,
      (Type, const ValueDecl *))
 NOTE(note_unsafe_storage,none,
      "%kindbase1 involves unsafe type %2", (bool, const ValueDecl *, Type))
+NOTE(note_unsafe_temporarily_escaping,none,
+     "'withoutActuallyEscaping' function of type %0 is unsafe", (Type))
 
 GROUPED_WARNING(decl_unsafe_storage,StrictMemorySafety,none,
       "%kindbase0 has storage involving unsafe types",

include/swift/AST/UnsafeUse.h

@@ -14,6 +14,7 @@
 #define SWIFT_AST_UNSAFEUSE_H
 
 #include "swift/AST/Decl.h"
+#include "swift/AST/Expr.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/ProtocolConformanceRef.h"
 #include "swift/AST/Type.h"
@@ -54,6 +55,9 @@ class UnsafeUse {
     CallToUnsafe,
     /// A @preconcurrency import.
     PreconcurrencyImport,
+    /// A use of withoutActuallyEscaping that lacks enforcement that the
+    /// value didn't actually escape.
+    TemporarilyEscaping,
   };
 
 private:
@@ -85,6 +89,8 @@ class UnsafeUse {
       const void *location;
     } entity;
 
+    MakeTemporarilyEscapableExpr *temporarilyEscaping;
+
     const ImportDecl *importDecl;
   } storage;
 
@@ -195,6 +201,12 @@ class UnsafeUse {
                         decl, type, location);
   }
 
+  static UnsafeUse forTemporarilyEscaping(MakeTemporarilyEscapableExpr *expr) {
+    UnsafeUse result(TemporarilyEscaping);
+    result.storage.temporarilyEscaping = expr;
+    return result;
+  }
+
   static UnsafeUse forPreconcurrencyImport(const ImportDecl *importDecl) {
     UnsafeUse result(PreconcurrencyImport);
     result.storage.importDecl = importDecl;
@@ -230,6 +242,9 @@ class UnsafeUse {
       return SourceLoc(
           llvm::SMLoc::getFromPointer((const char *)storage.entity.location));
 
+    case TemporarilyEscaping:
+      return storage.temporarilyEscaping->getLoc();
+
     case PreconcurrencyImport:
       return storage.importDecl->getLoc();
     }
@@ -240,6 +255,7 @@ class UnsafeUse {
     switch (getKind()) {
     case Override:
     case Witness:
+    case TemporarilyEscaping:
     case PreconcurrencyImport:
       // Cannot replace location.
       return;
@@ -283,6 +299,7 @@ class UnsafeUse {
       return storage.entity.decl;
 
     case UnsafeConformance:
+    case TemporarilyEscaping:
       return nullptr;
 
     case PreconcurrencyImport:
@@ -315,6 +332,7 @@ class UnsafeUse {
     case CallToUnsafe:
     case UnsafeConformance:
     case PreconcurrencyImport:
+    case TemporarilyEscaping:
       return nullptr;
     }
   }
@@ -341,6 +359,9 @@ class UnsafeUse {
     case ReferenceToUnsafeThroughTypealias:
     case CallToUnsafe:
       return storage.entity.type;
+
+    case TemporarilyEscaping:
+      return storage.temporarilyEscaping->getOpaqueValue()->getType();
     }
   }
 
@@ -365,6 +386,7 @@ class UnsafeUse {
     case ReferenceToUnsafeStorage:
     case ReferenceToUnsafeThroughTypealias:
     case CallToUnsafe:
+    case TemporarilyEscaping:
     case PreconcurrencyImport:
       return ProtocolConformanceRef::forInvalid();
     }

lib/Sema/TypeCheckEffects.cpp

@@ -644,7 +644,10 @@ class EffectsHandlingWalker : public ASTWalker {
           break;
         }
       }
+    } else if (auto TE = dyn_cast<MakeTemporarilyEscapableExpr>(E)) {
+      recurse = asImpl().checkTemporarilyEscapable(TE);
     }
+
     // Error handling validation (via checkTopLevelEffects) happens after
     // type checking. If an unchecked expression is still around, the code was
     // invalid.
@@ -1090,6 +1093,15 @@ class Classification {
     return result;
   }
 
+  /// Return a potentially-unsafe classification by recording all of the
+  /// provided unsafe uses.
+  static Classification forUnsafe(ArrayRef<UnsafeUse> unsafeUses) {
+    Classification result;
+    for (const auto &unsafeUse : unsafeUses)
+      result.recordUnsafeUse(unsafeUse);
+    return result;
+  }
+
   /// Return a classification for a given declaration reference.
   static Classification
   forDeclRef(ConcreteDeclRef declRef, ConditionalEffectKind conditionalKind,
@@ -1856,6 +1868,38 @@ class ApplyClassifier {
     return result;
   }
 
+  Classification classifyTemporarilyEscapable(MakeTemporarilyEscapableExpr *E) {
+    auto opaqueValue = E->getOpaqueValue();
+    if (!opaqueValue)
+      return Classification();
+
+    auto type = opaqueValue->getType();
+    if (!type)
+      return Classification();
+
+    auto fnType = type->getAs<AnyFunctionType>();
+    if (!fnType)
+      return Classification();
+
+    // Runtime safety checks for temporarily-escapable functions aren't
+    // always available.
+    switch (fnType->getRepresentation()) {
+    case FunctionTypeRepresentation::Swift:
+      // Swift performs runtime checking to ensure that the function did not
+      // escape.
+      return Classification();
+
+    case FunctionTypeRepresentation::Block:
+      return Classification::forUnsafe({UnsafeUse::forTemporarilyEscaping(E)});
+
+    case FunctionTypeRepresentation::Thin:
+    case FunctionTypeRepresentation::CFunctionPointer:
+      // Thin and C functions cannot carry any state with them, so escaping
+      // the pointers does not introduce a memory-safety issue.
+      return Classification();
+    }
+  }
+
 private:
   /// Classify a throwing or async function according to our local
   /// knowledge of its implementation.
@@ -2073,6 +2117,10 @@ class ApplyClassifier {
       return ShouldRecurse;
     }
 
+    ShouldRecurse_t checkTemporarilyEscapable(MakeTemporarilyEscapableExpr *E) {
+      return ShouldRecurse;
+    }
+
     ConditionalEffectKind checkExhaustiveDoBody(DoCatchStmt *S) {
       // All errors thrown by the do body are caught, but any errors thrown
       // by the catch bodies are bounded by the throwing kind of the do body.
@@ -2214,6 +2262,10 @@ class ApplyClassifier {
       return ShouldRecurse;
     }
 
+    ShouldRecurse_t checkTemporarilyEscapable(MakeTemporarilyEscapableExpr *E) {
+      return ShouldRecurse;
+    }
+
     void visitExprPre(Expr *expr) { return; }
   };
 
@@ -2327,6 +2379,11 @@ class ApplyClassifier {
       return ShouldRecurse;
     }
 
+    ShouldRecurse_t checkTemporarilyEscapable(MakeTemporarilyEscapableExpr *E) {
+      classification.merge(Self.classifyTemporarilyEscapable(E));
+      return ShouldRecurse;
+    }
+
     void visitExprPre(Expr *expr) { return; }
   };
 
@@ -3774,6 +3831,13 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     return ShouldRecurse;
   }
 
+  ShouldRecurse_t checkTemporarilyEscapable(MakeTemporarilyEscapableExpr *E) {
+    auto classification = getApplyClassifier().
+        classifyTemporarilyEscapable(E);
+    checkEffectSite(E, /*requiresTry=*/false, classification);
+    return ShouldRecurse;
+  }
+
   ConditionalEffectKind checkExhaustiveDoBody(DoCatchStmt *S) {
     // This is a context where errors are handled.
     ContextScope scope(*this, CurContext.withHandlesErrors());

lib/Sema/TypeCheckUnsafe.cpp

@@ -154,6 +154,14 @@ void swift::diagnoseUnsafeUse(const UnsafeUse &use) {
     return;
   }
 
+  case UnsafeUse::TemporarilyEscaping: {
+    Type type = use.getType();
+    ASTContext &ctx = type->getASTContext();
+    ctx.Diags.diagnose(
+        use.getLocation(), diag::note_unsafe_temporarily_escaping, type);
+    return;
+  }
+
   case UnsafeUse::PreconcurrencyImport: {
     auto importDecl = cast<ImportDecl>(use.getDecl());
     importDecl->diagnose(diag::preconcurrency_import_unsafe);

test/Unsafe/unsafe_withoutactuallyescaping.swift

@@ -0,0 +1,19 @@
+// RUN: %target-typecheck-verify-swift -strict-memory-safety
+
+// REQUIRES: objc_interop
+
+func callMe(body: @escaping @convention(block) () -> Void) { }
+
+func callMeSwiftly(body: @escaping () -> Void) { }
+
+func testCall(body: @convention(block) () -> Void, swiftBody: () -> Void) {
+  // expected-warning@+2{{unsafe}}{{3-3=unsafe }}
+  // expected-note@+1{{'withoutActuallyEscaping' function of type '@convention(block) () -> Void' is unsafe}}
+  withoutActuallyEscaping(body) { nonescapingBody in
+    callMe(body: nonescapingBody)
+  }
+
+  withoutActuallyEscaping(swiftBody) { nonescapingBody in
+    callMeSwiftly(body: nonescapingBody)
+  }
+}