From 350e92119b1a7f616dc1cce46536cbaf4b04dd6a Mon Sep 17 00:00:00 2001
From: Joe Shajrawi <shajrawi@apple.com>
Date: Thu, 25 Jan 2018 17:22:26 -0800
Subject: [PATCH] IRGen: fix a corner-case wherein a partial_apply was not
 converted in LoadableByAddress It cleans up shouldTransformParameter, fixing
 a corner-case for optional function parameter, and re-creates *all* partial
 applies that only contain function/optional function parameters

---
 lib/IRGen/LoadableByAddress.cpp         | 24 ++++++++++--------------
 test/IRGen/big_types_corner_cases.swift | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/lib/IRGen/LoadableByAddress.cpp b/lib/IRGen/LoadableByAddress.cpp
index 4054a8e8afb45..ecdeae69ab77e 100644
--- a/lib/IRGen/LoadableByAddress.cpp
+++ b/lib/IRGen/LoadableByAddress.cpp
@@ -79,24 +79,16 @@ static bool shouldTransformFunctionType(GenericEnvironment *env,
                                         CanSILFunctionType fnType,
                                         irgen::IRGenModule &IGM);
 
+static SILParameterInfo getNewParameter(GenericEnvironment *env,
+                                        SILParameterInfo param,
+                                        irgen::IRGenModule &IGM);
+
 static bool shouldTransformParameter(GenericEnvironment *env,
                                      SILParameterInfo param,
                                      irgen::IRGenModule &IGM) {
-  SILType storageType = param.getSILStorageType();
-
-  // FIXME: only function types and not recursively-transformable types?
-  if (auto fnType = storageType.getAs<SILFunctionType>())
-    return shouldTransformFunctionType(env, fnType, IGM);
 
-  switch (param.getConvention()) {
-  case ParameterConvention::Indirect_In_Guaranteed:
-  case ParameterConvention::Indirect_Inout:
-  case ParameterConvention::Indirect_InoutAliasable:
-  case ParameterConvention::Indirect_In:
-    return false;
-  default:
-    return isLargeLoadableType(env, storageType, IGM);
-  }
+  auto newParam = getNewParameter(env, param, IGM);
+  return (param != newParam);
 }
 
 static bool shouldTransformFunctionType(GenericEnvironment *env,
@@ -2558,6 +2550,10 @@ void LoadableByAddress::run() {
           if (isa<ProjectBlockStorageInst>(dest)) {
             storeToBlockStorageInstrs.insert(SI);
           }
+        } else if (auto *PAI = dyn_cast<PartialApplyInst>(&I)) {
+          if (modApplies.count(PAI) == 0) {
+            modApplies.insert(PAI);
+          }
         }
       }
     }
diff --git a/test/IRGen/big_types_corner_cases.swift b/test/IRGen/big_types_corner_cases.swift
index ae2cfdb3ba7ee..7d6a771877fa0 100644
--- a/test/IRGen/big_types_corner_cases.swift
+++ b/test/IRGen/big_types_corner_cases.swift
@@ -205,3 +205,17 @@ func bigStructGet() -> BigStruct {
 public func testGetFunc() {
   let testGetPtr: @convention(thin) () -> BigStruct = bigStructGet
 }
+
+// CHECK-LABEL: define{{( protected)?}} hidden swiftcc void @"$S22big_types_corner_cases7TestBigC4testyyF"(%T22big_types_corner_cases7TestBigC* swiftself)
+// CHECK: [[CALL1:%.*]] = call %swift.type* @"$SSayy22big_types_corner_cases9BigStructVcSgGMa"
+// CHECK: [[CALL2:%.*]] = call i8** @"$SSayy22big_types_corner_cases9BigStructVcSgGSayxGs10CollectionsWl
+// CHECK: call swiftcc void @"$Ss10CollectionPsE5index5where5IndexQzSgSb7ElementQzKc_tKF"(%TSq.0* noalias nocapture sret {{.*}}, i8* bitcast (i1 (%T22big_types_corner_cases9BigStructVytIegir_Sg*, %swift.refcounted*, %swift.error**)* @"$S22big_types_corner_cases9BigStructVIegy_SgSbs5Error_pIgxdzo_ACytIegir_SgSbsAE_pIgidzo_TRTA" to i8*), %swift.refcounted* {{.*}}, %swift.type* [[CALL1]], i8** [[CALL2]], %swift.opaque* noalias nocapture swiftself
+// CHECK: ret void
+class TestBig {
+    typealias Handler = (BigStruct) -> Void
+
+    func test() {
+        let arr = [Handler?]()
+        let d = arr.index(where: { _ in true })
+    }
+}
\ No newline at end of file