Try an experiment

include/swift/AST/Builtins.def

@@ -880,6 +880,11 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildDefaultActorExecutorRef,
 BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildMainActorExecutorRef,
                                    "buildMainActorExecutorRef", "n", Special)
 
+/// isOnStack : (Builtin.NativeObject) -> Builtin.Int1
+///
+/// Returns true if the optimizer successfully moved this object from the stack onto the heap.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(IsOnStack, "isOnStack", "n", Special)
+
 #undef BUILTIN_MISC_OPERATION_WITH_SILGEN
 
 #undef BUILTIN_MISC_OPERATION

include/swift/IRGen/IRGenSILPasses.h

@@ -16,9 +16,9 @@ class SILTransform;
 
 namespace irgen {
 
-/// Create a pass to hoist alloc_stack instructions with non-fixed size.
 SILTransform *createAllocStackHoisting();
 SILTransform *createLoadableByAddress();
+SILTransform *createStackSizeLimitTransform();
 
 } // end namespace irgen
 } // end namespace swift

include/swift/SIL/SILBuilder.h

@@ -414,6 +414,14 @@ class SILBuilder {
                                        Var, hasDynamicLifetime));
   }
 
+  /// Helper method for creating alloc_box from an instance value.
+  AllocBoxInst *createAllocBox(SILLocation Loc, SILType InstanceType,
+                               Optional<SILDebugVariable> Var = None,
+                               bool hasDynamicLifetime = false) {
+    auto BoxType = SILBoxType::get(InstanceType.getASTType());
+    return createAllocBox(Loc, BoxType, Var, hasDynamicLifetime);
+  }
+
   AllocExistentialBoxInst *
   createAllocExistentialBox(SILLocation Loc, SILType ExistentialType,
                             CanType ConcreteType,
@@ -610,6 +618,15 @@ class SILBuilder {
                                    getModule()));
   }
 
+  /// Returns integer_literal Builtin.Int1, 1
+  IntegerLiteralInst *createIntegerLiteralBool(SILLocation Loc, bool Value) {
+    auto builtinType =
+        SILType::getBuiltinIntegerType(1, getModule().getASTContext());
+
+    return insert(IntegerLiteralInst::create(getSILDebugLocation(Loc),
+                                             builtinType, Value, getModule()));
+  }
+
   FloatLiteralInst *createFloatLiteral(FloatLiteralExpr *E);
 
   FloatLiteralInst *createFloatLiteral(SILLocation Loc, SILType Ty,

include/swift/SILOptimizer/PassManager/Passes.def

@@ -342,6 +342,8 @@ PASS(LateReleaseHoisting, "late-release-hoisting",
      "Late SIL release Hoisting Preserving Epilogues")
 IRGEN_PASS(LoadableByAddress, "loadable-address",
      "SIL Large Loadable type by-address lowering.")
+IRGEN_PASS(StackSizeLimitTransform, "stack-size-limit",
+           "Limit the amount of stack growth possible in a single SIL function by converting stack memory to heap memory")
 PASS(MandatorySILLinker, "mandatory-linker",
      "Deserialize all referenced SIL functions that are shared or transparent")
 PASS(PerformanceSILLinker, "performance-linker",

lib/AST/Builtins.cpp

@@ -963,6 +963,11 @@ static ValueDecl *getEndUnpairedAccessOperation(ASTContext &ctx,
                             _void);
 }
 
+static ValueDecl *getIsOnStackOperation(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin, _parameters(_nativeObject),
+                            _int(1));
+}
+
 static ValueDecl *getSizeOrAlignOfOperation(ASTContext &ctx,
                                             Identifier id) {
   return getBuiltinFunction(ctx, id, _thin,
@@ -2492,6 +2497,11 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     if (!Types.empty()) return nullptr;
     return getEndUnpairedAccessOperation(Context, Id);
 
+  case BuiltinValueKind::IsOnStack:
+    if (!Types.empty())
+      return nullptr;
+    return getIsOnStackOperation(Context, Id);
+
 #define BUILTIN(id, name, Attrs)
 #define BUILTIN_BINARY_OPERATION(id, name, attrs)
 #define BUILTIN_BINARY_OPERATION_OVERLOADED_STATIC(id, name, attrs, overload)  \

lib/IRGen/CMakeLists.txt

@@ -54,6 +54,7 @@ add_swift_host_library(swiftIRGen STATIC
   MetadataRequest.cpp
   Outlining.cpp
   StructLayout.cpp
+  StackSizeLimitTransform.cpp
   SwiftTargetInfo.cpp
   TypeLayout.cpp
   TypeLayoutDumper.cpp

lib/IRGen/GenBuiltin.cpp

@@ -124,6 +124,9 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
                             ArrayRef<SILType> argTypes,
                             Explosion &args, Explosion &out,
                             SubstitutionMap substitutions) {
+  assert(Builtin.ID != BuiltinValueKind::IsOnStack &&
+         "Builtin should not exist at IRGen time");
+
   if (Builtin.ID == BuiltinValueKind::COWBufferForReading) {
     // Just forward the incoming argument.
     assert(args.size() == 1 && "Expecting one incoming argument");

lib/SIL/IR/OperandOwnership.cpp

@@ -763,6 +763,7 @@ BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, StartAsyncLetWithLocalBuffer)
 BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, EndAsyncLetLifetime)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroup)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, DestroyTaskGroup)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, IsOnStack)
 
 BUILTIN_OPERAND_OWNERSHIP(ForwardingConsume, COWBufferForReading)
 BUILTIN_OPERAND_OWNERSHIP(ForwardingConsume, UnsafeGuaranteed)

lib/SIL/IR/ValueOwnership.cpp

@@ -555,6 +555,7 @@ CONSTANT_OWNERSHIP_BUILTIN(None, StartAsyncLetWithLocalBuffer)
 CONSTANT_OWNERSHIP_BUILTIN(None, EndAsyncLetLifetime)
 CONSTANT_OWNERSHIP_BUILTIN(None, CreateTaskGroup)
 CONSTANT_OWNERSHIP_BUILTIN(None, DestroyTaskGroup)
+CONSTANT_OWNERSHIP_BUILTIN(None, IsOnStack)
 
 #undef CONSTANT_OWNERSHIP_BUILTIN
 

lib/SILGen/SILGenBuiltin.cpp

@@ -1657,6 +1657,21 @@ static ManagedValue emitBuiltinBuildMainActorExecutorRef(
                               BuiltinValueKind::BuildMainActorExecutorRef);
 }
 
+static ManagedValue emitBuiltinIsOnStack(SILGenFunction &SGF, SILLocation loc,
+                                         SubstitutionMap substitutions,
+                                         ArrayRef<ManagedValue> args,
+                                         SGFContext C) {
+  auto &ctx = SGF.getASTContext();
+  auto id = ctx.getIdentifier(getBuiltinName(BuiltinValueKind::IsOnStack));
+  SmallVector<SILValue, 1> argValues;
+  if (!args.empty())
+    argValues.push_back(args[0].getValue());
+  auto builtinType = SILType::getBuiltinIntegerType(1, ctx);
+  auto builtin =
+      SGF.B.createBuiltin(loc, id, builtinType, substitutions, argValues);
+  return ManagedValue::forUnmanaged(builtin);
+}
+
 Optional<SpecializedEmitter>
 SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) {
   // Only consider standalone declarations in the Builtin module.

lib/SILOptimizer/Analysis/EscapeAnalysis.cpp

@@ -2373,6 +2373,18 @@ void EscapeAnalysis::analyzeInstruction(SILInstruction *I,
       }
       return;
     }
+    case SILInstructionKind::BuiltinInst: {
+      // Some builtins do not escape.
+      if (auto *bi = dyn_cast<BuiltinInst>(I))
+        if (auto kind = bi->getBuiltinKind())
+          // For now we only do is on stack.
+          if (*kind == BuiltinValueKind::IsOnStack)
+            return;
+
+      // If we don't have one of those, be conservative.
+      setAllEscaping(I, ConGraph);
+      return;
+    }
     default:
       // We handle all other instructions conservatively.
       setAllEscaping(I, ConGraph);

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -764,6 +764,10 @@ SILPassPipelinePlan::getLoweringPassPipeline(const SILOptions &Options) {
   return P;
 }
 
+static llvm::cl::opt<bool>
+    EnableStackSizeLimitTransform("swift-irgen-enable-stack-size-limit",
+                                  llvm::cl::init(true), llvm::cl::Hidden);
+
 SILPassPipelinePlan
 SILPassPipelinePlan::getIRGenPreparePassPipeline(const SILOptions &Options) {
   SILPassPipelinePlan P(Options);
@@ -773,6 +777,8 @@ SILPassPipelinePlan::getIRGenPreparePassPipeline(const SILOptions &Options) {
   // llvm-ir generation for dynamic alloca instructions.
   P.addAllocStackHoisting();
   P.addLoadableByAddress();
+  if (EnableStackSizeLimitTransform)
+    P.addStackSizeLimitTransform();
 
   return P;
 }

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -113,6 +113,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::Sizeof:
     case BuiltinValueKind::Strideof:
     case BuiltinValueKind::IsPOD:
+    case BuiltinValueKind::IsOnStack:
     case BuiltinValueKind::IsConcrete:
     case BuiltinValueKind::IsBitwiseTakable:
     case BuiltinValueKind::IsSameMetatype:

stdlib/public/core/ContiguousArrayBuffer.swift

@@ -276,8 +276,8 @@ internal struct _ContiguousArrayBuffer<Element>: _ArrayBufferProtocol {
          _ContiguousArrayStorage<Element>.self,
          realMinimumCapacity._builtinWordValue, Element.self)
 
-      let storageAddr = UnsafeMutableRawPointer(Builtin.bridgeToRawPointer(_storage))
-      if let allocSize = _mallocSize(ofAllocation: storageAddr) {
+      if let allocSize = _mallocSizeIfHeap(of: _storage) {
+        let storageAddr = UnsafeMutableRawPointer(Builtin.bridgeToRawPointer(_storage))
         let endAddr = storageAddr + allocSize
         let realCapacity = endAddr.assumingMemoryBound(to: Element.self) - firstElementAddress
         _initStorageHeader(

stdlib/public/core/Shims.swift

@@ -37,7 +37,19 @@ internal let _fastEnumerationStorageMutationsPtr =
   UnsafeMutablePointer<CUnsignedLong>(Builtin.addressof(&_fastEnumerationStorageMutationsTarget))
 #endif
 
-@usableFromInline @_alwaysEmitIntoClient
+
+
+
+@usableFromInline @_alwaysEmitIntoClient @_effects(readonly)
 internal func _mallocSize(ofAllocation ptr: UnsafeRawPointer) -> Int? {
   return _swift_stdlib_has_malloc_size() ? _swift_stdlib_malloc_size(ptr) : nil
 }
+
+@_transparent @inlinable @_effects(readonly)
+internal func _mallocSizeIfHeap(of object: AnyObject) -> Int? {
+  if Bool(Builtin.isOnStack(_nativeObject(fromNative: object))) { return nil }
+  return _mallocSize(
+    ofAllocation: UnsafeRawPointer(Builtin.bridgeToRawPointer(object))
+  )
+}
+