Merge pull request #37666 from grynspan/jgrynspan/temporary-buffers

[SE-0322] Temporary uninitialized buffers

Author: grynspan

include/swift/AST/Builtins.def

@@ -603,6 +603,29 @@ BUILTIN_MISC_OPERATION(AllocRaw, "allocRaw", "", Special)
 /// was allocated.
 BUILTIN_MISC_OPERATION(DeallocRaw, "deallocRaw", "", Special)
 
+/// StackAlloc has type (Int, Int, Int) -> Builtin.RawPointer
+///
+/// Parameters: capacity, stride, alignment
+///
+/// The resulting pointer comes from the stack (as in the non-standard C
+/// extension `alloca()`.) It is at least as aligned as specified and is valid
+/// until the end of the calling scope.
+///
+/// The count and stride are multiplied together to get the byte count to use
+/// for the allocation.
+///
+/// The passed alignment must be a positive power of two. If the alignment value
+/// is not known at compile time, MaximumAlignment is assumed.
+BUILTIN_MISC_OPERATION(StackAlloc, "stackAlloc", "", Special)
+
+/// StackDealloc has type (Builtin.RawPointer) -> ()
+///
+/// Parameters: address.
+///
+/// The range starting at `address`, previously allocated with
+/// Builtin.stackAlloc(), is deallocated from the stack.
+BUILTIN_MISC_OPERATION(StackDealloc, "stackDealloc", "", Special)
+
 /// Fence has type () -> ().
 BUILTIN_MISC_OPERATION(Fence, "fence", "", None)
 

include/swift/AST/DiagnosticsIRGen.def

@@ -50,5 +50,14 @@ ERROR(alignment_more_than_maximum,none,
       "@_alignment cannot increase alignment above maximum alignment of %0",
       (unsigned))
 
+ERROR(temporary_allocation_size_negative,none,
+      "allocation capacity must be greater than or equal to zero", ())
+ERROR(temporary_allocation_size_overflow,none,
+      "allocation byte count too large", ())
+ERROR(temporary_allocation_alignment_not_positive,none,
+      "alignment value must be greater than zero", ())
+ERROR(temporary_allocation_alignment_not_power_of_2,none,
+      "alignment value must be a power of two", ())
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

lib/AST/Builtins.cpp

@@ -1024,6 +1024,18 @@ static ValueDecl *getDeallocOperation(ASTContext &ctx, Identifier id) {
                             _void);
 }
 
+static ValueDecl *getStackAllocOperation(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin,
+                            _parameters(_word, _word, _word),
+                            _rawPointer);
+}
+
+static ValueDecl *getStackDeallocOperation(ASTContext &ctx, Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin,
+                            _parameters(_rawPointer),
+                            _void);
+}
+
 static ValueDecl *getFenceOperation(ASTContext &ctx, Identifier id) {
   return getBuiltinFunction(ctx, id, _thin, _parameters(), _void);
 }
@@ -2641,6 +2653,11 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
   case BuiltinValueKind::DeallocRaw:
     return getDeallocOperation(Context, Id);
 
+  case BuiltinValueKind::StackAlloc:
+    return getStackAllocOperation(Context, Id);
+  case BuiltinValueKind::StackDealloc:
+    return getStackDeallocOperation(Context, Id);
+
   case BuiltinValueKind::CastToNativeObject:
   case BuiltinValueKind::UnsafeCastToNativeObject:
   case BuiltinValueKind::CastFromNativeObject:

lib/IRGen/GenOpaque.cpp

@@ -533,15 +533,16 @@ irgen::emitInitializeBufferWithCopyOfBufferCall(IRGenFunction &IGF,
 StackAddress IRGenFunction::emitDynamicAlloca(SILType T,
                                               const llvm::Twine &name) {
   llvm::Value *size = emitLoadOfSize(*this, T);
-  return emitDynamicAlloca(IGM.Int8Ty, size, Alignment(16), name);
+  return emitDynamicAlloca(IGM.Int8Ty, size, Alignment(16), true, name);
 }
 
 StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
                                               llvm::Value *arraySize,
                                               Alignment align,
+                                              bool allowTaskAlloc,
                                               const llvm::Twine &name) {
   // Async functions call task alloc.
-  if (isAsync()) {
+  if (allowTaskAlloc && isAsync()) {
     llvm::Value *byteCount;
     auto eltSize = IGM.DataLayout.getTypeAllocSize(eltTy);
     if (eltSize == 1) {
@@ -556,6 +557,8 @@ StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
     return {address, address.getAddress()};
   // In coroutines, call llvm.coro.alloca.alloc.
   } else if (isCoroutine()) {
+    // NOTE: llvm does not support dynamic allocas in coroutines.
+
     // Compute the number of bytes to allocate.
     llvm::Value *byteCount;
     auto eltSize = IGM.DataLayout.getTypeAllocSize(eltTy);
@@ -606,9 +609,10 @@ StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
 
 /// Deallocate dynamic alloca's memory if requested by restoring the stack
 /// location before the dynamic alloca's call.
-void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address) {
+void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address,
+                                                bool allowTaskDealloc) {
   // Async function use taskDealloc.
-  if (isAsync() && address.getAddress().isValid()) {
+  if (allowTaskDealloc && isAsync() && address.getAddress().isValid()) {
     emitTaskDealloc(Address(address.getExtraInfo(), address.getAlignment()));
     return;
   }
@@ -617,6 +621,8 @@ void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address) {
   // for a partial_apply [stack] that did not need a context object on the
   // stack.
   else if (isCoroutine() && address.getAddress().isValid()) {
+    // NOTE: llvm does not support dynamic allocas in coroutines.
+
     auto allocToken = address.getExtraInfo();
     assert(allocToken && "dynamic alloca in coroutine without alloc token?");
     auto freeFn = llvm::Intrinsic::getDeclaration(

lib/IRGen/IRGenFunction.h

@@ -222,9 +222,10 @@ class IRGenFunction {
 
   StackAddress emitDynamicAlloca(SILType type, const llvm::Twine &name = "");
   StackAddress emitDynamicAlloca(llvm::Type *eltTy, llvm::Value *arraySize,
-                                 Alignment align,
+                                 Alignment align, bool allowTaskAlloc = true,
                                  const llvm::Twine &name = "");
-  void emitDeallocateDynamicAlloca(StackAddress address);
+  void emitDeallocateDynamicAlloca(StackAddress address,
+                                   bool allowTaskDealloc = true);
 
   llvm::BasicBlock *createBasicBlock(const llvm::Twine &Name);
   const TypeInfo &getTypeInfoForUnlowered(Type subst);

lib/IRGen/IRGenSIL.cpp

@@ -18,6 +18,7 @@
 #define DEBUG_TYPE "irgensil"
 
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/DiagnosticsIRGen.h"
 #include "swift/AST/IRGenOptions.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Pattern.h"
@@ -57,6 +58,7 @@
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Transforms/Utils/Local.h"
 
@@ -436,6 +438,11 @@ class IRGenSILFunction :
   /// Calculates EstimatedStackSize.
   void estimateStackSize();
 
+  inline bool isAddress(SILValue v) const {
+    SILType type = v->getType();
+    return type.isAddress() || type.getASTType() == IGM.Context.TheRawPointerType;
+  }
+
   void setLoweredValue(SILValue v, LoweredValue &&lv) {
     auto inserted = LoweredValues.insert({v, std::move(lv)});
     assert(inserted.second && "already had lowered value for sil value?!");
@@ -444,31 +451,31 @@ class IRGenSILFunction :
 
   /// Create a new Address corresponding to the given SIL address value.
   void setLoweredAddress(SILValue v, const Address &address) {
-    assert(v->getType().isAddress() && "address for non-address value?!");
+    assert(isAddress(v) && "address for non-address value?!");
     setLoweredValue(v, address);
   }
 
   void setLoweredStackAddress(SILValue v, const StackAddress &address) {
-    assert(v->getType().isAddress() && "address for non-address value?!");
+    assert(isAddress(v) && "address for non-address value?!");
     setLoweredValue(v, address);
   }
 
   void setLoweredDynamicallyEnforcedAddress(SILValue v,
                                             const Address &address,
                                             llvm::Value *scratch) {
-    assert(v->getType().isAddress() && "address for non-address value?!");
+    assert(isAddress(v) && "address for non-address value?!");
     setLoweredValue(v, DynamicallyEnforcedAddress{address, scratch});
   }
 
   void setContainerOfUnallocatedAddress(SILValue v,
                                         const Address &buffer) {
-    assert(v->getType().isAddress() && "address for non-address value?!");
+    assert(isAddress(v) && "address for non-address value?!");
     setLoweredValue(v,
       LoweredValue(buffer, LoweredValue::ContainerForUnallocatedAddress));
   }
 
   void overwriteAllocatedAddress(SILValue v, const Address &address) {
-    assert(v->getType().isAddress() && "address for non-address value?!");
+    assert(isAddress(v) && "address for non-address value?!");
     auto it = LoweredValues.find(v);
     assert(it != LoweredValues.end() && "no existing entry for overwrite?");
     assert(it->second.isUnallocatedAddressInBuffer() &&
@@ -1623,6 +1630,9 @@ void LoweredValue::getExplosion(IRGenFunction &IGF, SILType type,
                                 Explosion &ex) const {
   switch (kind) {
   case Kind::StackAddress:
+    ex.add(Storage.get<StackAddress>(kind).getAddressPointer());
+    return;
+
   case Kind::ContainedAddress:
   case Kind::DynamicallyEnforcedAddress:
   case Kind::CoroutineState:
@@ -2958,9 +2968,127 @@ static std::unique_ptr<CallEmission> getCallEmissionForLoweredValue(
   return callEmission;
 }
 
+/// Get the size passed to stackAlloc().
+static llvm::Value *getStackAllocationSize(IRGenSILFunction &IGF,
+                                           SILValue vCapacity,
+                                           SILValue vStride,
+                                           SourceLoc loc) {
+  auto &Diags = IGF.IGM.Context.Diags;
+
+  // Check for a negative capacity, which is invalid.
+  auto capacity = IGF.getLoweredSingletonExplosion(vCapacity);
+  Optional<int64_t> capacityValue;
+  if (auto capacityConst = dyn_cast<llvm::ConstantInt>(capacity)) {
+    capacityValue = capacityConst->getSExtValue();
+    if (*capacityValue < 0) {
+      Diags.diagnose(loc, diag::temporary_allocation_size_negative);
+    }
+  }
+
+  // Check for a negative stride, which should never occur because the caller
+  // should always be using MemoryLayout<T>.stride to produce this value.
+  auto stride = IGF.getLoweredSingletonExplosion(vStride);
+  Optional<int64_t> strideValue;
+  if (auto strideConst = dyn_cast<llvm::ConstantInt>(stride)) {
+    strideValue = strideConst->getSExtValue();
+    if (*strideValue < 0) {
+      llvm_unreachable("Builtin.stackAlloc() caller passed an invalid stride");
+    }
+  }
+
+  // Get the byte count (the product of capacity and stride.)
+  llvm::Value *result = nullptr;
+  if (capacityValue && strideValue) {
+    int64_t byteCount = 0;
+    auto overflow = llvm::MulOverflow(*capacityValue, *strideValue, byteCount);
+    if (overflow) {
+      Diags.diagnose(loc, diag::temporary_allocation_size_overflow);
+    }
+    result = llvm::ConstantInt::get(IGF.IGM.SizeTy, byteCount);
+
+  } else {
+    // If either value is not known at compile-time, preconditions must be
+    // tested at runtime by Builtin.stackAlloc()'s caller. See
+    // _byteCountForTemporaryAllocation(of:capacity:).
+    result = IGF.Builder.CreateMul(capacity, stride);
+  }
+
+  // If the caller requests a zero-byte allocation, allocate one byte instead
+  // to ensure that the resulting pointer is valid and unique on the stack.
+  return IGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::umax,
+    {IGF.IGM.SizeTy}, {llvm::ConstantInt::get(IGF.IGM.SizeTy, 1), result});
+}
+
+/// Get the alignment passed to stackAlloc() as a compile-time constant.
+///
+/// If the specified alignment is not known at compile time or is not valid,
+/// the default maximum alignment is substituted.
+static Alignment getStackAllocationAlignment(IRGenSILFunction &IGF,
+                                             SILValue v,
+                                             SourceLoc loc) {
+  auto &Diags = IGF.IGM.Context.Diags;
+
+  // Check for a non-positive alignment, which is invalid.
+  auto align = IGF.getLoweredSingletonExplosion(v);
+  if (auto alignConst = dyn_cast<llvm::ConstantInt>(align)) {
+    auto alignValue = alignConst->getSExtValue();
+    if (alignValue <= 0) {
+      Diags.diagnose(loc, diag::temporary_allocation_alignment_not_positive);
+    } else if (!llvm::isPowerOf2_64(alignValue)) {
+      Diags.diagnose(loc, diag::temporary_allocation_alignment_not_power_of_2);
+    } else {
+      return Alignment(alignValue);
+    }
+  }
+
+  // If the alignment is not known at compile-time, preconditions must be tested
+  // at runtime by Builtin.stackAlloc()'s caller. See
+  // _isStackAllocationSafe(byteCount:alignment:).
+  return Alignment(MaximumAlignment);
+}
+
+/// Emit a call to a stack allocation builtin (stackAlloc() or stackDealloc().)
+///
+/// Returns whether or not `i` was such a builtin (true if so, false if it was
+/// some other builtin.)
+static bool emitStackAllocBuiltinCall(IRGenSILFunction &IGF,
+                                      swift::BuiltinInst *i) {
+  if (i->getBuiltinKind() == BuiltinValueKind::StackAlloc) {
+    // Stack-allocate a buffer with the specified size/alignment.
+    auto loc = i->getLoc().getSourceLoc();
+    auto size = getStackAllocationSize(
+      IGF, i->getOperand(0), i->getOperand(1), loc);
+    auto align = getStackAllocationAlignment(IGF, i->getOperand(2), loc);
+
+    auto stackAddress = IGF.emitDynamicAlloca(IGF.IGM.Int8Ty, size, align,
+                                              false, "temp_alloc");
+    IGF.setLoweredStackAddress(i, stackAddress);
+
+    return true;
+
+  } else if (i->getBuiltinKind() == BuiltinValueKind::StackDealloc) {
+    // Deallocate a stack address previously allocated with the StackAlloc
+    // builtin above.
+    auto address = i->getOperand(0);
+    auto stackAddress = IGF.getLoweredStackAddress(address);
+
+    if (stackAddress.getAddress().isValid()) {
+      IGF.emitDeallocateDynamicAlloca(stackAddress, false);
+    }
+
+    return true;
+  }
+
+  return false;
+}
+
 void IRGenSILFunction::visitBuiltinInst(swift::BuiltinInst *i) {
   const BuiltinInfo &builtin = getSILModule().getBuiltinInfo(i->getName());
 
+  if (emitStackAllocBuiltinCall(*this, i)) {
+    return;
+  }
+
   auto argValues = i->getArguments();
   Explosion args;
   SmallVector<SILType, 4> argTypes;

lib/SIL/IR/OperandOwnership.cpp

@@ -733,6 +733,8 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SMulOver)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SRem)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, GenericSRem)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SSubOver)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, StackAlloc)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, StackDealloc)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SToSCheckedTrunc)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, SToUCheckedTrunc)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, Expect)

lib/SIL/IR/SILInstruction.cpp

@@ -1264,6 +1264,12 @@ bool SILInstruction::isAllocatingStack() const {
   if (auto *PA = dyn_cast<PartialApplyInst>(this))
     return PA->isOnStack();
 
+  if (auto *BI = dyn_cast<BuiltinInst>(this)) {
+    if (BI->getBuiltinKind() == BuiltinValueKind::StackAlloc) {
+      return true;
+    }
+  }
+
   return false;
 }
 
@@ -1275,6 +1281,13 @@ bool SILInstruction::isDeallocatingStack() const {
     if (DRI->canAllocOnStack())
       return true;
   }
+
+  if (auto *BI = dyn_cast<BuiltinInst>(this)) {
+    if (BI->getBuiltinKind() == BuiltinValueKind::StackDealloc) {
+      return true;
+    }
+  }
+
   return false;
 }
 

lib/SIL/IR/ValueOwnership.cpp

@@ -487,6 +487,8 @@ CONSTANT_OWNERSHIP_BUILTIN(None, Alignof)
 CONSTANT_OWNERSHIP_BUILTIN(None, AllocRaw)
 CONSTANT_OWNERSHIP_BUILTIN(None, AssertConf)
 CONSTANT_OWNERSHIP_BUILTIN(None, UToSCheckedTrunc)
+CONSTANT_OWNERSHIP_BUILTIN(None, StackAlloc)
+CONSTANT_OWNERSHIP_BUILTIN(None, StackDealloc)
 CONSTANT_OWNERSHIP_BUILTIN(None, SToSCheckedTrunc)
 CONSTANT_OWNERSHIP_BUILTIN(None, SToUCheckedTrunc)
 CONSTANT_OWNERSHIP_BUILTIN(None, UToUCheckedTrunc)

lib/SIL/Utils/MemAccessUtils.cpp

@@ -2154,6 +2154,8 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
     case BuiltinValueKind::CondFailMessage:
     case BuiltinValueKind::AllocRaw:
     case BuiltinValueKind::DeallocRaw:
+    case BuiltinValueKind::StackAlloc:
+    case BuiltinValueKind::StackDealloc:
     case BuiltinValueKind::Fence:
     case BuiltinValueKind::StaticReport:
     case BuiltinValueKind::Once: