[IRGen] Added fn ptr to async FunctionPointer.

The address of the function to be called when generating code to invoke
the function associated with FunctionPointer which is produced via
direct reference is by definition statically known; it is neither necessary
nor desireable to load this address out of the AsyncFunctionPointer
corresponding to the function.

Here, that spurious additional work is skipped.  The approach is to add
a second value to the FunctionPointer struct.  For FunctionPointers
whose kind is AsyncFunctionPointer, this value is either null or else
the address of the corresponding function.

rdar://71376092

Author: nate-chandler

lib/IRGen/Callee.h

@@ -234,6 +234,12 @@ namespace irgen {
     /// The actual pointer, either to the function or to its descriptor.
     llvm::Value *Value;
 
+    /// An additional value whose meaning varies by the FunctionPointer's Kind:
+    /// - Kind::AsyncFunctionPointer -> pointer to the corresponding function
+    ///                                 if the FunctionPointer was created via
+    ///                                 forDirect; nullptr otherwise. 
+    llvm::Value *SecondaryValue;
+
     PointerAuthInfo AuthInfo;
 
     Signature Sig;
@@ -243,9 +249,11 @@ namespace irgen {
     /// We may add more arguments to this; try to use the other
     /// constructors/factories if possible.
     explicit FunctionPointer(Kind kind, llvm::Value *value,
+                             llvm::Value *secondaryValue,
                              PointerAuthInfo authInfo,
                              const Signature &signature)
-        : kind(kind), Value(value), AuthInfo(authInfo), Sig(signature) {
+        : kind(kind), Value(value), SecondaryValue(secondaryValue),
+          AuthInfo(authInfo), Sig(signature) {
       // The function pointer should have function type.
       assert(value->getType()->getPointerElementType()->isFunctionTy());
       // TODO: maybe assert similarity to signature.getType()?
@@ -258,18 +266,25 @@ namespace irgen {
       }
     }
 
+    explicit FunctionPointer(Kind kind, llvm::Value *value,
+                             PointerAuthInfo authInfo,
+                             const Signature &signature)
+        : FunctionPointer(kind, value, nullptr, authInfo, signature){};
+
     // Temporary only!
     explicit FunctionPointer(Kind kind, llvm::Value *value,
                              const Signature &signature)
       : FunctionPointer(kind, value, PointerAuthInfo(), signature) {}
 
-    static FunctionPointer forDirect(IRGenModule &IGM,
-                                     llvm::Constant *value,
+    static FunctionPointer forDirect(IRGenModule &IGM, llvm::Constant *value,
+                                     llvm::Constant *secondaryValue,
                                      CanSILFunctionType fnType);
 
     static FunctionPointer forDirect(Kind kind, llvm::Constant *value,
+                                     llvm::Constant *secondaryValue,
                                      const Signature &signature) {
-      return FunctionPointer(kind, value, PointerAuthInfo(), signature);
+      return FunctionPointer(kind, value, secondaryValue, PointerAuthInfo(),
+                             signature);
     }
 
     static FunctionPointer forExplosionValue(IRGenFunction &IGF,
@@ -295,6 +310,13 @@ namespace irgen {
     /// Return the actual function pointer.
     llvm::Value *getRawPointer() const { return Value; }
 
+    /// Assuming that the receiver is of kind AsyncFunctionPointer, returns the
+    /// pointer to the corresponding function if available.
+    llvm::Value *getRawAsyncFunction() const {
+      assert(kind.isAsyncFunctionPointer());
+      return SecondaryValue;
+    }
+
     /// Given that this value is known to have been constructed from
     /// a direct function, return the function pointer.
     llvm::Constant *getDirectPointer() const {

lib/IRGen/GenCall.cpp

@@ -1985,18 +1985,33 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
   assert(functionPointer.getKind() != FunctionPointer::Kind::Function);
   bool emitFunction = values.first;
   bool emitSize = values.second;
-  auto *ptr = functionPointer.getRawPointer();
-  if (auto authInfo = functionPointer.getAuthInfo()) {
-    ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
-  }
-  auto *afpPtr =
-      IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
+  // Ensure that the AsyncFunctionPointer is not auth'd if it is not used and
+  // that it is not auth'd more than once if it is needed.
+  //
+  // The AsyncFunctionPointer is not needed in the case where only the function
+  // is being loaded and the FunctionPointer was created from a function_ref
+  // instruction.
+  llvm::Optional<llvm::Value *> afpPtrValue = llvm::None;
+  auto getAFPPtr = [&]() {
+    if (!afpPtrValue) {
+      auto *ptr = functionPointer.getRawPointer();
+      if (auto authInfo = functionPointer.getAuthInfo()) {
+        ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
+      }
+      auto *afpPtr =
+          IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
+      afpPtrValue = afpPtr;
+    }
+    return *afpPtrValue;
+  };
   llvm::Value *fn = nullptr;
   if (emitFunction) {
     if (functionPointer.useStaticContextSize()) {
       fn = functionPointer.getRawPointer();
+    } else if (auto *function = functionPointer.getRawAsyncFunction()) {
+      fn = function;
     } else {
-      llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
+      llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(getAFPPtr(), 0);
       fn = IGF.emitLoadOfRelativePointer(
           Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
           /*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
@@ -2014,7 +2029,7 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
                                     initialContextSize.getValue());
     } else {
       assert(!functionPointer.useStaticContextSize());
-      auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
+      auto *sizePtr = IGF.Builder.CreateStructGEP(getAFPPtr(), 1);
       size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
     }
   }
@@ -4667,8 +4682,9 @@ Callee irgen::getCFunctionPointerCallee(IRGenFunction &IGF,
 
 FunctionPointer FunctionPointer::forDirect(IRGenModule &IGM,
                                            llvm::Constant *fnPtr,
+                                           llvm::Constant *secondaryValue,
                                            CanSILFunctionType fnType) {
-  return forDirect(fnType, fnPtr, IGM.getSignature(fnType));
+  return forDirect(fnType, fnPtr, secondaryValue, IGM.getSignature(fnType));
 }
 
 StringRef FunctionPointer::getName(IRGenModule &IGM) const {
@@ -4690,6 +4706,14 @@ llvm::Value *FunctionPointer::getPointer(IRGenFunction &IGF) const {
   case BasicKind::Function:
     return Value;
   case BasicKind::AsyncFunctionPointer: {
+    if (auto *rawFunction = getRawAsyncFunction()) {
+      // If the pointer to the underlying function is available, it means that
+      // this FunctionPointer instance was created via 
+      // FunctionPointer::forDirect and as such has no AuthInfo.
+      assert(!AuthInfo && "have PointerAuthInfo for an async FunctionPointer "
+                          "for which the raw function is known");
+      return rawFunction;
+    }
     auto *fnPtr = Value;
     if (auto authInfo = AuthInfo) {
       fnPtr = emitPointerAuthAuth(IGF, fnPtr, authInfo);

lib/IRGen/GenClass.cpp

@@ -2653,7 +2653,15 @@ FunctionPointer irgen::emitVirtualMethodValue(IRGenFunction &IGF,
   case ClassMetadataLayout::MethodInfo::Kind::DirectImpl: {
     auto fnPtr = llvm::ConstantExpr::getBitCast(methodInfo.getDirectImpl(),
                                            signature.getType()->getPointerTo());
-    return FunctionPointer::forDirect(methodType, fnPtr, signature);
+    llvm::Constant *secondaryValue = nullptr;
+    if (cast<AbstractFunctionDecl>(method.getDecl())->hasAsync()) {
+      auto *silFn = IGF.IGM.getSILFunctionForAsyncFunctionPointer(
+          methodInfo.getDirectImpl());
+      secondaryValue = cast<llvm::Constant>(
+          IGF.IGM.getAddrOfSILFunction(silFn, NotForDefinition));
+    }
+    return FunctionPointer::forDirect(methodType, fnPtr, secondaryValue,
+                                      signature);
   }
   }
 

lib/IRGen/GenKeyPath.cpp

@@ -252,8 +252,9 @@ getAccessorForComputedComponent(IRGenModule &IGM,
                                &ignoreWitnessMetadata,
                                forwardedArgs);
     }
-    auto fnPtr = FunctionPointer::forDirect(IGM, accessorFn,
-                                            accessor->getLoweredFunctionType());
+    auto fnPtr =
+        FunctionPointer::forDirect(IGM, accessorFn, /*secondaryValue*/ nullptr,
+                                   accessor->getLoweredFunctionType());
     auto call = IGF.Builder.CreateCall(fnPtr, forwardedArgs.claimAll());
 
     if (call->getType()->isVoidTy())

lib/IRGen/GenObjC.cpp

@@ -658,8 +658,9 @@ Callee irgen::getObjCMethodCallee(IRGenFunction &IGF,
   Selector selector(method);
   llvm::Value *selectorValue = IGF.emitObjCSelectorRefLoad(selector.str());
 
-  auto fn = FunctionPointer::forDirect(FunctionPointer::Kind::Function,
-                                       messenger, sig);
+  auto fn =
+      FunctionPointer::forDirect(FunctionPointer::Kind::Function, messenger,
+                                 /*secondaryValue*/ nullptr, sig);
   return Callee(std::move(info), fn, receiverValue, selectorValue);
 }
 

lib/IRGen/IRGenSIL.cpp

@@ -2526,13 +2526,17 @@ void IRGenSILFunction::visitFunctionRefBaseInst(FunctionRefBaseInst *i) {
       fn, NotForDefinition, false /*isDynamicallyReplaceableImplementation*/,
       isa<PreviousDynamicFunctionRefInst>(i));
   llvm::Constant *value;
+  llvm::Constant *secondaryValue;
   if (fpKind.isAsyncFunctionPointer()) {
     value = IGM.getAddrOfAsyncFunctionPointer(fn);
     value = llvm::ConstantExpr::getBitCast(value, fnPtr->getType());
+    secondaryValue = IGM.getAddrOfSILFunction(fn, NotForDefinition);
   } else {
     value = fnPtr;
+    secondaryValue = nullptr;
   }
-  FunctionPointer fp = FunctionPointer(fpKind, value, sig);
+  FunctionPointer fp =
+      FunctionPointer::forDirect(fpKind, value, secondaryValue, sig);
 
   // Store the function as a FunctionPointer so we can avoid bitcasting
   // or thunking if we don't need to.
@@ -6525,16 +6529,18 @@ void IRGenSILFunction::visitWitnessMethodInst(swift::WitnessMethodInst *i) {
   if (IGM.isResilient(conformance.getRequirement(),
                       ResilienceExpansion::Maximal)) {
     llvm::Constant *fnPtr = IGM.getAddrOfDispatchThunk(member, NotForDefinition);
+    llvm::Constant *secondaryValue = nullptr;
 
     if (fnType->isAsync()) {
+      secondaryValue = fnPtr;
       auto *fnPtrType = fnPtr->getType();
       fnPtr = IGM.getAddrOfAsyncFunctionPointer(
           LinkEntity::forDispatchThunk(member));
       fnPtr = llvm::ConstantExpr::getBitCast(fnPtr, fnPtrType);
     }
 
     auto sig = IGM.getSignature(fnType);
-    auto fn = FunctionPointer::forDirect(fnType, fnPtr, sig);
+    auto fn = FunctionPointer::forDirect(fnType, fnPtr, secondaryValue, sig);
 
     setLoweredFunctionPointer(i, fn);
     return;