[Async CC] Pull poly params from entry point emission.

Previously, EmitPolymorphicParameters dealt directly with an Explosion
from which it pulled values.  In one place, there was a conditional
check for async which handled some cases.  There was however another
place where the polymorphic parameter was pulled directly from the
explosion.  That missed case resulted in attempting to pull a
polymorphic parameter directly from an Explosion which contains only a
%swift.context* per the async calling convention.

Here, those parameters are now pulled from an EntryPointArgumentEmission
subclasses of which are able to provide the relevant definition of what
pulling a parameter means.

rdar://problem/70144083

Author: nate-chandler

lib/IRGen/EntryPointArgumentEmission.h

@@ -20,6 +20,7 @@ namespace swift {
 namespace irgen {
 
 class Explosion;
+struct GenericRequirement;
 
 class EntryPointArgumentEmission {
 
@@ -28,6 +29,9 @@ class EntryPointArgumentEmission {
   virtual bool requiresIndirectResult(SILType retType) = 0;
   virtual llvm::Value *getIndirectResultForFormallyDirectResult() = 0;
   virtual llvm::Value *getIndirectResult(unsigned index) = 0;
+  virtual llvm::Value *getNextPolymorphicParameterAsMetadata() = 0;
+  virtual llvm::Value *
+  getNextPolymorphicParameter(GenericRequirement &requirement) = 0;
 };
 
 class NativeCCEntryPointArgumentEmission

lib/IRGen/GenProto.cpp

@@ -489,7 +489,8 @@ class EmitPolymorphicParameters : public PolymorphicConvention {
 public:
   EmitPolymorphicParameters(IRGenFunction &IGF, SILFunction &Fn);
 
-  void emit(Explosion &in, WitnessMetadata *witnessMetadata,
+  void emit(EntryPointArgumentEmission &emission,
+            WitnessMetadata *witnessMetadata,
             const GetParameterFn &getParameter);
 
 private:
@@ -499,7 +500,8 @@ class EmitPolymorphicParameters : public PolymorphicConvention {
 
   /// Fulfill local type data from any extra information associated with
   /// the given source.
-  void bindExtraSource(const MetadataSource &source, Explosion &in,
+  void bindExtraSource(const MetadataSource &source,
+                       EntryPointArgumentEmission &emission,
                        WitnessMetadata *witnessMetadata);
 
   void bindParameterSources(const GetParameterFn &getParameter);
@@ -528,9 +530,9 @@ CanType EmitPolymorphicParameters::getArgTypeInContext(unsigned paramIndex) cons
       IGM.getSILModule(), FnType, IGM.getMaximalTypeExpansionContext()));
 }
 
-void EmitPolymorphicParameters::bindExtraSource(const MetadataSource &source,
-                                                Explosion &in,
-                                         WitnessMetadata *witnessMetadata) {
+void EmitPolymorphicParameters::bindExtraSource(
+    const MetadataSource &source, EntryPointArgumentEmission &emission,
+    WitnessMetadata *witnessMetadata) {
   switch (source.getKind()) {
     case MetadataSource::Kind::Metadata:
     case MetadataSource::Kind::ClassPointer:
@@ -540,7 +542,7 @@ void EmitPolymorphicParameters::bindExtraSource(const MetadataSource &source,
     case MetadataSource::Kind::GenericLValueMetadata: {
       CanType argTy = getArgTypeInContext(source.getParamIndex());
 
-      llvm::Value *metadata = in.claimNext();
+      llvm::Value *metadata = emission.getNextPolymorphicParameterAsMetadata();
       setTypeMetadataName(IGF.IGM, metadata, argTy);
 
       IGF.bindLocalTypeDataFromTypeMetadata(argTy, IsExact, metadata,
@@ -2223,55 +2225,23 @@ bool irgen::hasPolymorphicParameters(CanSILFunctionType ty) {
 }
 
 /// Emit a polymorphic parameters clause, binding all the metadata necessary.
-void EmitPolymorphicParameters::emit(Explosion &in,
+void EmitPolymorphicParameters::emit(EntryPointArgumentEmission &emission,
                                      WitnessMetadata *witnessMetadata,
                                      const GetParameterFn &getParameter) {
   // Collect any early sources and bind local type data from them.
   for (auto &source : getSources()) {
-    bindExtraSource(source, in, witnessMetadata);
+    bindExtraSource(source, emission, witnessMetadata);
   }
 
   auto getInContext = [&](CanType type) -> CanType {
     return getTypeInContext(type);
   };
 
-  unsigned index = 0;
   // Collect any concrete type metadata that's been passed separately.
   enumerateUnfulfilledRequirements([&](GenericRequirement requirement) {
-    llvm::Value *value;
-    if (Fn.isAsync()) {
-      auto *context = in.peek(
-          /*the index of the swift.context in the async function CC*/ 0);
-      auto layout = getAsyncContextLayout(IGF, &Fn);
-      Address dataAddr = layout.emitCastTo(IGF, context);
-      assert(layout.hasBindings());
-
-      auto bindingLayout = layout.getBindingsLayout();
-      auto bindingsAddr =
-          bindingLayout.project(IGF, dataAddr, /*offsets*/ None);
-      auto erasedBindingsAddr =
-          IGF.Builder.CreateBitCast(bindingsAddr, IGF.IGM.Int8PtrPtrTy);
-      auto uncastBindingAddr = IGF.Builder.CreateConstArrayGEP(
-          erasedBindingsAddr, index, IGF.IGM.getPointerSize());
-      if (requirement.Protocol) {
-        auto bindingAddrAddr = IGF.Builder.CreateBitCast(
-            uncastBindingAddr.getAddress(), IGF.IGM.WitnessTablePtrPtrTy);
-        auto bindingAddr = IGF.Builder.CreateLoad(
-            bindingAddrAddr, IGF.IGM.getPointerAlignment());
-        value = bindingAddr;
-      } else {
-        auto bindingAddrAddr = IGF.Builder.CreateBitCast(
-            uncastBindingAddr.getAddress(), IGF.IGM.TypeMetadataPtrPtrTy);
-        auto bindingAddr = IGF.Builder.CreateLoad(
-            bindingAddrAddr, IGF.IGM.getPointerAlignment());
-        value = bindingAddr;
-      }
-    } else {
-      value = in.claimNext();
-    }
+    llvm::Value *value = emission.getNextPolymorphicParameter(requirement);
     bindGenericRequirement(IGF, requirement, value, MetadataState::Complete,
                            getInContext);
-    ++index;
   });
 
   // Bind all the fulfillments we can from the formal parameters.
@@ -2689,12 +2659,12 @@ void irgen::collectTrailingWitnessMetadata(
 }
 
 /// Perform all the bindings necessary to emit the given declaration.
-void irgen::emitPolymorphicParameters(IRGenFunction &IGF,
-                                      SILFunction &Fn,
-                                      Explosion &in,
+void irgen::emitPolymorphicParameters(IRGenFunction &IGF, SILFunction &Fn,
+                                      EntryPointArgumentEmission &emission,
                                       WitnessMetadata *witnessMetadata,
                                       const GetParameterFn &getParameter) {
-  EmitPolymorphicParameters(IGF, Fn).emit(in, witnessMetadata, getParameter);
+  EmitPolymorphicParameters(IGF, Fn).emit(emission, witnessMetadata,
+                                          getParameter);
 }
 
 /// Given an array of polymorphic arguments as might be set up by

lib/IRGen/GenProto.h

@@ -40,6 +40,7 @@ namespace swift {
 namespace irgen {
   class Address;
   class DynamicMetadataRequest;
+  class EntryPointArgumentEmission;
   class Explosion;
   class FunctionPointer;
   class IRGenFunction;
@@ -133,12 +134,11 @@ namespace irgen {
   ///
   /// \param witnessMetadata - can be omitted if the function is
   ///   definitely not a witness method
-  void emitPolymorphicParameters(IRGenFunction &IGF,
-                                 SILFunction &Fn,
-                                 Explosion &args,
+  void emitPolymorphicParameters(IRGenFunction &IGF, SILFunction &Fn,
+                                 EntryPointArgumentEmission &emission,
                                  WitnessMetadata *witnessMetadata,
                                  const GetParameterFn &getParameter);
- 
+
   void emitPolymorphicParametersFromArray(IRGenFunction &IGF,
                                           NominalTypeDecl *typeDecl,
                                           Address array,

lib/IRGen/IRGenSIL.cpp

@@ -1147,6 +1147,13 @@ class SyncEntryPointArgumentEmission
   llvm::Value *getIndirectResult(unsigned index) override {
     return allParamValues.claimNext();
   };
+  llvm::Value *
+  getNextPolymorphicParameter(GenericRequirement &requirement) override {
+    return allParamValues.claimNext();
+  }
+  llvm::Value *getNextPolymorphicParameterAsMetadata() override {
+    return allParamValues.claimNext();
+  }
 };
 class AsyncEntryPointArgumentEmission
     : public virtual EntryPointArgumentEmission {
@@ -1206,6 +1213,7 @@ class AsyncNativeCCEntryPointArgumentEmission final
   llvm::Value *context;
   /*const*/ AsyncContextLayout layout;
   const Address dataAddr;
+  unsigned polymorphicParameterIndex = 0;
 
   llvm::Value *loadValue(ElementLayout layout) {
     Address addr = layout.project(IGF, dataAddr, /*offsets*/ llvm::None);
@@ -1248,6 +1256,46 @@ class AsyncNativeCCEntryPointArgumentEmission final
                      "into indirect IR results because all results are already "
                      "indirected through the context");
   }
+  Address getNextUncastBinding() {
+    auto index = polymorphicParameterIndex;
+    ++polymorphicParameterIndex;
+
+    assert(layout.hasBindings());
+
+    auto bindingLayout = layout.getBindingsLayout();
+    auto bindingsAddr = bindingLayout.project(IGF, dataAddr, /*offsets*/ None);
+    auto erasedBindingsAddr =
+        IGF.Builder.CreateBitCast(bindingsAddr, IGF.IGM.Int8PtrPtrTy);
+    auto uncastBindingAddr = IGF.Builder.CreateConstArrayGEP(
+        erasedBindingsAddr, index, IGF.IGM.getPointerSize());
+    return uncastBindingAddr;
+  }
+  llvm::Value *castUncastBindingToMetadata(Address uncastBindingAddr) {
+    auto bindingAddrAddr = IGF.Builder.CreateBitCast(
+        uncastBindingAddr.getAddress(), IGF.IGM.TypeMetadataPtrPtrTy);
+    auto bindingAddr =
+        IGF.Builder.CreateLoad(bindingAddrAddr, IGF.IGM.getPointerAlignment());
+    return bindingAddr;
+  }
+  llvm::Value *castUncastBindingToWitnessTable(Address uncastBindingAddr) {
+    auto bindingAddrAddr = IGF.Builder.CreateBitCast(
+        uncastBindingAddr.getAddress(), IGF.IGM.WitnessTablePtrPtrTy);
+    auto bindingAddr =
+        IGF.Builder.CreateLoad(bindingAddrAddr, IGF.IGM.getPointerAlignment());
+    return bindingAddr;
+  }
+  llvm::Value *
+  getNextPolymorphicParameter(GenericRequirement &requirement) override {
+    auto uncastBindingAddr = getNextUncastBinding();
+    if (requirement.Protocol) {
+      return castUncastBindingToWitnessTable(uncastBindingAddr);
+    } else {
+      return castUncastBindingToMetadata(uncastBindingAddr);
+    }
+  }
+  llvm::Value *getNextPolymorphicParameterAsMetadata() override {
+    return castUncastBindingToMetadata(getNextUncastBinding());
+  }
   llvm::Value *getIndirectResult(unsigned index) override {
     auto fieldLayout = layout.getIndirectReturnLayout(index);
     return loadValue(fieldLayout);
@@ -1675,13 +1723,13 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
   // Bind polymorphic arguments.  This can only be done after binding
   // all the value parameters.
   if (hasPolymorphicParameters(funcTy)) {
-    emitPolymorphicParameters(IGF, *IGF.CurSILFn, allParamValues,
-                              &witnessMetadata,
-      [&](unsigned paramIndex) -> llvm::Value* {
-        SILValue parameter =
-          IGF.CurSILFn->getArgumentsWithoutIndirectResults()[paramIndex];
-        return IGF.getLoweredSingletonExplosion(parameter);
-      });
+    emitPolymorphicParameters(
+        IGF, *IGF.CurSILFn, *emission, &witnessMetadata,
+        [&](unsigned paramIndex) -> llvm::Value * {
+          SILValue parameter =
+              IGF.CurSILFn->getArgumentsWithoutIndirectResults()[paramIndex];
+          return IGF.getLoweredSingletonExplosion(parameter);
+        });
   }
 
   assert(allParamValues.empty() && "didn't claim all parameters!");
@@ -1776,7 +1824,7 @@ static void emitEntryPointArgumentsCOrObjC(IRGenSILFunction &IGF,
   // all the value parameters, and must be done even for non-polymorphic
   // functions because of imported Objective-C generics.
   emitPolymorphicParameters(
-      IGF, *IGF.CurSILFn, params, nullptr,
+      IGF, *IGF.CurSILFn, *emission, nullptr,
       [&](unsigned paramIndex) -> llvm::Value * {
         SILValue parameter = entry->getArguments()[paramIndex];
         return IGF.getLoweredSingletonExplosion(parameter);

validation-test/compiler_crashers_2_fixed/rdar70144083.swift

@@ -0,0 +1,37 @@
+// RUN: %target-swift-frontend -emit-ir %s -enable-experimental-concurrency
+
+public protocol AsyncIteratorProtocol {
+    associatedtype Element
+    associatedtype Failure: Error
+
+    mutating func nextResult() async -> Result<Element, Failure>?
+    mutating func cancel()
+}
+
+public protocol AsyncSequence {
+    associatedtype Element
+    associatedtype Failure: Error
+    associatedtype AsyncIterator: AsyncIteratorProtocol where AsyncIterator.Element == Element, AsyncIterator.Failure == Failure
+    
+    func makeAsyncIterator() -> AsyncIterator
+}
+
+struct Just<Element>: AsyncSequence {
+    typealias Failure = Never
+
+    struct AsyncIterator: AsyncIteratorProtocol {
+        var value: Element?
+        mutating func nextResult() async -> Result<Element, Never>? {
+            defer { value = nil }
+            return value.map { .success($0) }
+        }
+
+        mutating func cancel() {
+            value = nil
+        }
+    }
+    var value: Element
+    func makeAsyncIterator() -> AsyncIterator {
+        return AsyncIterator(value: value)
+    }
+}