SILOptimizer: avoid stack protection in two cases where it's not needed

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -486,13 +486,23 @@ private extension Instruction {
         if !atp.needsStackProtection {
           return nil
         }
+        var hasNoStores = NoStores()
+        if hasNoStores.walkDownUses(ofValue: atp, path: SmallProjectionPath()) == .continueWalk {
+          return nil
+        }
+
         // The result of an `address_to_pointer` may be used in any unsafe way, e.g.
         // passed to a C function.
         baseAddr = atp.operand
       case let ia as IndexAddrInst:
         if !ia.needsStackProtection {
           return nil
         }
+        var hasNoStores = NoStores()
+        if hasNoStores.walkDownUses(ofAddress: ia, path: SmallProjectionPath()) == .continueWalk {
+          return nil
+        }
+
         // `index_addr` is unsafe if not used for tail-allocated elements (e.g. in Array).
         baseAddr = ia.base
       default:
@@ -509,6 +519,29 @@ private extension Instruction {
   }
 }
 
+/// Checks if there are no stores to an address or raw pointer.
+private struct NoStores : ValueDefUseWalker, AddressDefUseWalker {
+  var walkDownCache = WalkerCache<SmallProjectionPath>()
+
+  mutating func leafUse(value: Operand, path: SmallProjectionPath) -> WalkResult {
+    if let ptai = value.instruction as? PointerToAddressInst {
+      return walkDownUses(ofAddress: ptai, path: path)
+    }
+    return .abortWalk
+  }
+
+  mutating func leafUse(address: Operand, path: SmallProjectionPath) -> WalkResult {
+    switch address.instruction {
+    case is LoadInst:
+      return .continueWalk
+    case let cai as CopyAddrInst:
+      return address == cai.sourceOperand ? .continueWalk : .abortWalk
+    default:
+      return .abortWalk
+    }
+  }
+}
+
 private extension Function {
   func setNeedsStackProtection(_ context: FunctionPassContext) {
     if !needsStackProtection {

include/swift/AST/Builtins.def

@@ -629,6 +629,10 @@ BUILTIN_MISC_OPERATION(DeallocRaw, "deallocRaw", "", Special)
 /// is not known at compile time, MaximumAlignment is assumed.
 BUILTIN_MISC_OPERATION(StackAlloc, "stackAlloc", "", Special)
 
+/// Like `stackAlloc`, but doesn't set the `[stack_protection]` flag on its
+/// containing function.
+BUILTIN_MISC_OPERATION(UnprotectedStackAlloc, "unprotectedStackAlloc", "", Special)
+
 /// StackDealloc has type (Builtin.RawPointer) -> ()
 ///
 /// Parameters: address.

include/swift/Basic/Features.def

@@ -82,6 +82,7 @@ LANGUAGE_FEATURE(BuiltinBuildMainExecutor, 0, "MainActor executor building built
 LANGUAGE_FEATURE(BuiltinCreateAsyncTaskInGroup, 0, "MainActor executor building builtin", true)
 LANGUAGE_FEATURE(BuiltinCopy, 0, "Builtin.copy()", true)
 LANGUAGE_FEATURE(BuiltinStackAlloc, 0, "Builtin.stackAlloc", true)
+LANGUAGE_FEATURE(BuiltinUnprotectedStackAlloc, 0, "Builtin.unprotectedStackAlloc", true)
 LANGUAGE_FEATURE(BuiltinTaskRunInline, 0, "Builtin.taskRunInline", true)
 LANGUAGE_FEATURE(BuiltinUnprotectedAddressOf, 0, "Builtin.unprotectedAddressOf", true)
 SUPPRESSIBLE_LANGUAGE_FEATURE(SpecializeAttributeWithAvailability, 0, "@_specialize attribute with availability", true)

lib/AST/ASTPrinter.cpp

@@ -3020,6 +3020,10 @@ static bool usesFeatureBuiltinStackAlloc(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureBuiltinUnprotectedStackAlloc(Decl *decl) {
+  return false;
+}
+
 static bool usesFeatureBuiltinAssumeAlignment(Decl *decl) {
   return false;
 }

lib/AST/Builtins.cpp

@@ -2707,6 +2707,7 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     return getDeallocOperation(Context, Id);
 
   case BuiltinValueKind::StackAlloc:
+  case BuiltinValueKind::UnprotectedStackAlloc:
     return getStackAllocOperation(Context, Id);
   case BuiltinValueKind::StackDealloc:
     return getStackDeallocOperation(Context, Id);

lib/IRGen/IRGenSIL.cpp

@@ -3237,7 +3237,8 @@ static Alignment getStackAllocationAlignment(IRGenSILFunction &IGF,
 /// some other builtin.)
 static bool emitStackAllocBuiltinCall(IRGenSILFunction &IGF,
                                       swift::BuiltinInst *i) {
-  if (i->getBuiltinKind() == BuiltinValueKind::StackAlloc) {
+  if (i->getBuiltinKind() == BuiltinValueKind::StackAlloc ||
+      i->getBuiltinKind() == BuiltinValueKind::UnprotectedStackAlloc) {
     // Stack-allocate a buffer with the specified size/alignment.
     auto loc = i->getLoc().getSourceLoc();
     auto size = getStackAllocationSize(

lib/SIL/IR/OperandOwnership.cpp

@@ -783,6 +783,7 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SRem)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, GenericSRem)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SSubOver)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, StackAlloc)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, UnprotectedStackAlloc)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, StackDealloc)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SToSCheckedTrunc)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SToUCheckedTrunc)

lib/SIL/IR/SILInstruction.cpp

@@ -1243,7 +1243,8 @@ bool SILInstruction::isAllocatingStack() const {
     return PA->isOnStack();
 
   if (auto *BI = dyn_cast<BuiltinInst>(this)) {
-    if (BI->getBuiltinKind() == BuiltinValueKind::StackAlloc) {
+    if (BI->getBuiltinKind() == BuiltinValueKind::StackAlloc ||
+        BI->getBuiltinKind() == BuiltinValueKind::UnprotectedStackAlloc) {
       return true;
     }
   }

lib/SIL/IR/ValueOwnership.cpp

@@ -498,6 +498,7 @@ CONSTANT_OWNERSHIP_BUILTIN(None, AllocRaw)
 CONSTANT_OWNERSHIP_BUILTIN(None, AssertConf)
 CONSTANT_OWNERSHIP_BUILTIN(None, UToSCheckedTrunc)
 CONSTANT_OWNERSHIP_BUILTIN(None, StackAlloc)
+CONSTANT_OWNERSHIP_BUILTIN(None, UnprotectedStackAlloc)
 CONSTANT_OWNERSHIP_BUILTIN(None, StackDealloc)
 CONSTANT_OWNERSHIP_BUILTIN(None, SToSCheckedTrunc)
 CONSTANT_OWNERSHIP_BUILTIN(None, SToUCheckedTrunc)

lib/SIL/Utils/MemAccessUtils.cpp

@@ -2519,6 +2519,7 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
     case BuiltinValueKind::AllocRaw:
     case BuiltinValueKind::DeallocRaw:
     case BuiltinValueKind::StackAlloc:
+    case BuiltinValueKind::UnprotectedStackAlloc:
     case BuiltinValueKind::StackDealloc:
     case BuiltinValueKind::Fence:
     case BuiltinValueKind::StaticReport:

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -153,6 +153,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::CreateTaskGroupWithFlags:
     case BuiltinValueKind::DestroyTaskGroup:
     case BuiltinValueKind::StackAlloc:
+    case BuiltinValueKind::UnprotectedStackAlloc:
     case BuiltinValueKind::StackDealloc:
     case BuiltinValueKind::AssumeAlignment:
       return false;

stdlib/public/core/TemporaryAllocation.swift

@@ -129,16 +129,7 @@ internal func _withUnsafeTemporaryAllocation<T, R>(
   let byteCount = _byteCountForTemporaryAllocation(of: type, capacity: capacity)
 
   guard _isStackAllocationSafe(byteCount: byteCount, alignment: alignment) else {
-    // Fall back to the heap. This may still be optimizable if escape analysis
-    // shows that the allocated pointer does not escape.
-    let buffer = UnsafeMutableRawPointer.allocate(
-      byteCount: byteCount,
-      alignment: alignment
-    )
-    defer {
-      buffer.deallocate()
-    }
-    return try body(buffer._rawValue)
+    return try _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
   }
 
   // This declaration must come BEFORE Builtin.stackAlloc() or
@@ -170,6 +161,63 @@ internal func _withUnsafeTemporaryAllocation<T, R>(
 #endif
 }
 
+#if $BuiltinUnprotectedStackAlloc
+@_alwaysEmitIntoClient @_transparent
+internal func _withUnprotectedUnsafeTemporaryAllocation<T, R>(
+  of type: T.Type,
+  capacity: Int,
+  alignment: Int,
+  _ body: (Builtin.RawPointer) throws -> R
+) rethrows -> R {
+  // How many bytes do we need to allocate?
+  let byteCount = _byteCountForTemporaryAllocation(of: type, capacity: capacity)
+
+  guard _isStackAllocationSafe(byteCount: byteCount, alignment: alignment) else {
+    return try _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
+  }
+
+  // This declaration must come BEFORE Builtin.unprotectedStackAlloc() or
+  // Builtin.stackDealloc() will end up blowing it away (and the verifier will
+  // notice and complain.)
+  let result: R
+
+  let stackAddress = Builtin.unprotectedStackAlloc(
+    capacity._builtinWordValue,
+    MemoryLayout<T>.stride._builtinWordValue,
+    alignment._builtinWordValue
+  )
+
+  // The multiple calls to Builtin.stackDealloc() are because defer { } produces
+  // a child function at the SIL layer and that conflicts with the verifier's
+  // idea of a stack allocation's lifetime.
+  do {
+    result = try body(stackAddress)
+    Builtin.stackDealloc(stackAddress)
+    return result
+
+  } catch {
+    Builtin.stackDealloc(stackAddress)
+    throw error
+  }
+}
+#endif
+
+@_alwaysEmitIntoClient @_transparent
+internal func _fallBackToHeapAllocation<R>(
+  byteCount: Int,
+  alignment: Int,
+  _ body: (Builtin.RawPointer) throws -> R
+) rethrows -> R {
+  let buffer = UnsafeMutableRawPointer.allocate(
+    byteCount: byteCount,
+    alignment: alignment
+  )
+  defer {
+    buffer.deallocate()
+  }
+  return try body(buffer._rawValue)
+}
+
 // MARK: - Public interface
 
 /// Provides scoped access to a raw buffer pointer with the specified byte count
@@ -222,6 +270,34 @@ public func withUnsafeTemporaryAllocation<R>(
   }
 }
 
+/// Provides scoped access to a raw buffer pointer with the specified byte count
+/// and alignment.
+///
+/// This function is similar to `withUnsafeTemporaryAllocation`, except that it
+/// doesn't trigger stack protection for the stack allocated memory.
+@_alwaysEmitIntoClient @_transparent
+public func _withUnprotectedUnsafeTemporaryAllocation<R>(
+  byteCount: Int,
+  alignment: Int,
+  _ body: (UnsafeMutableRawBufferPointer) throws -> R
+) rethrows -> R {
+  return try _withUnsafeTemporaryAllocation(
+    of: Int8.self,
+    capacity: byteCount,
+    alignment: alignment
+  ) { pointer in
+#if $BuiltinUnprotectedStackAlloc
+    let buffer = UnsafeMutableRawBufferPointer(
+      start: .init(pointer),
+      count: byteCount
+    )
+    return try body(buffer)
+#else
+    return try withUnsafeTemporaryAllocation(byteCount: byteCount, alignment: alignment, body)
+#endif
+  }
+}
+
 /// Provides scoped access to a buffer pointer to memory of the specified type
 /// and with the specified capacity.
 ///
@@ -272,3 +348,32 @@ public func withUnsafeTemporaryAllocation<T, R>(
     return try body(buffer)
   }
 }
+
+/// Provides scoped access to a buffer pointer to memory of the specified type
+/// and with the specified capacity.
+///
+/// This function is similar to `withUnsafeTemporaryAllocation`, except that it
+/// doesn't trigger stack protection for the stack allocated memory.
+@_alwaysEmitIntoClient @_transparent
+public func _withUnprotectedUnsafeTemporaryAllocation<T, R>(
+  of type: T.Type,
+  capacity: Int,
+  _ body: (UnsafeMutableBufferPointer<T>) throws -> R
+) rethrows -> R {
+  return try _withUnprotectedUnsafeTemporaryAllocation(
+    of: type,
+    capacity: capacity,
+    alignment: MemoryLayout<T>.alignment
+  ) { pointer in
+#if $BuiltinUnprotectedStackAlloc
+    Builtin.bindMemory(pointer, capacity._builtinWordValue, type)
+    let buffer = UnsafeMutableBufferPointer<T>(
+      start: .init(pointer),
+      count: capacity
+    )
+    return try body(buffer)
+#else
+    return try withUnsafeTemporaryAllocation(of: type, capacity: capacity, body)
+#endif
+  }
+}

stdlib/public/core/UnsafeRawPointer.swift

@@ -463,7 +463,7 @@ public struct UnsafeRawPointer: _Pointer {
     as type: T.Type
   ) -> T {
     _debugPrecondition(_isPOD(T.self))
-    return withUnsafeTemporaryAllocation(of: T.self, capacity: 1) {
+    return _withUnprotectedUnsafeTemporaryAllocation(of: T.self, capacity: 1) {
       let temporary = $0.baseAddress._unsafelyUnwrappedUnchecked
       Builtin.int_memcpy_RawPointer_RawPointer_Int64(
         temporary._rawValue,
@@ -1245,7 +1245,7 @@ public struct UnsafeMutableRawPointer: _Pointer {
     as type: T.Type
   ) -> T {
     _debugPrecondition(_isPOD(T.self))
-    return withUnsafeTemporaryAllocation(of: T.self, capacity: 1) {
+    return _withUnprotectedUnsafeTemporaryAllocation(of: T.self, capacity: 1) {
       let temporary = $0.baseAddress._unsafelyUnwrappedUnchecked
       Builtin.int_memcpy_RawPointer_RawPointer_Int64(
         temporary._rawValue,

test/IRGen/temporary_allocation/codegen.swift

@@ -59,6 +59,13 @@ withUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
 // CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint [16 x i8]* [[INT_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[INT_PTR]])
 
+_withUnprotectedUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
+  blackHole(buffer.baseAddress)
+}
+// CHECK: [[INT_PTR_RAW:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
+// CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint [16 x i8]* [[INT_PTR_RAW]] to [[WORD]]
+// CHECK: call swiftcc void @blackHole([[WORD]] [[INT_PTR]])
+
 withUnsafeTemporaryAllocation(of: Void.self, capacity: 2) { buffer in
   blackHole(buffer.baseAddress)
 }