Reuse _ContiguousArrayStorage<AnyObject> metadata for any class or objc generic type

Reduces the number of _ContiguousArrayStorage metadata.

In order to support constant time bridging we do need to set the correct
metadata when we bridge to Objective-C. This is so that the type check
succeeds when bridging back from Objective-C to reuse the storage
instance rather than bridging the elements.

To support dynamically setting the `_ContiguousArrayStorage` element
type i needed to add support for optimizing `alloc_ref_dynamic`
throughout the optimizer.

Possible future improvements:
* Use different metadata such that we can disambiguate native Swift
  classes during destruction -- allowing native release rather then unknown
  release usage.
* Optimize the newly added semantic function
  getContiguousArrayStorageType

rdar://86171143

Author: aschwaighofer

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -220,7 +220,7 @@ final public class DeallocStackInst : Instruction, UnaryInstruction {
 }
 
 final public class DeallocStackRefInst : Instruction, UnaryInstruction {
-  public var allocRef: AllocRefInst { operand as! AllocRefInst }
+  public var allocRef: AllocRefInstBase { operand as! AllocRefInstBase }
 }
 
 final public class CondFailInst : Instruction, UnaryInstruction {

include/swift/AST/KnownStdlibTypes.def

@@ -49,6 +49,7 @@ KNOWN_STDLIB_TYPE_DECL(String, NominalTypeDecl, 0)
 KNOWN_STDLIB_TYPE_DECL(StaticString, NominalTypeDecl, 0)
 KNOWN_STDLIB_TYPE_DECL(Substring, NominalTypeDecl, 0)
 KNOWN_STDLIB_TYPE_DECL(Array, NominalTypeDecl, 1)
+KNOWN_STDLIB_TYPE_DECL(_ContiguousArrayStorage, NominalTypeDecl, 1)
 KNOWN_STDLIB_TYPE_DECL(Set, NominalTypeDecl, 1)
 KNOWN_STDLIB_TYPE_DECL(Sequence, NominalTypeDecl, 1)
 KNOWN_STDLIB_TYPE_DECL(Dictionary, NominalTypeDecl, 2)

include/swift/SIL/SILBuilder.h

@@ -408,14 +408,15 @@ class SILBuilder {
 
   AllocRefDynamicInst *createAllocRefDynamic(SILLocation Loc, SILValue operand,
                                              SILType type, bool objc,
+                                             bool canAllocOnStack,
                                     ArrayRef<SILType> ElementTypes,
                                     ArrayRef<SILValue> ElementCountOperands) {
     // AllocRefDynamicInsts expand to function calls and can therefore
     // not be counted towards the function prologue.
     assert(!Loc.isInPrologue());
     return insert(AllocRefDynamicInst::create(
-        getSILDebugLocation(Loc), *F, operand, type, objc, ElementTypes,
-        ElementCountOperands));
+        getSILDebugLocation(Loc), *F, operand, type, objc, canAllocOnStack,
+        ElementTypes, ElementCountOperands));
   }
 
   AllocBoxInst *createAllocBox(SILLocation Loc, CanSILBoxType BoxType,

include/swift/SIL/SILCloner.h

@@ -839,6 +839,7 @@ SILCloner<ImplClass>::visitAllocRefDynamicInst(AllocRefDynamicInst *Inst) {
                                       getOpValue(Inst->getMetatypeOperand()),
                                       getOpType(Inst->getType()),
                                       Inst->isObjC(),
+                                      Inst->canAllocOnStack(),
                                       ElemTypes, CountArgs);
   recordClonedInstruction(Inst, NewInst);
 }

include/swift/SIL/SILInstruction.h

@@ -2102,6 +2102,26 @@ class AllocRefInstBase : public AllocationInst {
   bool isObjC() const {
     return SILNode::Bits.AllocRefInstBase.ObjC;
   }
+
+  static bool classof(SILNodePointer node) {
+    if (auto *i = dyn_cast<SILInstruction>(node.get()))
+      return classof(i);
+    return false;
+  }
+
+  static bool classof(const SILInstruction *inst) {
+    return classof(inst->getKind());
+  }
+
+  static bool classof(SILInstructionKind kind) {
+    switch (kind) {
+    case SILInstructionKind::AllocRefInst:
+    case SILInstructionKind::AllocRefDynamicInst:
+      return true;
+    default:
+      return false;
+    }
+  }
 };
 
 /// AllocRefInst - This represents the primitive allocation of an instance
@@ -2162,10 +2182,11 @@ class AllocRefDynamicInst final
   AllocRefDynamicInst(SILDebugLocation DebugLoc,
                       SILType ty,
                       bool objc,
+                      bool canBeOnStack,
                       ArrayRef<SILType> ElementTypes,
                       ArrayRef<SILValue> AllOperands)
       : InstructionBaseWithTrailingOperands(AllOperands, DebugLoc, ty, objc,
-                                            false, ElementTypes) {
+                                            canBeOnStack, ElementTypes) {
     assert(AllOperands.size() >= ElementTypes.size() + 1);
     std::uninitialized_copy(ElementTypes.begin(), ElementTypes.end(),
                             getTrailingObjects<SILType>());
@@ -2174,6 +2195,7 @@ class AllocRefDynamicInst final
   static AllocRefDynamicInst *
   create(SILDebugLocation DebugLoc, SILFunction &F,
          SILValue metatypeOperand, SILType ty, bool objc,
+         bool canBeOnStack,
          ArrayRef<SILType> ElementTypes,
          ArrayRef<SILValue> ElementCountOperands);
 
@@ -2189,6 +2211,9 @@ class AllocRefDynamicInst final
   MutableArrayRef<Operand> getTypeDependentOperands() {
     return getAllOperands().slice(getNumTailTypes() + 1);
   }
+  // Is the deinit and the size of the dynamic type known to be equivalent to
+  // the the base type (i.e `this->getType()`).
+  bool isDynamicTypeDeinitAndSizeKnownEquivalentToBaseType() const;
 };
 
 /// This represents the allocation of a heap box for a Swift value of some type.
@@ -7649,7 +7674,9 @@ class DeallocStackRefInst
   DeallocStackRefInst(SILDebugLocation DebugLoc, SILValue Operand)
       : UnaryInstructionBase(DebugLoc, Operand) {}
 public:
-  AllocRefInst *getAllocRef() { return cast<AllocRefInst>(getOperand()); }
+  AllocRefInstBase *getAllocRef() {
+    return cast<AllocRefInstBase>(getOperand());
+  }
 };
 
 /// Deallocate memory for a reference type instance from a destructor or

lib/IRGen/GenClass.cpp

@@ -828,12 +828,32 @@ llvm::Value *irgen::emitClassAllocationDynamic(IRGenFunction &IGF,
                                                llvm::Value *metadata,
                                                SILType selfType,
                                                bool objc,
+                                               int &StackAllocSize,
                                                TailArraysRef TailArrays) {
   // If we need to use Objective-C allocation, do so.
   if (objc) {
+    StackAllocSize = -1;
     return emitObjCAllocObjectCall(IGF, metadata, selfType);
   }
 
+  llvm::Value *Promoted;
+  auto &classTI = IGF.getTypeInfo(selfType).as<ClassTypeInfo>();
+  auto &classLayout = classTI.getClassLayout(IGF.IGM, selfType,
+                                          /*forBackwardDeployment=*/false);
+
+  // If we are allowed to allocate on the stack we are allowed to use
+  // `selfType`'s size assumptions.
+  if (StackAllocSize >= 0 &&
+      (Promoted = stackPromote(IGF, classLayout, StackAllocSize,
+                                           TailArrays))) {
+    llvm::Value *val = IGF.Builder.CreateBitCast(Promoted,
+                                                 IGF.IGM.RefCountedPtrTy);
+    val = IGF.emitInitStackObjectCall(metadata, val, "reference.new");
+
+    llvm::Type *destType = classLayout.getType()->getPointerTo();
+    return IGF.Builder.CreateBitCast(val, destType);
+  }
+
   // Otherwise, allocate using Swift's routines.
   llvm::Value *size, *alignMask;
   std::tie(size, alignMask)
@@ -845,10 +865,8 @@ llvm::Value *irgen::emitClassAllocationDynamic(IRGenFunction &IGF,
 
   llvm::Value *val = IGF.emitAllocObjectCall(metadata, size, alignMask,
                                              "reference.new");
-  auto &classTI = IGF.getTypeInfo(selfType).as<ClassTypeInfo>();
-  auto &layout = classTI.getClassLayout(IGF.IGM, selfType,
-                                        /*forBackwardDeployment=*/false);
-  llvm::Type *destType = layout.getType()->getPointerTo();
+  StackAllocSize = -1;
+  llvm::Type *destType = classLayout.getType()->getPointerTo();
   return IGF.Builder.CreateBitCast(val, destType);
 }
 

lib/IRGen/GenClass.h

@@ -167,10 +167,12 @@ namespace irgen {
                   bool objc, int &StackAllocSize, TailArraysRef TailArrays);
 
   /// Emit an allocation of a class using a metadata value.
-  llvm::Value *emitClassAllocationDynamic(IRGenFunction &IGF, 
+  llvm::Value *emitClassAllocationDynamic(IRGenFunction &IGF,
                                           llvm::Value *metadata,
                                           SILType selfType,
-                                          bool objc, TailArraysRef TailArrays);
+                                          bool objc,
+                                          int &StackAllocSize,
+                                          TailArraysRef TailArrays);
 
   /// Emit class deallocation.
   void emitClassDeallocation(IRGenFunction &IGF, SILType selfType,

lib/IRGen/IRGenSIL.cpp

@@ -5425,14 +5425,30 @@ void IRGenSILFunction::visitAllocRefInst(swift::AllocRefInst *i) {
 }
 
 void IRGenSILFunction::visitAllocRefDynamicInst(swift::AllocRefDynamicInst *i) {
+  int StackAllocSize = -1;
+  if (i->canAllocOnStack()) {
+    assert(i->isDynamicTypeDeinitAndSizeKnownEquivalentToBaseType());
+    estimateStackSize();
+    // Is there enough space for stack allocation?
+    StackAllocSize = IGM.IRGen.Opts.StackPromotionSizeLimit - EstimatedStackSize;
+  }
+
   SmallVector<std::pair<SILType, llvm::Value *>, 4> TailArrays;
   buildTailArrays(*this, TailArrays, i);
 
   Explosion metadata = getLoweredExplosion(i->getMetatypeOperand());
   auto metadataValue = metadata.claimNext();
   llvm::Value *alloced = emitClassAllocationDynamic(*this, metadataValue,
                                                     i->getType(), i->isObjC(),
+                                                    StackAllocSize,
                                                     TailArrays);
+
+  if (StackAllocSize >= 0) {
+    // Remember that this alloc_ref_dynamic allocates the object on the stack.
+    StackAllocs.insert(i);
+    EstimatedStackSize += StackAllocSize;
+  }
+
   Explosion e;
   e.add(alloced);
   setLoweredExplosion(i, e);
@@ -5476,7 +5492,7 @@ void IRGenSILFunction::visitDeallocRefInst(swift::DeallocRefInst *i) {
   // Lower the operand.
   Explosion self = getLoweredExplosion(i->getOperand());
   auto selfValue = self.claimNext();
-  auto *ARI = dyn_cast<AllocRefInst>(i->getOperand());
+  auto *ARI = dyn_cast<AllocRefInstBase>(i->getOperand());
   if (ARI && StackAllocs.count(ARI)) {
     // We can ignore dealloc_refs for stack allocated objects.
     //

lib/SIL/IR/Linker.cpp

@@ -327,6 +327,20 @@ void SILLinkerVisitor::visitInitExistentialRefInst(
     visitProtocolConformance(C, Optional<SILDeclRef>());
   }
 }
+void SILLinkerVisitor::visitAllocRefDynamicInst(AllocRefDynamicInst *ARI) {
+  if (!isLinkAll())
+    return;
+
+  if (!ARI->isDynamicTypeDeinitAndSizeKnownEquivalentToBaseType())
+    return;
+
+  // Grab the class decl from the alloc ref inst.
+  ClassDecl *D = ARI->getType().getClassOrBoundGenericClass();
+  if (!D)
+    return;
+
+  linkInVTable(D);
+}
 
 void SILLinkerVisitor::visitAllocRefInst(AllocRefInst *ARI) {
   if (!isLinkAll())

lib/SIL/IR/Linker.h

@@ -72,6 +72,7 @@ class SILLinkerVisitor : public SILInstructionVisitor<SILLinkerVisitor, void> {
   void visitInitExistentialAddrInst(InitExistentialAddrInst *IEI);
   void visitInitExistentialRefInst(InitExistentialRefInst *IERI);
   void visitAllocRefInst(AllocRefInst *ARI);
+  void visitAllocRefDynamicInst(AllocRefDynamicInst *ARI);
   void visitMetatypeInst(MetatypeInst *MI);
 
 private: