IRGen support for GuaranteedAddress result convention

Author: meg-gupta

lib/IRGen/CallEmission.h

@@ -103,6 +103,7 @@ class CallEmission {
   virtual void emitCallToUnmappedExplosion(llvm::CallBase *call,
                                            Explosion &out) = 0;
   void emitYieldsToExplosion(Explosion &out);
+  void emitGuaranteedAddressToExplosion(Explosion &out);
   void setKeyPathAccessorArguments(Explosion &in, bool isOutlined,
                                    Explosion &out);
   virtual FunctionPointer getCalleeFunctionPointer() = 0;

lib/IRGen/GenCall.cpp

@@ -591,6 +591,10 @@ namespace {
     /// function type.
     void expandCoroutineContinuationType();
 
+    /// Initializes the result type for functions with @guaranteed_addr return
+    /// convention.
+    void expandGuaranteedAddressResult();
+
     // Expand the components for the async continuation entrypoint of the
     // function type (the function to be called on returning).
     void expandAsyncReturnType();
@@ -694,6 +698,10 @@ void SignatureExpansion::expandResult(
 
   auto fnConv = getSILFuncConventions();
 
+  if (fnConv.hasGuaranteedAddressResult()) {
+    return expandGuaranteedAddressResult();
+  }
+
   // Disable the use of sret if we have multiple indirect results.
   if (fnConv.getNumIndirectSILResults() > 1)
     CanUseSRet = false;
@@ -911,6 +919,11 @@ void SignatureExpansion::expandCoroutineContinuationParameters() {
   }
 }
 
+void SignatureExpansion::expandGuaranteedAddressResult() {
+  CanUseSRet = false;
+  ResultIRType = IGM.PtrTy;
+}
+
 void SignatureExpansion::addAsyncParameters() {
   // using TaskContinuationFunction =
   //   SWIFT_CC(swift)
@@ -3933,6 +3946,11 @@ void CallEmission::emitYieldsToExplosion(Explosion &out) {
   }
 }
 
+void CallEmission::emitGuaranteedAddressToExplosion(Explosion &out) {
+  auto call = emitCallSite();
+  out.add(call);
+}
+
 /// Emit the result of this call to an explosion.
 void CallEmission::emitToExplosion(Explosion &out, bool isOutlined) {
   assert(state == State::Emitting);
@@ -3948,6 +3966,14 @@ void CallEmission::emitToExplosion(Explosion &out, bool isOutlined) {
 
   SILFunctionConventions fnConv(getCallee().getSubstFunctionType(),
                                 IGF.getSILModule());
+
+  if (fnConv.hasGuaranteedAddressResult()) {
+    assert(LastArgWritten == 0 &&
+           "@guaranteed_addr along with indirect result?");
+    emitGuaranteedAddressToExplosion(out);
+    return;
+  }
+
   SILType substResultType =
       fnConv.getSILResultType(IGF.IGM.getMaximalTypeExpansionContext());
 
@@ -7001,6 +7027,16 @@ void irgen::emitYieldOnceCoroutineResult(IRGenFunction &IGF, Explosion &result,
   }
 }
 
+void irgen::emitGuaranteedAddressResult(IRGenFunction &IGF, Explosion &result,
+                                        SILType funcResultType,
+                                        SILType returnResultType) {
+  assert(funcResultType == returnResultType);
+  assert(funcResultType.isAddress());
+  auto &Builder = IGF.Builder;
+  Builder.CreateRet(result.claimNext());
+  assert(result.empty());
+}
+
 FunctionPointer
 IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) {
   auto *fnTy = llvm::FunctionType::get(

lib/IRGen/GenCall.h

@@ -281,6 +281,10 @@ namespace irgen {
   void emitYieldOnceCoroutineResult(IRGenFunction &IGF, Explosion &result,
                                     SILType funcResultType, SILType returnResultType);
 
+  void emitGuaranteedAddressResult(IRGenFunction &IGF, Explosion &result,
+                                   SILType funcResultType,
+                                   SILType returnResultType);
+
   Address emitAutoDiffCreateLinearMapContextWithType(
       IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype);
 

lib/IRGen/IRGenSIL.cpp

@@ -2032,6 +2032,11 @@ static ArrayRef<SILArgument *> emitEntryPointIndirectReturn(
   SILFunctionConventions fnConv(funcTy, IGF.getSILModule());
   SILType directResultType = IGF.CurSILFn->mapTypeIntoContext(
       fnConv.getSILResultType(IGF.IGM.getMaximalTypeExpansionContext()));
+
+  if (fnConv.hasGuaranteedAddressResult()) {
+    return entry->getArguments();
+  }
+
   if (requiresIndirectResult(directResultType)) {
     auto &paramTI = IGF.IGM.getTypeInfo(directResultType);
     auto &retTI =
@@ -3934,7 +3939,11 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
 
   // For a simple apply, just bind the apply result to the result of the call.
   if (auto apply = dyn_cast<ApplyInst>(i)) {
-    setLoweredExplosion(apply, result);
+    if (apply->hasGuaranteedAddressResult()) {
+      setCorrespondingLoweredValues(apply->getResults(), result);
+    } else {
+      setLoweredExplosion(apply, result);
+    }
     emission->end();
 
   // For begin_apply, we have to destructure the call.
@@ -4451,6 +4460,16 @@ static void emitReturnInst(IRGenSILFunction &IGF,
     return;
   }
 
+  if (conv.hasGuaranteedAddressResult()) {
+    assert(IGF.CurSILFn->getLoweredFunctionType()->getLanguage() ==
+           SILFunctionLanguage::Swift);
+    auto funcResultType = IGF.CurSILFn->mapTypeIntoContext(
+        conv.getSILResultType(IGF.IGM.getMaximalTypeExpansionContext()));
+
+    emitGuaranteedAddressResult(IGF, result, funcResultType, resultTy);
+    return;
+  }
+
   // The invariant on the out-parameter is that it's always zeroed, so
   // there's nothing to do here.