IRGen: Use the header size for the initial offset in non fixed heap layouts

When the first element in the heap layout was non fixed we would use the
mininum size of the total heap layout for the initial offset. This would
create unneccessary large heap layouts.

rdar://61716736

Author: aschwaighofer

lib/IRGen/GenHeap.cpp

@@ -285,7 +285,7 @@ static llvm::Value *calcInitOffset(swift::irgen::IRGenFunction &IGF,
                                    const swift::irgen::HeapLayout &layout) {
   llvm::Value *offset = nullptr;
   if (i == 0) {
-    auto startoffset = layout.getSize();
+    auto startoffset = layout.getHeaderSize();
     offset = llvm::ConstantInt::get(IGF.IGM.SizeTy, startoffset.getValue());
     return offset;
   }

lib/IRGen/StructLayout.cpp

@@ -70,6 +70,7 @@ StructLayout::StructLayout(IRGenModule &IGM,
     assert(!builder.empty() == requiresHeapHeader(layoutKind));
     MinimumAlign = Alignment(1);
     MinimumSize = Size(0);
+    headerSize = builder.getHeaderSize();
     SpareBits.clear();
     IsFixedLayout = true;
     IsKnownPOD = IsPOD;
@@ -79,6 +80,7 @@ StructLayout::StructLayout(IRGenModule &IGM,
   } else {
     MinimumAlign = builder.getAlignment();
     MinimumSize = builder.getSize();
+    headerSize = builder.getHeaderSize();
     SpareBits = builder.getSpareBits();
     IsFixedLayout = builder.isFixedLayout();
     IsKnownPOD = builder.isPOD();
@@ -186,13 +188,15 @@ void StructLayoutBuilder::addHeapHeader() {
   CurSize = IGM.RefCountedStructSize;
   CurAlignment = IGM.getPointerAlignment();
   StructFields.push_back(IGM.RefCountedStructTy);
+  headerSize = CurSize;
 }
 
 void StructLayoutBuilder::addNSObjectHeader() {
   assert(StructFields.empty() && "adding heap header at a non-zero offset");
   CurSize = IGM.getPointerSize();
   CurAlignment = IGM.getPointerAlignment();
   StructFields.push_back(IGM.ObjCClassPtrTy);
+  headerSize = CurSize;
 }
 
 bool StructLayoutBuilder::addFields(llvm::MutableArrayRef<ElementLayout> elts,

lib/IRGen/StructLayout.h

@@ -253,6 +253,7 @@ class StructLayoutBuilder {
   IRGenModule &IGM;
   SmallVector<llvm::Type*, 8> StructFields;
   Size CurSize = Size(0);
+  Size headerSize = Size(0);
 private:
   Alignment CurAlignment = Alignment(1);
   SmallVector<SpareBitVector, 8> CurSpareBits;
@@ -314,6 +315,9 @@ class StructLayoutBuilder {
   /// Return the size of the structure built so far.
   Size getSize() const { return CurSize; }
 
+  // Return the size of the header.
+  Size getHeaderSize() const { return headerSize; }
+
   /// Return the alignment of the structure built so far.
   Alignment getAlignment() const { return CurAlignment; }
 
@@ -350,6 +354,9 @@ class StructLayout {
 
   /// The statically-known minimum bound on the size.
   Size MinimumSize;
+
+  /// The size of a header if present.
+  Size headerSize;
 
   /// The statically-known spare bit mask.
   SpareBitVector SpareBits;
@@ -386,6 +393,7 @@ class StructLayout {
                ArrayRef<ElementLayout> elements)
     : MinimumAlign(builder.getAlignment()),
       MinimumSize(builder.getSize()),
+      headerSize(builder.getHeaderSize()),
       SpareBits(builder.getSpareBits()),
       IsFixedLayout(builder.isFixedLayout()),
       IsKnownPOD(builder.isPOD()),
@@ -401,6 +409,7 @@ class StructLayout {
 
   llvm::Type *getType() const { return Ty; }
   Size getSize() const { return MinimumSize; }
+  Size getHeaderSize() const { return headerSize; }
   Alignment getAlignment() const { return MinimumAlign; }
   const SpareBitVector &getSpareBits() const { return SpareBits; }
   SpareBitVector &getSpareBits() { return SpareBits; }

test/IRGen/partial_apply.sil

@@ -1,9 +1,12 @@
-// RUN: %target-swift-frontend %s -emit-ir | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -enable-library-evolution -emit-module-path=%t/resilient_struct.swiftmodule -module-name=resilient_struct %S/../Inputs/resilient_struct.swift
+// RUN: %target-swift-frontend -I %t -emit-ir %s | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
 
 // REQUIRES: CPU=x86_64
 
 import Builtin
 import Swift
+import resilient_struct
 
 class SwiftClass {}
 sil_vtable SwiftClass {}
@@ -738,3 +741,29 @@ bb0(%x : $*SwiftClassPair):
   %t = tuple()
   return %t : $()
 }
+
+sil public_external @closure : $@convention(thin) (@in_guaranteed ResilientInt, @guaranteed SwiftClass) -> ()
+
+// Make sure that we use the heap header size (16) for the initial offset.
+// CHECK: define swiftcc void @test_initial_offset(%swift.opaque* noalias nocapture %0, %T13partial_apply10SwiftClassC* %1)
+// CHECK:   [[T0:%.*]] = call swiftcc %swift.metadata_response @"$s16resilient_struct12ResilientIntVMa"
+// CHECK:   [[MD:%.*]] = extractvalue %swift.metadata_response [[T0]], 0
+// CHECK:   [[CAST:%.*]] = bitcast %swift.type* [[MD]] to i8***
+// CHECK:   [[VWT_PTR:%.*]] = getelementptr inbounds i8**, i8*** [[CAST]], i64 -1
+// CHECK:   [[VWT_TMP:%.*]] = load i8**, i8*** [[VWT_PTR]]
+// CHECK:   [[VWT:%.*]] = bitcast i8** [[VWT_TMP]] to %swift.vwtable*
+// CHECK:   [[FLAGS_PTR:%.*]] = getelementptr inbounds %swift.vwtable, %swift.vwtable* [[VWT]], i32 0, i32 10
+// CHECK:   [[FLAGS:%.*]] = load i32, i32* [[FLAGS_PTR]]
+// CHECK:   [[FLAGS2:%.*]] = zext i32 [[FLAGS]] to i64
+// CHECK:   [[ALIGNMASK:%.*]] = and i64 [[FLAGS2]], 255
+// CHECK:   = xor i64 [[ALIGNMASK]], -1
+// CHECK:   = add i64 16, [[ALIGNMASK]]
+
+sil @test_initial_offset : $@convention(thin) (@in_guaranteed ResilientInt, @guaranteed SwiftClass) -> () {
+bb0(%x : $*ResilientInt, %y : $SwiftClass):
+  %f = function_ref @closure : $@convention(thin) (@in_guaranteed ResilientInt, @guaranteed SwiftClass) -> ()
+  %p = partial_apply [callee_guaranteed] %f(%x, %y) : $@convention(thin) (@in_guaranteed ResilientInt, @guaranteed SwiftClass) -> ()
+  release_value %p : $@callee_guaranteed () ->()
+  %t = tuple()
+  return %t : $()
+}