WIP for [non_nested] support for alloc_stack

include/swift/Runtime/RuntimeFunctions.def

@@ -2570,7 +2570,7 @@ FUNCTION(AsyncLetBegin,
          ),
          ATTRS(NoUnwind),
          EFFECT(RuntimeEffect::Concurrency),
-         MEMEFFECTS(ArgMemOnly))
+         UNKNOWN_MEMEFFECTS)
 
 /// void swift_asyncLet_finish(
 ///     AsyncLet *alet,
@@ -2592,7 +2592,7 @@ FUNCTION(AsyncLetFinish,
          ),
          ATTRS(NoUnwind),
          EFFECT(RuntimeEffect::Concurrency),
-         MEMEFFECTS(ArgMemOnly))
+         UNKNOWN_MEMEFFECTS)
 
 /// void swift_task_run_inline(
 ///     OpaqueValue *result,

include/swift/SIL/SILCloner.h

@@ -1062,6 +1062,7 @@ SILCloner<ImplClass>::visitAllocStackInst(AllocStackInst *Inst) {
       true
 #endif
   );
+  NewInst->setStackAllocationIsNested(Inst->isStackAllocationNested());
   recordClonedInstruction(Inst, NewInst);
 }
 

include/swift/SIL/SILInstruction.h

@@ -141,6 +141,15 @@ class SILPrintContext;
 
 template <typename ImplClass> class SILClonerWithScopes;
 
+enum StackAllocationIsNested_t : bool {
+  /// The instruction may not obey the ABBA rule of stack allocation.
+  StackAllocationIsNotNested = false,
+
+  /// The instruction obeys the ABBA rule of stack allocation and can
+  /// allocate memory on the stack normally.
+  StackAllocationIsNested = true,
+};
+
 enum class MemoryBehavior {
   None,
   /// The instruction may read memory.
@@ -858,6 +867,15 @@ class SILInstruction : public llvm::ilist_node<SILInstruction> {
   /// The stack allocation produced by the instruction, if any.
   SILValue getStackAllocation() const;
 
+  /// Returns the kind of stack memory that should be allocated. There are
+  /// certain (unfortunate) situations in which "stack" allocations may become
+  /// unnested and must use alternative allocation strategies. Rather than
+  /// requiring all of these to explicitly use heap allocation, which may be
+  /// to be significantly less efficient (e.g. )
+  StackAllocationIsNested_t isStackAllocationNested() const;
+
+  void setStackAllocationIsNested(StackAllocationIsNested_t isNested);
+
   /// Returns true if this is the deallocation of a stack allocating instruction.
   /// The first operand must be the allocating instruction.
   bool isDeallocatingStack() const;
@@ -2055,6 +2073,14 @@ class AllocStackInst final
     }
   }
 
+  StackAllocationIsNested_t isStackAllocationNested() const {
+    return StackAllocationIsNested_t(!sharedUInt8().AllocStackInst.isNested);
+  }
+
+  void setStackAllocationIsNested(StackAllocationIsNested_t isNested) {
+    sharedUInt8().AllocStackInst.isNested = bool(isNested);
+  }
+
   void markUsesMoveableValueDebugInfo() {
     sharedUInt8().AllocStackInst.usesMoveableValueDebugInfo =
         (bool)UsesMoveableValueDebugInfo;

include/swift/SIL/SILNode.h

@@ -229,7 +229,8 @@ class alignas(8) SILNode :
                  lexical : 1,
                  fromVarDecl : 1,
                  usesMoveableValueDebugInfo : 1,
-                 hasInvalidatedVarInfo : 1);
+                 hasInvalidatedVarInfo : 1,
+                 isNested : 1);
 
     SHARED_FIELD(AllocBoxInst, uint8_t
                  dynamicLifetime : 1,

lib/IRGen/FixedTypeInfo.h

@@ -82,8 +82,12 @@ class FixedTypeInfo : public TypeInfo {
   }
 
   StackAddress allocateStack(IRGenFunction &IGF, SILType T,
-                             const llvm::Twine &name) const override;
-  void deallocateStack(IRGenFunction &IGF, StackAddress addr, SILType T) const override;
+                             const llvm::Twine &name,
+                             StackAllocationIsNested_t isNested =
+                               StackAllocationIsNested) const override;
+  void deallocateStack(IRGenFunction &IGF, StackAddress addr, SILType T,
+                       StackAllocationIsNested_t isNested =
+                         StackAllocationIsNested) const override;
   void destroyStack(IRGenFunction &IGF, StackAddress addr, SILType T,
                     bool isOutlined) const override;
 

lib/IRGen/GenArray.cpp

@@ -639,17 +639,18 @@ class NonFixedArrayTypeInfo final
   }
 
   StackAddress allocateStack(IRGenFunction &IGF, SILType T,
-                             const llvm::Twine &name) const override {
+                             const llvm::Twine &name,
+                             StackAllocationIsNested_t isNested) const override {
     // Allocate memory on the stack.
-    auto alloca = IGF.emitDynamicAlloca(T, name);
+    auto alloca = IGF.emitDeallocStackDynamic(T, name, isNested);
     IGF.Builder.CreateLifetimeStart(alloca.getAddressPointer());
     return alloca.withAddress(getAddressForPointer(alloca.getAddressPointer()));
   }
 
   void deallocateStack(IRGenFunction &IGF, StackAddress stackAddress,
-                       SILType T) const override {
+                       SILType T, StackAllocationIsNested_t isNested) const override {
     IGF.Builder.CreateLifetimeEnd(stackAddress.getAddress().getAddress());
-    IGF.emitDeallocateDynamicAlloca(stackAddress);
+    IGF.emitDeallocStackDynamic(stackAddress, isNested);
   }
 
   void destroyStack(IRGenFunction &IGF, StackAddress stackAddress, SILType T,

lib/IRGen/GenInit.cpp

@@ -60,7 +60,8 @@ void IRGenModule::emitSILGlobalVariable(SILGlobalVariable *var) {
 }
 
 StackAddress FixedTypeInfo::allocateStack(IRGenFunction &IGF, SILType T,
-                                          const Twine &name) const {
+                                          const Twine &name,
+                                          StackAllocationIsNested_t isNested) const {
   // If the type is known to be empty, don't actually allocate anything.
   if (isKnownEmpty(ResilienceExpansion::Maximal)) {
     auto addr = getUndefAddress();
@@ -81,7 +82,8 @@ void FixedTypeInfo::destroyStack(IRGenFunction &IGF, StackAddress addr,
 }
 
 void FixedTypeInfo::deallocateStack(IRGenFunction &IGF, StackAddress addr,
-                                    SILType T) const {
+                                    SILType T,
+                                    StackAllocationIsNested_t isNested) const {
   if (isKnownEmpty(ResilienceExpansion::Maximal))
     return;
   IGF.Builder.CreateLifetimeEnd(addr.getAddress(), getFixedSize());

lib/IRGen/GenOpaque.cpp

@@ -553,6 +553,27 @@ irgen::emitInitializeBufferWithCopyOfBufferCall(IRGenFunction &IGF,
   return call;
 }
 
+StackAddress
+IRGenFunction::emitDynamicStackAllocation(SILType T,
+                                          StackAllocationIsNested_t isNested,
+                                          const llvm::Twine &name) {
+  if (isNested) {
+    return emitDynamicAlloca(T, name);
+  }
+
+  FIXME: use malloc
+}
+
+void
+IRGenFunction::emitDynamicStackDeallocation(StackAddress address,
+                                            StackAllocationIsNested_t isNested) {
+  if (isNested) {
+    return emitDeallocateDynamicAlloca(address);
+  }
+
+  FIXME: use malloc
+}
+
 /// Emit a dynamic alloca call to allocate enough memory to hold an object of
 /// type 'T' and an optional llvm.stackrestore point if 'isInEntryBlock' is
 /// false.

lib/IRGen/GenType.cpp

@@ -1374,11 +1374,13 @@ namespace {
       return nullptr;
     }
     StackAddress allocateStack(IRGenFunction &IGF, SILType T,
-                               const llvm::Twine &name) const override {
+                               const llvm::Twine &name,
+                               StackAllocationIsNested_t isNested) const override {
       llvm_unreachable("should not call on an immovable opaque type");
     }
     void deallocateStack(IRGenFunction &IGF, StackAddress addr,
-                         SILType T) const override {
+                         SILType T,
+                         StackAllocationIsNested_t isNested) const override {
       llvm_unreachable("should not call on an immovable opaque type");
     }
     void destroyStack(IRGenFunction &IGF, StackAddress addr, SILType T,

lib/IRGen/IRGenFunction.h

@@ -317,6 +317,12 @@ class IRGenFunction {
                                    bool useTaskDeallocThrough = false,
                                    bool forCalleeCoroutineFrame = false);
 
+  StackAddress emitAllocStackDynamic(SILType type,
+                                     StackAllocationIsNested_t isNested,
+                                     const llvm::Twine &name = "");
+  void emitDeallocStackDynamic(StackAddress address,
+                               StackAllocationIsNested_t isNested);
+
   llvm::BasicBlock *createBasicBlock(const llvm::Twine &Name);
   const TypeInfo &getTypeInfoForUnlowered(Type subst);
   const TypeInfo &getTypeInfoForUnlowered(AbstractionPattern orig, Type subst);

lib/IRGen/IRGenSIL.cpp

@@ -6575,7 +6575,9 @@ void IRGenSILFunction::visitAllocStackInst(swift::AllocStackInst *i) {
   DebugTypeInfo DbgTy;
   emitDebugInfoBeforeAllocStack(i, type, DbgTy);
 
-  auto stackAddr = type.allocateStack(*this, i->getElementType(), dbgname);
+  auto stackAddr = type.allocateStack(*this, i->getElementType(),
+                                      i->isStackAllocationNested(),
+                                      dbgname);
   setLoweredStackAddress(i, stackAddr);
   Address addr = stackAddr.getAddress();
 
@@ -6668,23 +6670,27 @@ void IRGenSILFunction::visitDeallocStackInst(swift::DeallocStackInst *i) {
   if (auto *closure = dyn_cast<PartialApplyInst>(i->getOperand())) {
     assert(closure->isOnStack());
     auto stackAddr = LoweredPartialApplyAllocations[i->getOperand()];
-    emitDeallocateDynamicAlloca(stackAddr);
+    emitDeallocStackDynamic(stackAddr, closure->isStackAllocationNested());
     return;
   }
   if (isaResultOf<BeginApplyInst>(i->getOperand())) {
     auto *mvi = getAsResultOf<BeginApplyInst>(i->getOperand());
     auto *bai = cast<BeginApplyInst>(mvi->getParent());
+    // FIXME: [non_nested]
     const auto &coroutine = getLoweredCoroutine(bai->getTokenResult());
     emitDeallocYieldOnce2CoroutineFrame(*this,
                                         coroutine.getCalleeAllocatedFrame());
     return;
   }
 
-  auto allocatedType = i->getOperand()->getType();
+  auto *asi = cast<AllocStackInst>(i->getOperand());
+
+  auto allocatedType = asi->getType();
   const TypeInfo &allocatedTI = getTypeInfo(allocatedType);
-  StackAddress stackAddr = getLoweredStackAddress(i->getOperand());
+  StackAddress stackAddr = getLoweredStackAddress(asi);
+  auto isNested = asi->isStackAllocationNested();
 
-  allocatedTI.deallocateStack(*this, stackAddr, allocatedType);
+  allocatedTI.deallocateStack(*this, stackAddr, allocatedType, isNested);
 }
 
 void IRGenSILFunction::visitDeallocStackRefInst(DeallocStackRefInst *i) {

lib/IRGen/NonFixedTypeInfo.h

@@ -63,18 +63,21 @@ class WitnessSizedTypeInfo : public IndirectTypeInfo<Impl, TypeInfo> {
   static bool isFixed() { return false; }
 
   StackAddress allocateStack(IRGenFunction &IGF, SILType T,
-                             const llvm::Twine &name) const override {
+                             const llvm::Twine &name,
+                             StackAllocationIsNested_t isNested =
+                               StackAllocationIsNested) const override {
     // Allocate memory on the stack.
-    auto alloca = IGF.emitDynamicAlloca(T, name);
+    auto alloca = IGF.emitDynamicStackAllocation(T, name, isNested);
     IGF.Builder.CreateLifetimeStart(alloca.getAddressPointer());
     return alloca.withAddress(
              getAsBitCastAddress(IGF, alloca.getAddressPointer()));
   }
 
   void deallocateStack(IRGenFunction &IGF, StackAddress stackAddress,
-                       SILType T) const override {
+                       SILType T, StackAllocationIsNested_t isNested =
+                         StackAllocationIsNested) const override {
     IGF.Builder.CreateLifetimeEnd(stackAddress.getAddress().getAddress());
-    IGF.emitDeallocateDynamicAlloca(stackAddress);
+    IGF.emitDynamicStackDeallocation(stackAddress, isNested);
   }
 
   void destroyStack(IRGenFunction &IGF, StackAddress stackAddress, SILType T,

lib/IRGen/TypeInfo.h

@@ -355,11 +355,15 @@ class TypeInfo {
 
   /// Allocate a variable of this type on the stack.
   virtual StackAddress allocateStack(IRGenFunction &IGF, SILType T,
-                                     const llvm::Twine &name) const = 0;
+                                     const llvm::Twine &name,
+                                     StackAllocationIsNested_t isNested =
+                                       StackAllocationIsNested) const = 0;
 
   /// Deallocate a variable of this type.
   virtual void deallocateStack(IRGenFunction &IGF, StackAddress addr,
-                               SILType T) const = 0;
+                               SILType T,
+                               StackAllocationIsNested_t isNested =
+                                 StackAllocationIsNested) const = 0;
 
   /// Destroy the value of a variable of this type, then deallocate its
   /// memory.

lib/SIL/IR/SILInstruction.cpp

@@ -1339,6 +1339,25 @@ SILValue SILInstruction::getStackAllocation() const {
   return cast<SingleValueInstruction>(this);
 }
 
+StackAllocationIsNested_t SILInstruction::isStackAllocationNested() const {
+  assert(isAllocatingStack());
+  if (auto ASI = dyn_cast<AllocStackInst>(this)) {
+    return ASI->isStackAllocationNested();
+  } else {
+    // TODO: implement for all remaining allocations
+    return StackAllocationIsNested;    
+  }
+}
+
+void SILInstruction::setStackAllocationIsNested(StackAllocationIsNested_t nested) {
+  assert(isAllocatingStack());
+  if (auto ASI = dyn_cast<AllocStackInst>(this)) {
+    ASI->setStackAllocationIsNested(nested);
+  } else if (!nested) {
+    llvm_unreachable("unimplemented");
+  }
+}
+
 bool SILInstruction::isDeallocatingStack() const {
   // NOTE: If you're adding a new kind of deallocating instruction,
   // there are several places scattered around the SIL optimizer which

lib/SIL/IR/SILInstructions.cpp

@@ -231,6 +231,7 @@ AllocStackInst::AllocStackInst(
   sharedUInt8().AllocStackInst.fromVarDecl = (bool)isFromVarDecl;
   sharedUInt8().AllocStackInst.usesMoveableValueDebugInfo =
       (bool)usesMoveableValueDebugInfo || elementType.isMoveOnly();
+  sharedUInt8().AllocStackInst.isNested = true;
   sharedUInt32().AllocStackInst.numOperands = TypeDependentOperands.size();
 
   // VarInfo must be initialized after

lib/SIL/IR/SILPrinter.cpp

@@ -1568,7 +1568,17 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
       printDebugInfoExpression(Var->DIExpr);
   }
 
+  template <class T>
+  void printNonNested(T *inst) {
+    static_assert(&T::isStackAllocationNested !=
+                    &SILInstruction::isStackAllocationNested,
+                  "Type doesn't override isStackAllocationNested");
+    if (!inst->isStackAllocationNested())
+      *this << "[non_nested] ";
+  }
+
   void visitAllocStackInst(AllocStackInst *AVI) {
+    printNonNested(AVI);
     if (AVI->hasDynamicLifetime())
       *this << "[dynamic_lifetime] ";
     if (AVI->isLexical())

lib/SIL/Parser/ParseSIL.cpp

@@ -4994,6 +4994,7 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
     auto isFromVarDecl = IsNotFromVarDecl;
     UsesMoveableValueDebugInfo_t usesMoveableValueDebugInfo =
         DoesNotUseMoveableValueDebugInfo;
+    auto isNested = StackAllocationIsNested;
 
     StringRef attributeName;
     SourceLoc attributeLoc;
@@ -5006,6 +5007,8 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
         isFromVarDecl = IsFromVarDecl;
       else if (attributeName == "moveable_value_debuginfo")
         usesMoveableValueDebugInfo = UsesMoveableValueDebugInfo;
+      else if (attributeName == "non_nested")
+        isNested = StackAllocationIsNotNested;
       else {
         P.diagnose(attributeLoc, diag::sil_invalid_attribute_for_instruction,
                    attributeName, "alloc_stack");
@@ -5025,14 +5028,16 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
       usesMoveableValueDebugInfo = UsesMoveableValueDebugInfo;
 
     // It doesn't make sense to attach a debug var info if the name is empty
+    AllocStackInst *ASI;
     if (VarInfo.Name.size())
-      ResultVal = B.createAllocStack(InstLoc, Ty, VarInfo, hasDynamicLifetime,
-                                     isLexical, isFromVarDecl,
-                                     usesMoveableValueDebugInfo);
+      ASI = B.createAllocStack(InstLoc, Ty, VarInfo, hasDynamicLifetime,
+                               isLexical, isFromVarDecl,
+                               usesMoveableValueDebugInfo);
     else
-      ResultVal =
-          B.createAllocStack(InstLoc, Ty, {}, hasDynamicLifetime, isLexical,
-                             isFromVarDecl, usesMoveableValueDebugInfo);
+      ASI = B.createAllocStack(InstLoc, Ty, {}, hasDynamicLifetime, isLexical,
+                               isFromVarDecl, usesMoveableValueDebugInfo);
+    ASI->setStackAllocationIsNested(isNested);
+    ResultVal = ASI;
     break;
   }
   case SILInstructionKind::MetatypeInst: {