Large loadable types pass: resolves a case wherein a tuple’s extracted type is a function signature

Author: Joe Shajrawi

lib/IRGen/LoadableByAddress.cpp

@@ -307,6 +307,8 @@ struct StructLoweringState {
   SmallVector<SwitchEnumInst *, 16> switchEnumInstsToMod;
   // All struct_extract instrs that should be converted to struct_element_addr
   SmallVector<StructExtractInst *, 16> structExtractInstsToMod;
+  // All tuple instructions for which the return type is a function type
+  SmallVector<SILInstruction *, 8> tupleInstsToMod;
   // All Retain and release instrs should be replaced with _addr version
   SmallVector<RetainValueInst *, 16> retainInstsToMod;
   SmallVector<ReleaseValueInst *, 16> releaseInstsToMod;
@@ -350,6 +352,7 @@ class LargeValueVisitor {
   void visitReleaseInst(ReleaseValueInst *instr);
   void visitResultTyInst(SILInstruction *instr);
   void visitDebugValueInst(DebugValueInst *instr);
+  void visitTupleInst(SILInstruction *instr);
   void visitInstr(SILInstruction *instr);
 };
 } // end anonymous namespace
@@ -411,6 +414,11 @@ void LargeValueVisitor::mapValueStorage() {
         visitSwitchEnumInst(SEI);
         break;
       }
+      case ValueKind::TupleElementAddrInst:
+      case ValueKind::TupleExtractInst: {
+        visitTupleInst(currIns);
+        break;
+      }
       default: {
         assert(!ApplySite::isa(currIns) && "Did not expect an ApplySite");
         assert(!dyn_cast<MethodInst>(currIns) && "Unhandled Method Inst");
@@ -607,6 +615,24 @@ void LargeValueVisitor::visitResultTyInst(SILInstruction *instr) {
   }
 }
 
+void LargeValueVisitor::visitTupleInst(SILInstruction *instr) {
+  SILType currSILType = instr->getType().getObjectType();
+  CanType currCanType = currSILType.getSwiftRValueType();
+  if (auto funcType = dyn_cast<SILFunctionType>(currCanType)) {
+    CanSILFunctionType canFuncType = CanSILFunctionType(funcType);
+    GenericEnvironment *genEnv = instr->getFunction()->getGenericEnvironment();
+    if (!genEnv && canFuncType->isPolymorphic()) {
+      genEnv = getGenericEnvironment(instr->getModule(), canFuncType);
+    }
+    SILFunctionType *newSILFunctionType =
+        getNewSILFunctionTypePtr(genEnv, funcType, pass.Mod);
+    if (funcType != newSILFunctionType) {
+      pass.tupleInstsToMod.push_back(instr);
+    }
+  }
+  visitInstr(instr);
+}
+
 void LargeValueVisitor::visitInstr(SILInstruction *instr) {
   for (Operand &operand : instr->getAllOperands()) {
     if (std::find(pass.largeLoadableArgs.begin(), pass.largeLoadableArgs.end(),
@@ -1208,6 +1234,40 @@ static bool allUsesAreReplaceable(SILInstruction *instr,
   return allUsesAreReplaceable;
 }
 
+static void castTupleInstr(SILInstruction *instr, IRGenModule &Mod) {
+  SILType currSILType = instr->getType().getObjectType();
+  CanType currCanType = currSILType.getSwiftRValueType();
+  SILFunctionType *funcType = dyn_cast<SILFunctionType>(currCanType);
+  assert(funcType && "Expcted SILFunctionType as tuple's return");
+  CanSILFunctionType canFuncType = CanSILFunctionType(funcType);
+  GenericEnvironment *genEnv = instr->getFunction()->getGenericEnvironment();
+  if (!genEnv && canFuncType->isPolymorphic()) {
+    genEnv = getGenericEnvironment(instr->getModule(), canFuncType);
+  }
+  SILType newSILType = getNewSILFunctionType(genEnv, funcType, Mod);
+  auto II = instr->getIterator();
+  ++II;
+  SILBuilder castBuilder(II);
+  SILInstruction *castInstr = nullptr;
+  switch (instr->getKind()) {
+  // Add cast to the new sil function type:
+  case ValueKind::TupleExtractInst: {
+    castInstr = castBuilder.createUncheckedBitCast(instr->getLoc(), instr,
+                                                   newSILType.getObjectType());
+    break;
+  }
+  case ValueKind::TupleElementAddrInst: {
+    castInstr = castBuilder.createUncheckedAddrCast(
+        instr->getLoc(), instr, newSILType.getAddressType());
+    break;
+  }
+  default:
+    llvm_unreachable("Unexpected instruction inside tupleInstsToMod");
+  }
+  instr->replaceAllUsesWith(castInstr);
+  castInstr->setOperand(0, instr);
+}
+
 static void rewriteFunction(StructLoweringState &pass,
                             LoadableStorageAllocation &allocator) {
 
@@ -1350,6 +1410,10 @@ static void rewriteFunction(StructLoweringState &pass,
     }
   }
 
+  for (SILInstruction *instr : pass.tupleInstsToMod) {
+    castTupleInstr(instr, pass.Mod);
+  }
+
   for (SILInstruction *instr : pass.debugInstsToMod) {
     assert(instr->getAllOperands().size() == 1 &&
            "Debug instructions have one operand");

lib/SIL/SILVerifier.cpp

@@ -1681,8 +1681,10 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
               "member variables of struct");
 
       SILType loweredType = structTy.getFieldType(field, F.getModule());
-      require((*opi)->getType() == loweredType,
-              "struct operand type does not match field type");
+      if (SI->getModule().getStage() != SILStage::Lowered) {
+        require((*opi)->getType() == loweredType,
+                "struct operand type does not match field type");
+      }
       ++opi;
     }
   }
@@ -1781,10 +1783,12 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     require(TI->getElements().size() == ResTy->getNumElements(),
             "Tuple field count mismatch!");
 
-    for (size_t i = 0, size = TI->getElements().size(); i < size; ++i) {
-      require(TI->getElement(i)->getType().getSwiftRValueType()
-              == ResTy.getElementType(i),
-              "Tuple element arguments do not match tuple type!");
+    if (TI->getModule().getStage() != SILStage::Lowered) {
+      for (size_t i = 0, size = TI->getElements().size(); i < size; ++i) {
+        require(TI->getElement(i)->getType().getSwiftRValueType() ==
+                    ResTy.getElementType(i),
+                "Tuple element arguments do not match tuple type!");
+      }
     }
   }
 
@@ -2009,9 +2013,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
 
     require(EI->getFieldNo() < operandTy->getNumElements(),
             "invalid field index for tuple_extract instruction");
-    require(EI->getType().getSwiftRValueType()
-            == operandTy.getElementType(EI->getFieldNo()),
-            "type of tuple_extract does not match type of element");
+    if (EI->getModule().getStage() != SILStage::Lowered) {
+      require(EI->getType().getSwiftRValueType() ==
+                  operandTy.getElementType(EI->getFieldNo()),
+              "type of tuple_extract does not match type of element");
+    }
   }
 
   void checkStructExtractInst(StructExtractInst *EI) {
@@ -2030,10 +2036,12 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     require(EI->getField()->getDeclContext() == sd,
             "struct_extract field is not a member of the struct");
 
-    SILType loweredFieldTy = operandTy.getFieldType(EI->getField(),
-                                                    F.getModule());
-    require(loweredFieldTy == EI->getType(),
-            "result of struct_extract does not match type of field");
+    if (EI->getModule().getStage() != SILStage::Lowered) {
+      SILType loweredFieldTy =
+          operandTy.getFieldType(EI->getField(), F.getModule());
+      require(loweredFieldTy == EI->getType(),
+              "result of struct_extract does not match type of field");
+    }
   }
 
   void checkTupleElementAddrInst(TupleElementAddrInst *EI) {
@@ -2048,9 +2056,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     ArrayRef<TupleTypeElt> fields = operandTy.castTo<TupleType>()->getElements();
     require(EI->getFieldNo() < fields.size(),
             "invalid field index for element_addr instruction");
-    require(EI->getType().getSwiftRValueType()
-              == CanType(fields[EI->getFieldNo()].getType()),
-            "type of tuple_element_addr does not match type of element");
+    if (EI->getModule().getStage() != SILStage::Lowered) {
+      require(EI->getType().getSwiftRValueType() ==
+                  CanType(fields[EI->getFieldNo()].getType()),
+              "type of tuple_element_addr does not match type of element");
+    }
   }
 
   void checkStructElementAddrInst(StructElementAddrInst *EI) {
@@ -2069,10 +2079,12 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     require(EI->getField()->getDeclContext() == sd,
             "struct_element_addr field is not a member of the struct");
 
-    SILType loweredFieldTy = operandTy.getFieldType(EI->getField(),
-                                                    F.getModule());
-    require(loweredFieldTy == EI->getType(),
-            "result of struct_element_addr does not match type of field");
+    if (EI->getModule().getStage() != SILStage::Lowered) {
+      SILType loweredFieldTy =
+          operandTy.getFieldType(EI->getField(), F.getModule());
+      require(loweredFieldTy == EI->getType(),
+              "result of struct_element_addr does not match type of field");
+    }
   }
 
   void checkRefElementAddrInst(RefElementAddrInst *EI) {

test/IRGen/big_types_corner_cases.swift

@@ -32,6 +32,20 @@ public func f3_uses_f2() {
 // CHECK: call swiftcc void %16(%T22big_types_corner_cases9BigStructV* noalias nocapture sret %call.aggresult1, %T22big_types_corner_cases9BigStructV* noalias nocapture dereferenceable
 // CHECK: ret void
 
+public func f4_tuple_use_of_f2() {
+  let x = BigStruct()
+  let tupleWithFunc = (f2_returns_f1(), x)
+  let useOfF2 = tupleWithFunc.0
+  let _ = useOfF2(x)
+}
+
+// CHECK-LABEL: define{{( protected)?}} swiftcc void @_T022big_types_corner_cases18f4_tuple_use_of_f2yyF()
+// CHECK: [[TUPLE:%.*]] = call swiftcc { i8*, %swift.refcounted* } @_T022big_types_corner_cases13f2_returns_f1AA9BigStructVADcyF()
+// CHECK: [[TUPLE_EXTRACT:%.*]] = extractvalue { i8*, %swift.refcounted* } [[TUPLE]], 0
+// CHECK: [[CAST_EXTRACT:%.*]] = bitcast i8* [[TUPLE_EXTRACT]] to void (%T22big_types_corner_cases9BigStructV*, %T22big_types_corner_cases9BigStructV*, %swift.refcounted*)*
+// CHECK:  call swiftcc void [[CAST_EXTRACT]](%T22big_types_corner_cases9BigStructV* noalias nocapture sret %call.aggresult1, %T22big_types_corner_cases9BigStructV* noalias nocapture dereferenceable
+// CHECK: ret void
+
 public struct MyStruct {
   public let a: Int
   public let b: String?