Transform tuple return types that contain function signatures.

Author: dingobye

lib/IRGen/LoadableByAddress.cpp

@@ -296,9 +296,15 @@ bool LargeSILTypeMapper::shouldTransformResults(GenericEnvironment *genEnv,
   if (!modifiableFunction(loweredTy)) {
     return false;
   }
+
   if (loweredTy->getNumResults() != 1) {
-    return false;
+    auto resultType = loweredTy->getAllResultsType();
+    auto newResultType = getNewSILType(genEnv, resultType, Mod);
+    bool hasFuncSig = containsFunctionSignature(genEnv, Mod,
+                                                resultType, newResultType);
+    return hasFuncSig;
   }
+
   auto singleResult = loweredTy->getSingleResult();
   auto resultStorageType = singleResult.getSILStorageType();
   auto newResultStorageType = getNewSILType(genEnv, resultStorageType, Mod);
@@ -1287,8 +1293,7 @@ SILArgument *LoadableStorageAllocation::replaceArgType(SILBuilder &argBuilder,
 void LoadableStorageAllocation::insertIndirectReturnArgs() {
   GenericEnvironment *genEnv = pass.F->getGenericEnvironment();
   auto loweredTy = pass.F->getLoweredFunctionType();
-  auto singleResult = loweredTy->getSingleResult();
-  SILType resultStorageType = singleResult.getSILStorageType();
+  SILType resultStorageType = loweredTy->getAllResultsType();
   auto canType = resultStorageType.getASTType();
   if (canType->hasTypeParameter()) {
     assert(genEnv && "Expected a GenericEnv");
@@ -1368,19 +1373,26 @@ void LoadableStorageAllocation::convertApplyResults() {
                                               pass.Mod)) {
         continue;
       }
-      auto singleResult = origSILFunctionType->getSingleResult();
-      auto resultStorageType = singleResult.getSILStorageType();
+      auto resultStorageType = origSILFunctionType->getAllResultsType();
       if (!isLargeLoadableType(genEnv, resultStorageType, pass.Mod)) {
-        // Make sure it is a function type
-        if (!resultStorageType.is<SILFunctionType>()) {
-          // Check if it is an optional function type
-          auto optionalType = resultStorageType.getOptionalObjectType();
-          assert(optionalType &&
-                 "Expected SILFunctionType or Optional for the result type");
-          assert(optionalType.is<SILFunctionType>() &&
-                 "Expected a SILFunctionType inside the optional Type");
-          (void)optionalType;
-        }
+        // Make sure it contains a function type
+        auto numFuncTy = llvm::count_if(origSILFunctionType->getResults(),
+            [](const SILResultInfo &origResult) {
+              auto resultStorageTy = origResult.getSILStorageType();
+              // Check if it is a function type
+              if (resultStorageTy.is<SILFunctionType>()) {
+                return true;
+              }
+              // Check if it is an optional function type
+              auto optionalType = resultStorageTy.getOptionalObjectType();
+              if (optionalType && optionalType.is<SILFunctionType>()) {
+                return true;
+              }
+              return false;
+            });
+        assert(numFuncTy != 0 &&
+               "Expected a SILFunctionType inside the result Type");
+        (void)numFuncTy;
         continue;
       }
       auto newSILType =
@@ -2213,10 +2225,25 @@ static bool rewriteFunctionReturn(StructLoweringState &pass) {
     return true;
   } else if (containsFunctionSignature(genEnv, pass.Mod, resultTy,
                                        newSILType) &&
-             (resultTy != newSILType) && (loweredTy->getNumResults() == 1)) {
-    SILResultInfo origResultInfo = loweredTy->getSingleResult();
-    SILResultInfo newSILResultInfo(newSILType.getASTType(),
-                                   origResultInfo.getConvention());
+             (resultTy != newSILType)) {
+
+    llvm::SmallVector<SILResultInfo, 2> newSILResultInfo;
+    if (auto tupleType = newSILType.getAs<TupleType>()) {
+      auto originalResults = loweredTy->getResults();
+      for (unsigned int i = 0; i < originalResults.size(); ++i) {
+        auto origResultInfo = originalResults[i];
+        auto canElem = tupleType.getElementType(i);
+        SILType objectType = SILType::getPrimitiveObjectType(canElem);
+        auto newResult = SILResultInfo(objectType.getASTType(), origResultInfo.getConvention());
+        newSILResultInfo.push_back(newResult);
+      }
+    } else {
+      assert(loweredTy->getNumResults() == 1 && "Expected a single result");
+      auto origResultInfo = loweredTy->getSingleResult();
+      auto newResult = SILResultInfo(newSILType.getASTType(), origResultInfo.getConvention());
+      newSILResultInfo.push_back(newResult);
+    }
+
     auto NewTy = SILFunctionType::get(
         loweredTy->getGenericSignature(), loweredTy->getExtInfo(),
         loweredTy->getCoroutineKind(),
@@ -2302,9 +2329,6 @@ getOperandTypeWithCastIfNecessary(SILInstruction *containingInstr, SILValue op,
     if (!genEnv && funcType->isPolymorphic()) {
       genEnv = getGenericEnvironment(funcType);
     }
-    if (!Mapper.shouldTransformFunctionType(genEnv, funcType, Mod)) {
-      return op;
-    }
     auto newFnType = Mapper.getNewSILFunctionType(genEnv, funcType, Mod);
     SILType newSILType = SILType::getPrimitiveObjectType(newFnType);
     if (nonOptionalType.isAddress()) {

test/IRGen/big_types_corner_cases.swift

@@ -288,17 +288,41 @@ public protocol sr8076_Query {
 }
 
 public protocol sr8076_ProtoQueryHandler {
-    func forceHandle<Q: sr8076_Query>(query: Q) throws -> (Q.Returned, sr8076_Filter?)
+    func forceHandle_1<Q: sr8076_Query>(query: Q) -> Void
+    func forceHandle_2<Q: sr8076_Query>(query: Q) -> (Q.Returned, sr8076_BigStruct?)
+    func forceHandle_3<Q: sr8076_Query>(query: Q) -> (Q.Returned, sr8076_Filter?)
+    func forceHandle_4<Q: sr8076_Query>(query: Q) throws -> (Q.Returned, sr8076_Filter?)
 }
 
 public protocol sr8076_QueryHandler: sr8076_ProtoQueryHandler {
     associatedtype Handled: sr8076_Query
-    func handle(query: Handled) throws -> (Handled.Returned, sr8076_Filter?)
+    func handle_1(query: Handled) -> Void
+    func handle_2(query: Handled) -> (Handled.Returned, sr8076_BigStruct?)
+    func handle_3(query: Handled) -> (Handled.Returned, sr8076_Filter?)
+    func handle_4(query: Handled) throws -> (Handled.Returned, sr8076_Filter?)
 }
 
 public extension sr8076_QueryHandler {
-    func forceHandle<Q: sr8076_Query>(query: Q) throws -> (Q.Returned, sr8076_Filter?) {
-        guard let body = handle as? (Q) throws -> (Q.Returned, sr8076_Filter?) else {
+    func forceHandle_4<Q: sr8076_Query>(query: Q) -> Void {
+        guard let body = handle_1 as? (Q) -> Void else {
+            fatalError("handler \(self) is expected to handle query \(query)")
+        }
+        body(query)
+    }
+    func forceHandle_4<Q: sr8076_Query>(query: Q) -> (Q.Returned, sr8076_BigStruct?) {
+        guard let body = handle_2 as? (Q) -> (Q.Returned, sr8076_BigStruct?) else {
+            fatalError("handler \(self) is expected to handle query \(query)")
+        }
+        return body(query)
+    }
+    func forceHandle_3<Q: sr8076_Query>(query: Q) -> (Q.Returned, sr8076_Filter?) {
+        guard let body = handle_3 as? (Q) -> (Q.Returned, sr8076_Filter?) else {
+            fatalError("handler \(self) is expected to handle query \(query)")
+        }
+        return body(query)
+    }
+    func forceHandle_4<Q: sr8076_Query>(query: Q) throws -> (Q.Returned, sr8076_Filter?) {
+        guard let body = handle_4 as? (Q) throws -> (Q.Returned, sr8076_Filter?) else {
             fatalError("handler \(self) is expected to handle query \(query)")
         }
         return try body(query)