[IRGen] Single payload enums consumes call deinits

nate-chandler · nate-chandler · commit 0075682b4d4c · 2024-03-01T21:01:50.000-08:00
The include the necessary metadata.
diff --git a/lib/IRGen/GenEnum.cpp b/lib/IRGen/GenEnum.cpp
@@ -1639,10 +1639,11 @@ namespace {
       return {payload, extraTag};
     }
     std::pair<EnumPayload, llvm::Value *>
-    getPayloadAndExtraTagFromExplosionOutlined(IRGenFunction &IGF,
-                                               Explosion &src) const {
+    getPayloadAndExtraTagFromExplosionOutlined(
+        IRGenFunction &IGF, Explosion &src,
+        OutliningMetadataCollector *collector) const {
       EnumPayload payload;
-      unsigned claimSZ = src.size();
+      unsigned claimSZ = src.size() - (collector ? collector->size() : 0);
       if (ExtraTagBitCount > 0) {
         --claimSZ;
       }
@@ -1833,7 +1834,7 @@ namespace {
       EnumPayload payload;
       llvm::Value *extraTag;
       std::tie(payload, extraTag) =
-          getPayloadAndExtraTagFromExplosionOutlined(IGF, src);
+          getPayloadAndExtraTagFromExplosionOutlined(IGF, src, nullptr);
       llvm::BasicBlock *endBB =
           testFixedEnumContainsPayload(IGF, payload, extraTag);
 
@@ -1856,25 +1857,33 @@ namespace {
 
     void emitCallToConsumeEnumFunction(IRGenFunction &IGF, Explosion &src,
                                        SILType theEnumType) const {
+      OutliningMetadataCollector collector(theEnumType, IGF, LayoutIsNotNeeded,
+                                           DeinitIsNeeded);
+      IGF.getTypeInfo(theEnumType)
+          .collectMetadataForOutlining(collector, theEnumType);
       if (!consumeEnumFunction)
-        consumeEnumFunction = emitConsumeEnumFunction(IGF.IGM, loweredType);
+        consumeEnumFunction =
+            emitConsumeEnumFunction(IGF.IGM, theEnumType, collector);
       Explosion tmp;
-      fillExplosionForOutlinedCall(IGF, src, tmp);
+      fillExplosionForOutlinedCall(IGF, src, tmp, &collector);
       llvm::CallInst *call = IGF.Builder.CreateCallWithoutDbgLoc(
           consumeEnumFunction->getFunctionType(), consumeEnumFunction,
           tmp.claimAll());
       call->setCallingConv(IGM.DefaultCC);
     }
 
-    llvm::Function *emitConsumeEnumFunction(IRGenModule &IGM,
-                                            SILType theEnumType) const {
+    llvm::Function *
+    emitConsumeEnumFunction(IRGenModule &IGM, SILType theEnumType,
+                            OutliningMetadataCollector &collector) const {
       IRGenMangler Mangler;
       auto manglingBits =
         getTypeAndGenericSignatureForManglingOutlineFunction(theEnumType);
       std::string name =
         Mangler.mangleOutlinedConsumeFunction(manglingBits.first,
                                               manglingBits.second);
-      auto func = createOutlineLLVMFunction(IGM, name, PayloadTypesAndTagType);
+      SmallVector<llvm::Type *, 2> params(PayloadTypesAndTagType);
+      collector.addPolymorphicParameterTypes(params);
+      auto func = createOutlineLLVMFunction(IGM, name, params);
 
       IRGenFunction IGF(IGM, func);
       Explosion src = IGF.collectParameters();
@@ -1884,7 +1893,8 @@ namespace {
       EnumPayload payload;
       llvm::Value *extraTag;
       std::tie(payload, extraTag) =
-          getPayloadAndExtraTagFromExplosionOutlined(IGF, src);
+          getPayloadAndExtraTagFromExplosionOutlined(IGF, src, &collector);
+      collector.bindPolymorphicParameters(IGF, src);
       llvm::BasicBlock *endBB =
           testFixedEnumContainsPayload(IGF, payload, extraTag);
 
@@ -2678,8 +2688,10 @@ namespace {
       }
     }
 
-    void fillExplosionForOutlinedCall(IRGenFunction &IGF, Explosion &src,
-                                      Explosion &out) const {
+    void
+    fillExplosionForOutlinedCall(IRGenFunction &IGF, Explosion &src,
+                                 Explosion &out,
+                                 OutliningMetadataCollector *collector) const {
       assert(out.empty() && "Out explosion must be empty!");
       EnumPayload payload;
       llvm::Value *extraTag;
@@ -2688,8 +2700,16 @@ namespace {
       payload.explode(IGM, out);
       if (extraTag)
         out.add(extraTag);
+
+      if (!collector)
+        return;
+      llvm::SmallVector<llvm::Value *, 4> args;
+      collector->addPolymorphicArguments(args);
+      for (auto *arg : args) {
+        out.add(arg);
+      }
     }
-    
+
     void unpackIntoPayloadExplosion(IRGenFunction &IGF,
                                     Explosion &asEnumIn,
                                     Explosion &asPayloadOut) const {
@@ -2749,7 +2769,7 @@ namespace {
         if (!copyEnumFunction)
           copyEnumFunction = emitCopyEnumFunction(IGM, loweredType);
         Explosion tmp;
-        fillExplosionForOutlinedCall(IGF, src, tmp);
+        fillExplosionForOutlinedCall(IGF, src, tmp, nullptr);
         llvm::CallInst *call = IGF.Builder.CreateCallWithoutDbgLoc(
             copyEnumFunction->getFunctionType(), copyEnumFunction,
             tmp.getAll());
@@ -2818,7 +2838,7 @@ namespace {
           IGF.Builder.emitBlock(endBB);
           return;
         }
-        emitCallToConsumeEnumFunction(IGF, src, loweredType);
+        emitCallToConsumeEnumFunction(IGF, src, T);
         return;
       }
 
diff --git a/lib/IRGen/Outlining.h b/lib/IRGen/Outlining.h
@@ -80,6 +80,10 @@ class OutliningMetadataCollector {
                              LayoutIsNeeded_t needsLayout,
                              DeinitIsNeeded_t needsDeinitTypes);
 
+  unsigned size() const {
+    return Subs.has_value() ? Requirements.size() : Values.size();
+  }
+
   // If any local type data is needed for \p type, add it.
   //
   // NOTE: To be called from TypeData instances.
diff --git a/test/IRGen/moveonly_value_functions.swift b/test/IRGen/moveonly_value_functions.swift
@@ -49,6 +49,16 @@ public struct InnerDeinitingDestructableNC<T> : ~Copyable {
   }
 }
 
+public struct InnerDeinitingWithLayoutNC<T> : ~Copyable {
+  public let t: T
+  deinit {}
+}
+
+public struct InnerDeinitingWithoutLayoutNC<T>: ~Copyable {
+  public let ptr: Int
+  deinit {}
+}
+
 public struct OuterDeinitingNC_1<T> : ~Copyable {
   public let i1: Int
   public let c1: C<T>
@@ -95,6 +105,21 @@ extension GenericContext_1 : P where T : P {
   }
 }
 
+public enum OuterSinglePayloadNC_1<T>: ~Copyable {
+  case none
+  case some(InnerDeinitingWithoutLayoutNC<T>)
+}
+
+public enum OuterSinglePayloadNC_2<T>: ~Copyable {
+  case none
+  case some(InnerDeinitingReleasableNC<T>)
+}
+
+public enum OuterSinglePayloadNC_3<T>: ~Copyable {
+  case none
+  case some(InnerDeinitingDestructableNC<T>)
+}
+
 // Destroyed value:
 // - has deinit
 // On lifetime end:
@@ -218,3 +243,71 @@ public func takeOuterNC_2<T>(_ o: consuming OuterNC_2<T>) {
 // CHECK-SAME:        ptr noalias swiftself %0)
 // CHECK:       }
 public func takeGenericContext_1OuterNC_1<T : P>(_ e: consuming GenericContext_1<T>.OuterNC_1) {}
+
+// CHECK-LABEL: define{{.*}} @"$s24moveonly_value_functions26takeOuterSinglePayloadNC_1yyAA0efgH2_1OyxGnlF"(
+//           :      i64 %0, 
+//           :      i8 %1, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_1OyxGlWOe"(
+//           :        i64 %0, 
+//           :        i1 %2, 
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+// CHECK-LABEL: define{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_1OyxGlWOe"(
+//           :      i64 %0, 
+//           :      i1 %1, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s24moveonly_value_functions29InnerDeinitingWithoutLayoutNCVfD"(
+//           :        i64 %0,
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+public func takeOuterSinglePayloadNC_1<T>(_ e: consuming OuterSinglePayloadNC_1<T>) {}
+
+// CHECK-LABEL: define{{.*}} @"$s24moveonly_value_functions26takeOuterSinglePayloadNC_2yyAA0efgH2_2OyxGnlF"(
+// CHECK-SAME:      ptr noalias nocapture dereferenceable(64) %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_2OyxGlWOs"(
+// CHECK-SAME:        ptr %0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+// CHECK-LABEL: define{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_2OyxGlWOs"(
+// CHECK-SAME:      ptr %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_2OyxGlWOe"(
+//            :       i64 %2, 
+//            :       i64 %4, 
+//            :       i64 %6, 
+//            :       i64 %8, 
+//            :       i64 %10, 
+//            :       i64 %12, 
+//            :       i64 %14, 
+//            :       i64 %16, 
+//            :       ptr %T)
+// CHECK:       }
+public func takeOuterSinglePayloadNC_2<T>(_ e: consuming OuterSinglePayloadNC_2<T>) {}
+
+// CHECK-LABEL: define{{.*}} @"$s24moveonly_value_functions26takeOuterSinglePayloadNC_3yyAA0efgH2_3OyxGnlF"(
+// CHECK-SAME:      ptr noalias %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME: {
+// CHECK:         call{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_3OyxGlWOh"(
+// CHECK-SAME:        ptr %0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+// CHECK-LABEL: define{{.*}} @"$s24moveonly_value_functions22OuterSinglePayloadNC_3OyxGlWOh"(
+// CHECK-SAME:      ptr %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         [[RESPONSE:%[^,]+]] = call{{.*}} @"$s24moveonly_value_functions28InnerDeinitingDestructableNCVMa"(
+//           :        i64 0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:         [[METADATA:%[^,]+]] = extractvalue %swift.metadata_response [[RESPONSE:%[^,]+]], 0
+// CHECK:         call{{.*}} @"$s24moveonly_value_functions28InnerDeinitingDestructableNCVfD"(
+// CHECK-SAME:        ptr [[METADATA:%[^,]+]], 
+//           :        ptr noalias swiftself %0)
+// CHECK:       }
+public func takeOuterSinglePayloadNC_3<T>(_ e: consuming OuterSinglePayloadNC_3<T>) {}
diff --git a/test/IRGen/moveonly_value_functions_onone.swift b/test/IRGen/moveonly_value_functions_onone.swift
@@ -99,6 +99,21 @@ extension GenericContext_1 : P where T : P {
   }
 }
 
+public enum OuterSinglePayloadNC_1<T>: ~Copyable {
+  case none
+  case some(InnerDeinitingWithoutLayoutNC<T>)
+}
+
+public enum OuterSinglePayloadNC_2<T>: ~Copyable {
+  case none
+  case some(InnerDeinitingReleasableNC<T>)
+}
+
+public enum OuterSinglePayloadNC_3<T>: ~Copyable {
+  case none
+  case some(InnerDeinitingDestructableNC<T>)
+}
+
 // Destroyed value:
 // - has deinit
 // On lifetime end:
@@ -222,3 +237,68 @@ public func takeOuterNC_2<T>(_ o: consuming OuterNC_2<T>) {
 // CHECK-SAME:        ptr noalias swiftself %0)
 // CHECK:       }
 public func takeGenericContext_1OuterNC_1<T : P>(_ e: consuming GenericContext_1<T>.OuterNC_1) {}
+
+// CHECK-LABEL: define{{.*}} @"$s30moveonly_value_functions_onone26takeOuterSinglePayloadNC_1yyAA0fghI2_1OyxGnlF"(
+//           :      i64 %0, 
+//           :      i8 %1, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s30moveonly_value_functions_onone22OuterSinglePayloadNC_1OyxGlWOh"(
+//           :        i64 %0, 
+//           :        i1 %2, 
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+// CHECK-LABEL: define{{.*}} @"$s30moveonly_value_functions_onone22OuterSinglePayloadNC_1OyxGlWOh"(
+//           :      i64 %0, 
+//           :      i1 %1, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s30moveonly_value_functions_onone29InnerDeinitingWithoutLayoutNCVfD"(
+//           :        i64 %0,
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+public func takeOuterSinglePayloadNC_1<T>(_ e: consuming OuterSinglePayloadNC_1<T>) {}
+
+// CHECK-LABEL: define{{.*}} @"$s30moveonly_value_functions_onone26takeOuterSinglePayloadNC_2yyAA0fghI2_2OyxGnlF"(
+// CHECK-SAME:      ptr noalias nocapture dereferenceable(64) %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         call{{.*}} @"$s30moveonly_value_functions_onone22OuterSinglePayloadNC_2OyxGlWOh"(
+//           :        ptr %0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+// CHECK-LABEL: define{{.*}} @"$s30moveonly_value_functions_onone22OuterSinglePayloadNC_2OyxGlWOh"(
+// CHECK-SAME:      ptr %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         [[RESPONSE:%[^,]+]] = call swiftcc %swift.metadata_response @"$s30moveonly_value_functions_onone26InnerDeinitingReleasableNCVMa"(
+// CHECK-SAME:        i64 0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:         [[METADATA:%[^,]+]] = extractvalue %swift.metadata_response [[RESPONSE]], 0
+// CHECK:         call swiftcc void @"$s30moveonly_value_functions_onone26InnerDeinitingReleasableNCVfD"(
+// CHECK-SAME:        ptr [[METADATA]],
+// CHECK-SAME:        ptr noalias nocapture swiftself dereferenceable(64) %0)
+// CHECK:       }
+public func takeOuterSinglePayloadNC_2<T>(_ e: consuming OuterSinglePayloadNC_2<T>) {}
+
+// CHECK-LABEL: define{{.*}} @"$s30moveonly_value_functions_onone26takeOuterSinglePayloadNC_3yyAA0fghI2_3OyxGnlF"(
+// CHECK-SAME:      ptr noalias %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME: {
+// CHECK:         call{{.*}} @"$s30moveonly_value_functions_onone22OuterSinglePayloadNC_3OyxGlWOh"(
+//           :        ptr %0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:       }
+// CHECK-LABEL: define{{.*}} @"$s30moveonly_value_functions_onone22OuterSinglePayloadNC_3OyxGlWOh"(
+// CHECK-SAME:      ptr %0, 
+// CHECK-SAME:      ptr %T)
+// CHECK-SAME:  {
+// CHECK:         [[RESPONSE:%[^,]+]] = call{{.*}} @"$s30moveonly_value_functions_onone28InnerDeinitingDestructableNCVMa"(
+//           :        i64 0, 
+// CHECK-SAME:        ptr %T)
+// CHECK:         [[METADATA:%[^,]+]] = extractvalue %swift.metadata_response [[RESPONSE:%[^,]+]], 0
+// CHECK:         call{{.*}} @"$s30moveonly_value_functions_onone28InnerDeinitingDestructableNCVfD"(
+// CHECK-SAME:        ptr [[METADATA:%[^,]+]], 
+//           :        ptr noalias swiftself %0)
+// CHECK:       }
+public func takeOuterSinglePayloadNC_3<T>(_ e: consuming OuterSinglePayloadNC_3<T>) {}
diff --git a/validation-test/IRGen/moveonly_partial_consumption_linked_list.swift b/validation-test/IRGen/moveonly_partial_consumption_linked_list.swift
