[StrictMemorySafety] Check the safety of return types of calls

Previously, we skipped checking the return type of a function for safety
as we expected to warn at the use of the returned value:

  let x = returnsUnsafe()
  usesUnsafe(x) // warn here

Unfortunately, this resulted in missing some unsafe constructs that can
introduce memory safety issues when the use of the return value had a
different shape resulting in false negatives for cases like:

  return returnsUnsafe()

or

  usesUnsafe(returnsUnsafe())

This PR changes the analysis to always take return types of function
calls into account.

rdar://157237301

Author: Gabor Horvath

lib/ClangImporter/ClangImporter.cpp

@@ -8529,7 +8529,7 @@ static bool hasUnsafeType(Evaluator &evaluator, clang::QualType clangType) {
     // All other pointers are considered unsafe.
     return true;
   }
-  
+
   // Handle records recursively.
   if (auto recordDecl = clangType->getAsTagDecl()) {
     // If we reached this point the types is not imported as a shared reference,
@@ -8606,7 +8606,7 @@ ClangDeclExplicitSafety::evaluate(Evaluator &evaluator,
     if (hasUnsafeType(evaluator, field->getType()))
       return ExplicitSafety::Unsafe;
   }
-  
+
   // Okay, call it safe.
   return ExplicitSafety::Safe;
 }

lib/Sema/TypeCheckUnsafe.cpp

@@ -21,11 +21,12 @@
 #include "TypeCheckUnsafe.h"
 
 #include "swift/AST/ASTContext.h"
-#include "swift/AST/Effects.h"
-#include "swift/AST/UnsafeUse.h"
 #include "swift/AST/DiagnosticsSema.h"
+#include "swift/AST/Effects.h"
 #include "swift/AST/PackConformance.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/Types.h"
+#include "swift/AST/UnsafeUse.h"
 
 using namespace swift;
 
@@ -240,19 +241,29 @@ bool swift::enumerateUnsafeUses(ConcreteDeclRef declRef,
   // If the type of this declaration involves unsafe types, diagnose that.
   ASTContext &ctx = decl->getASTContext();
   auto subs = declRef.getSubstitutions();
-  if (!skipTypeCheck) {
-    auto type = decl->getInterfaceType();
-    if (subs) {
-      if (auto *genericFnType = type->getAs<GenericFunctionType>())
-        type = genericFnType->substGenericArgs(subs);
-      else
-        type = type.subst(subs);
-    }
+  auto type = decl->getInterfaceType();
+  if (subs) {
+    if (auto *genericFnType = type->getAs<GenericFunctionType>())
+      type = genericFnType->substGenericArgs(subs);
+    else
+      type = type.subst(subs);
+  }
 
+  if (skipTypeCheck) {
+    // We check the arguements instead of the funcion type for function calls.
+    // On the other hand, we still need to check the return type as we might not
+    // have a declRef with the type of the return type. E.g. in
+    // `return funcionCall()`.
+    if (auto *fnTy = type->getAs<AnyFunctionType>())
+      type = fnTy->getResult();
+    else
+      type = Type();
+  }
+  if (type) {
     bool shouldReturnTrue = false;
     diagnoseUnsafeType(ctx, loc, type, [&](Type unsafeType) {
       if (fn(UnsafeUse::forReferenceToUnsafe(
-               decl, isCall && !isa<ParamDecl>(decl), unsafeType, loc)))
+              decl, isCall && !isa<ParamDecl>(decl), unsafeType, loc)))
         shouldReturnTrue = true;
     });
 

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -162,7 +162,7 @@ public struct CheckedContinuation<T, E: Error>: Sendable {
   /// the caller. The task continues executing when its executor is
   /// able to reschedule it.
   public func resume(returning value: sending T) {
-    if let c: UnsafeContinuation<T, E> = canary.takeContinuation() {
+    if let c: UnsafeContinuation<T, E> = unsafe canary.takeContinuation() {
       unsafe c.resume(returning: value)
     } else {
       #if !$Embedded
@@ -186,7 +186,7 @@ public struct CheckedContinuation<T, E: Error>: Sendable {
   /// the caller. The task continues executing when its executor is
   /// able to reschedule it.
   public func resume(throwing error: __owned E) {
-    if let c: UnsafeContinuation<T, E> = canary.takeContinuation() {
+    if let c: UnsafeContinuation<T, E> = unsafe canary.takeContinuation() {
       unsafe c.resume(throwing: error)
     } else {
       #if !$Embedded

stdlib/public/Concurrency/Executor.swift

@@ -920,7 +920,7 @@ internal func _task_serialExecutor_getExecutorRef<E>(_ executor: E) -> Builtin.E
 @_silgen_name("_task_taskExecutor_getTaskExecutorRef")
 internal func _task_taskExecutor_getTaskExecutorRef<E>(_ taskExecutor: E) -> Builtin.Executor
     where E: TaskExecutor {
-  return taskExecutor.asUnownedTaskExecutor().executor
+  return unsafe taskExecutor.asUnownedTaskExecutor().executor
 }
 
 // Used by the concurrency runtime

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -327,7 +327,7 @@ public struct ExecutorJob: Sendable, ~Copyable {
   /// Returns the result of executing the closure.
   @available(StdlibDeploymentTarget 6.2, *)
   public func withUnsafeExecutorPrivateData<R, E>(body: (UnsafeMutableRawBufferPointer) throws(E) -> R) throws(E) -> R {
-    let base = _jobGetExecutorPrivateData(self.context)
+    let base = unsafe _jobGetExecutorPrivateData(self.context)
     let size = unsafe 2 * MemoryLayout<OpaquePointer>.stride
     return unsafe try body(UnsafeMutableRawBufferPointer(start: base,
                                                          count: size))
@@ -479,21 +479,21 @@ extension ExecutorJob {
 
     /// Allocate a specified number of bytes of uninitialized memory.
     public func allocate(capacity: Int) -> UnsafeMutableRawBufferPointer {
-      let base = _jobAllocate(context, capacity)
+      let base = unsafe _jobAllocate(context, capacity)
       return unsafe UnsafeMutableRawBufferPointer(start: base, count: capacity)
     }
 
     /// Allocate uninitialized memory for a single instance of type `T`.
     public func allocate<T>(as type: T.Type) -> UnsafeMutablePointer<T> {
-      let base = _jobAllocate(context, MemoryLayout<T>.size)
+      let base = unsafe _jobAllocate(context, MemoryLayout<T>.size)
       return unsafe base.bindMemory(to: type, capacity: 1)
     }
 
     /// Allocate uninitialized memory for the specified number of
     /// instances of type `T`.
     public func allocate<T>(capacity: Int, as type: T.Type)
       -> UnsafeMutableBufferPointer<T> {
-      let base = _jobAllocate(context, MemoryLayout<T>.stride * capacity)
+      let base = unsafe _jobAllocate(context, MemoryLayout<T>.stride * capacity)
       let typedBase = unsafe base.bindMemory(to: T.self, capacity: capacity)
       return unsafe UnsafeMutableBufferPointer<T>(start: typedBase, count: capacity)
     }

stdlib/public/Concurrency/Task+PriorityEscalation.swift

@@ -125,7 +125,7 @@ func __withTaskPriorityEscalationHandler0<T, E>(
   onPriorityEscalated handler0: @Sendable (UInt8, UInt8) -> Void,
   isolation: isolated (any Actor)? = #isolation
 ) async throws(E) -> T {
-  let record = _taskAddPriorityEscalationHandler(handler: handler0)
+  let record = unsafe _taskAddPriorityEscalationHandler(handler: handler0)
   defer { unsafe _taskRemovePriorityEscalationHandler(record: record) }
 
   return try await operation()

stdlib/public/Concurrency/Task+TaskExecutor.swift

@@ -146,9 +146,9 @@ public func withTaskExecutorPreference<T, Failure>(
   }
 
   let taskExecutorBuiltin: Builtin.Executor =
-    taskExecutor.asUnownedTaskExecutor().executor
+    unsafe taskExecutor.asUnownedTaskExecutor().executor
 
-  let record = _pushTaskExecutorPreference(taskExecutorBuiltin)
+  let record = unsafe _pushTaskExecutorPreference(taskExecutorBuiltin)
   defer {
     unsafe _popTaskExecutorPreference(record: record)
   }
@@ -177,9 +177,9 @@ public func _unsafeInheritExecutor_withTaskExecutorPreference<T: Sendable>(
   }
 
   let taskExecutorBuiltin: Builtin.Executor =
-    taskExecutor.asUnownedTaskExecutor().executor
+    unsafe taskExecutor.asUnownedTaskExecutor().executor
 
-  let record = _pushTaskExecutorPreference(taskExecutorBuiltin)
+  let record = unsafe _pushTaskExecutorPreference(taskExecutorBuiltin)
   defer {
     unsafe _popTaskExecutorPreference(record: record)
   }

stdlib/public/Concurrency/Task.swift

@@ -984,7 +984,7 @@ internal func _getCurrentTaskName() -> UnsafePointer<UInt8>?
 
 @available(SwiftStdlib 6.2, *)
 internal func _getCurrentTaskNameString() -> String? {
-  if let stringPtr = _getCurrentTaskName() {
+  if let stringPtr = unsafe _getCurrentTaskName() {
     unsafe String(cString: stringPtr)
   } else {
     nil

stdlib/public/Concurrency/TaskCancellation.swift

@@ -77,7 +77,7 @@ public func withTaskCancellationHandler<T>(
 ) async rethrows -> T {
   // unconditionally add the cancellation record to the task.
   // if the task was already cancelled, it will be executed right away.
-  let record = _taskAddCancellationHandler(handler: handler)
+  let record = unsafe _taskAddCancellationHandler(handler: handler)
   defer { unsafe _taskRemoveCancellationHandler(record: record) }
 
   return try await operation()
@@ -98,7 +98,7 @@ public func _unsafeInheritExecutor_withTaskCancellationHandler<T>(
 ) async rethrows -> T {
   // unconditionally add the cancellation record to the task.
   // if the task was already cancelled, it will be executed right away.
-  let record = _taskAddCancellationHandler(handler: handler)
+  let record = unsafe _taskAddCancellationHandler(handler: handler)
   defer { unsafe _taskRemoveCancellationHandler(record: record) }
 
   return try await operation()

stdlib/public/Concurrency/TaskLocal.swift

@@ -178,7 +178,7 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   /// or if the task-local has no value bound, this will return the `defaultValue`
   /// of the task local.
   public func get() -> Value {
-    guard let rawValue = _taskLocalValueGet(key: key) else {
+    guard let rawValue = unsafe _taskLocalValueGet(key: key) else {
       return self.defaultValue
     }