Merge pull request #35611 from slavapestov/emit-dispatch-thunk-async-fp

IRGen: Emit async function pointers for resilient class and protocol dispatch thunks

Author: slavapestov

include/swift/IRGen/Linking.h

@@ -130,10 +130,24 @@ class LinkEntity {
     /// is a ConstructorDecl* inside a class.
     DispatchThunkInitializer,
 
-    /// A method dispatch thunk for an allocating constructor.  The pointer is a
-    /// ConstructorDecl* inside a protocol or a class.
+    /// A method dispatch thunk for an allocating constructor.  The pointer is
+    /// a ConstructorDecl* inside a protocol or a class.
     DispatchThunkAllocator,
 
+    /// An async function pointer for a method dispatch thunk.  The pointer is
+    /// a FuncDecl* inside a protocol or a class.
+    DispatchThunkAsyncFunctionPointer,
+
+    /// An async function pointer for a method dispatch thunk for an
+    /// initializing constructor.  The pointer is a ConstructorDecl* inside a
+    /// class.
+    DispatchThunkInitializerAsyncFunctionPointer,
+
+    /// An async function pointer for a method dispatch thunk for an allocating
+    /// constructor.  The pointer is a ConstructorDecl* inside a protocol or
+    /// a class.
+    DispatchThunkAllocatorAsyncFunctionPointer,
+
     /// A method descriptor.  The pointer is a FuncDecl* inside a protocol
     /// or a class.
     MethodDescriptor,
@@ -295,7 +309,7 @@ class LinkEntity {
 
     /// The same as AsyncFunctionPointer but with a different stored value, for
     /// use by TBDGen.
-    /// The pointer is a AbstractStorageDecl*.
+    /// The pointer is an AbstractFunctionDecl*.
     AsyncFunctionPointerAST,
 
     /// The pointer is a SILFunction*.
@@ -861,7 +875,8 @@ class LinkEntity {
   }
 
   static LinkEntity
-  forSILFunction(SILFunction *F, bool IsDynamicallyReplaceableImplementation) {
+  forSILFunction(SILFunction *F,
+                 bool IsDynamicallyReplaceableImplementation=false) {
     LinkEntity entity;
     entity.Pointer = F;
     entity.SecondaryPointer = nullptr;
@@ -1100,12 +1115,36 @@ class LinkEntity {
     return entity;
   }
 
-  static LinkEntity forAsyncFunctionPointer(SILFunction *silFunction) {
+  static LinkEntity forAsyncFunctionPointer(LinkEntity other) {
     LinkEntity entity;
-    entity.Pointer = silFunction;
+    entity.Pointer = other.Pointer;
     entity.SecondaryPointer = nullptr;
-    entity.Data = LINKENTITY_SET_FIELD(
-        Kind, unsigned(LinkEntity::Kind::AsyncFunctionPointer));
+
+    switch (other.getKind()) {
+    case LinkEntity::Kind::SILFunction:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::AsyncFunctionPointer));
+      break;
+
+    case LinkEntity::Kind::DispatchThunk:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::DispatchThunkAsyncFunctionPointer));
+      break;
+
+    case LinkEntity::Kind::DispatchThunkInitializer:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::DispatchThunkInitializerAsyncFunctionPointer));
+      break;
+
+    case LinkEntity::Kind::DispatchThunkAllocator:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::DispatchThunkAllocatorAsyncFunctionPointer));
+      break;
+
+    default:
+      llvm_unreachable("Link entity kind cannot have an async function pointer");
+    }
+
     return entity;
   }
 
@@ -1115,6 +1154,39 @@ class LinkEntity {
     return entity;
   }
 
+  LinkEntity getUnderlyingEntityForAsyncFunctionPointer() const {
+    LinkEntity entity;
+    entity.Pointer = Pointer;
+    entity.SecondaryPointer = nullptr;
+
+    switch (getKind()) {
+    case LinkEntity::Kind::AsyncFunctionPointer:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::SILFunction));
+      break;
+
+    case LinkEntity::Kind::DispatchThunkAsyncFunctionPointer:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::DispatchThunk));
+      break;
+
+    case LinkEntity::Kind::DispatchThunkInitializerAsyncFunctionPointer:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::DispatchThunkInitializer));
+      break;
+
+    case LinkEntity::Kind::DispatchThunkAllocatorAsyncFunctionPointer:
+      entity.Data = LINKENTITY_SET_FIELD(
+          Kind, unsigned(LinkEntity::Kind::DispatchThunkAllocator));
+      break;
+
+    default:
+      llvm_unreachable("Link entity is not an async function pointer");
+    }
+
+    return entity;
+  }
+
   void mangle(llvm::raw_ostream &out) const;
   void mangle(SmallVectorImpl<char> &buffer) const;
   std::string mangleAsString() const;

lib/IRGen/Callee.h

@@ -174,8 +174,8 @@ namespace irgen {
 
     // Temporary only!
     explicit FunctionPointer(KindTy kind, llvm::Value *value,
-                             const Signature &signature, bool
-                             isWithoutCtxt = false)
+                             const Signature &signature,
+                             bool isWithoutCtxt = false)
         : FunctionPointer(kind, value, PointerAuthInfo(), signature, isWithoutCtxt) {}
 
     static FunctionPointer forDirect(IRGenModule &IGM,

lib/IRGen/GenCall.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/Runtime/Config.h"
+#include "swift/IRGen/Linking.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/SILType.h"
 #include "clang/AST/ASTContext.h"
@@ -379,7 +380,7 @@ llvm::CallInst *IRGenFunction::emitSuspendAsyncCall(ArrayRef<llvm::Value *> args
   auto *calleeContext = Builder.CreateExtractValue(id,
       (unsigned)AsyncFunctionArgumentIndex::Context);
   llvm::Constant *projectFn = cast<llvm::Constant>(args[1])->stripPointerCasts();
-  // Get the caller context from the calle context.
+  // Get the caller context from the callee context.
   llvm::Value *context = Builder.CreateCall(projectFn, {calleeContext});
   context = Builder.CreateBitCast(context, IGM.SwiftContextPtrTy);
   Builder.CreateStore(context, asyncContextLocation);
@@ -3586,9 +3587,10 @@ emitRetconCoroutineEntry(IRGenFunction &IGF, CanSILFunctionType fnType,
 }
 
 void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
-                                   SILFunction *asyncFunction) {
+                                   const AsyncContextLayout &layout,
+                                   LinkEntity asyncFunction) {
   auto &IGM = IGF.IGM;
-  auto size = getAsyncContextLayout(IGM, asyncFunction).getSize();
+  auto size = layout.getSize();
   auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
       IGM.getAddrOfAsyncFunctionPointer(asyncFunction), IGM.Int8PtrTy);
   auto *id = IGF.Builder.CreateIntrinsicCall(

lib/IRGen/GenCall.h

@@ -414,7 +414,9 @@ namespace irgen {
   Address emitAllocAsyncContext(IRGenFunction &IGF, llvm::Value *sizeValue);
   void emitDeallocAsyncContext(IRGenFunction &IGF, Address context);
 
-  void emitAsyncFunctionEntry(IRGenFunction &IGF, SILFunction *asyncFunc);
+  void emitAsyncFunctionEntry(IRGenFunction &IGF,
+                              const AsyncContextLayout &layout,
+                              LinkEntity asyncFunction);
 
   /// Yield the given values from the current continuation.
   ///

lib/IRGen/GenKeyPath.cpp

@@ -920,7 +920,7 @@ emitKeyPathComponent(IRGenModule &IGM,
       idKind = KeyPathComponentHeader::Pointer;
       // FIXME: Does this need to be signed?
       auto idRef = IGM.getAddrOfLLVMVariableOrGOTEquivalent(
-        LinkEntity::forSILFunction(id.getFunction(), false));
+        LinkEntity::forSILFunction(id.getFunction()));
 
       idValue = idRef.getValue();
       // If we got an indirect reference, we'll need to resolve it at

lib/IRGen/GenMeta.cpp

@@ -5353,14 +5353,14 @@ bool irgen::methodRequiresReifiedVTableEntry(IRGenModule &IGM,
 }
 
 llvm::GlobalValue *irgen::emitAsyncFunctionPointer(IRGenModule &IGM,
-                                                   SILFunction *function,
+                                                   llvm::Function *function,
+                                                   LinkEntity entity,
                                                    Size size) {
   ConstantInitBuilder initBuilder(IGM);
   ConstantStructBuilder builder(
       initBuilder.beginStruct(IGM.AsyncFunctionPointerTy));
-  builder.addRelativeAddress(
-      IGM.getAddrOfSILFunction(function, NotForDefinition));
+  builder.addRelativeAddress(function);
   builder.addInt32(size.getValue());
   return cast<llvm::GlobalValue>(IGM.defineAsyncFunctionPointer(
-      function, builder.finishAndCreateFuture()));
+      entity, builder.finishAndCreateFuture()));
 }

lib/IRGen/GenMeta.h

@@ -47,6 +47,7 @@ namespace irgen {
   class Size;
   class StructLayout;
   class ClassLayout;
+  class LinkEntity;
 
   bool requiresForeignTypeMetadata(CanType type);
   bool requiresForeignTypeMetadata(NominalTypeDecl *decl);
@@ -183,7 +184,9 @@ namespace irgen {
                                           ArrayRef<Requirement> requirements);
 
   llvm::GlobalValue *emitAsyncFunctionPointer(IRGenModule &IGM,
-                                              SILFunction *function, Size size);
+                                              llvm::Function *function,
+                                              LinkEntity entity,
+                                              Size size);
 } // end namespace irgen
 } // end namespace swift
 

lib/IRGen/GenThunk.cpp

@@ -81,6 +81,13 @@ static FunctionPointer lookupMethod(IRGenFunction &IGF, SILDeclRef declRef) {
     return layout;
   };
 
+  if (funcTy->isAsync()) {
+    auto layout = getAsyncContextLayout();
+    auto entity = LinkEntity::forDispatchThunk(declRef);
+    emitAsyncFunctionEntry(IGF, layout, entity);
+    emitAsyncFunctionPointer(IGF.IGM, IGF.CurFn, entity, layout.getSize());
+  }
+
   // Protocol case.
   if (isa<ProtocolDecl>(decl->getDeclContext())) {
     // Find the witness table.
@@ -162,18 +169,25 @@ void IRGenModule::emitDispatchThunk(SILDeclRef declRef) {
     IGF.Builder.CreateRet(result);
 }
 
+llvm::Constant *
+IRGenModule::getAddrOfAsyncFunctionPointer(LinkEntity entity) {
+  return getAddrOfLLVMVariable(
+    LinkEntity::forAsyncFunctionPointer(entity),
+    NotForDefinition, DebugTypeInfo());
+}
+
 llvm::Constant *
 IRGenModule::getAddrOfAsyncFunctionPointer(SILFunction *function) {
   (void)getAddrOfSILFunction(function, NotForDefinition);
-  auto entity = LinkEntity::forAsyncFunctionPointer(function);
-  return getAddrOfLLVMVariable(entity, NotForDefinition, DebugTypeInfo());
+  return getAddrOfAsyncFunctionPointer(
+      LinkEntity::forSILFunction(function));
 }
 
-llvm::Constant *IRGenModule::defineAsyncFunctionPointer(SILFunction *function,
+llvm::Constant *IRGenModule::defineAsyncFunctionPointer(LinkEntity entity,
                                                         ConstantInit init) {
-  auto entity = LinkEntity::forAsyncFunctionPointer(function);
+  auto asyncEntity = LinkEntity::forAsyncFunctionPointer(entity);
   auto *var = cast<llvm::GlobalVariable>(
-      getAddrOfLLVMVariable(entity, init, DebugTypeInfo()));
+      getAddrOfLLVMVariable(asyncEntity, init, DebugTypeInfo()));
   setTrueConstGlobal(var);
   return var;
 }

lib/IRGen/IRGenModule.h

@@ -1396,8 +1396,9 @@ private:                            \
   /// Cast the given constant to i8*.
   llvm::Constant *getOpaquePtr(llvm::Constant *pointer);
 
+  llvm::Constant *getAddrOfAsyncFunctionPointer(LinkEntity entity);
   llvm::Constant *getAddrOfAsyncFunctionPointer(SILFunction *function);
-  llvm::Constant *defineAsyncFunctionPointer(SILFunction *function,
+  llvm::Constant *defineAsyncFunctionPointer(LinkEntity entity,
                                              ConstantInit init);
   SILFunction *getSILFunctionForAsyncFunctionPointer(llvm::Constant *afp);
 

lib/IRGen/IRGenSIL.cpp

@@ -26,6 +26,7 @@
 #include "swift/Basic/ExternalUnion.h"
 #include "swift/Basic/Range.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/IRGen/Linking.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/InstructionUtils.h"
@@ -1721,7 +1722,9 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
   }
 
   if (funcTy->isAsync()) {
-    emitAsyncFunctionEntry(IGF, IGF.CurSILFn);
+    emitAsyncFunctionEntry(IGF,
+                           getAsyncContextLayout(IGF.IGM, IGF.CurSILFn),
+                           LinkEntity::forSILFunction(IGF.CurSILFn));
   }
 
   SILFunctionConventions conv(funcTy, IGF.getSILModule());
@@ -1943,9 +1946,12 @@ void IRGenSILFunction::emitSILFunction() {
     IGM.IRGen.addDynamicReplacement(CurSILFn);
 
   auto funcTy = CurSILFn->getLoweredFunctionType();
-  if (funcTy->isAsync() && funcTy->getLanguage() == SILFunctionLanguage::Swift)
-    emitAsyncFunctionPointer(IGM, CurSILFn,
+  if (funcTy->isAsync() && funcTy->getLanguage() == SILFunctionLanguage::Swift) {
+    emitAsyncFunctionPointer(IGM,
+                             CurFn,
+                             LinkEntity::forSILFunction(CurSILFn),
                              getAsyncContextLayout(*this).getSize());
+  }
 
   // Configure the dominance resolver.
   // TODO: consider re-using a dom analysis from the PassManager
@@ -2278,14 +2284,14 @@ void IRGenSILFunction::visitFunctionRefBaseInst(FunctionRefBaseInst *i) {
   auto *fnPtr = IGM.getAddrOfSILFunction(
       fn, NotForDefinition, false /*isDynamicallyReplaceableImplementation*/,
       isa<PreviousDynamicFunctionRefInst>(i));
-  llvm::Value *value;
+  llvm::Constant *value;
   auto isSpecialAsyncWithoutCtxtSize =
       fn->isAsync() && (
           fn->getName().equals("swift_task_future_wait") ||
           fn->getName().equals("swift_task_group_wait_next"));
   if (fn->isAsync() && !isSpecialAsyncWithoutCtxtSize) {
     value = IGM.getAddrOfAsyncFunctionPointer(fn);
-    value = Builder.CreateBitCast(value, fnPtr->getType());
+    value = llvm::ConstantExpr::getBitCast(value, fnPtr->getType());
   } else {
     value = fnPtr;
   }