SILGen/ClangImporter: Handle async imports with a boolean error flag argument.

docs/ABI/Mangling.rst

@@ -222,8 +222,8 @@ types where the metadata itself has unknown layout.)
   global ::= global 'Tm'                 // merged function
   global ::= entity                      // some identifiable thing
   global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
-  global ::= impl-function-type type 'Tz'     // objc-to-swift-async completion handler block implementation
-  global ::= impl-function-type type 'TZ'     // objc-to-swift-async completion handler block implementation (predefined by runtime)
+  global ::= impl-function-type type 'Tz' index? // objc-to-swift-async completion handler block implementation
+  global ::= impl-function-type type 'TZ' index? // objc-to-swift-async completion handler block implementation (predefined by runtime)
   global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
   global ::= impl-function-type type generic-signature? 'Tz'     // objc-to-swift-async completion handler block implementation
   global ::= impl-function-type type generic-signature? 'TZ'     // objc-to-swift-async completion handler block implementation (predefined by runtime)

include/swift/AST/ASTMangler.h

@@ -172,6 +172,7 @@ class ASTMangler : public Mangler {
   std::string mangleObjCAsyncCompletionHandlerImpl(CanSILFunctionType BlockType,
                                                    CanType ResultType,
                                                    CanGenericSignature Sig,
+                                                   Optional<bool> FlagParamIsZeroOnError,
                                                    bool predefined);
 
   /// Mangle the derivative function (JVP/VJP), or optionally its vtable entry

include/swift/SIL/AbstractionPattern.h

@@ -276,6 +276,20 @@ class AbstractionPattern {
       Async,
     };
 
+    enum {
+      AsyncCompletionParameterIndexMask = 0xFFEu,
+      AsyncCompletionParameterIndexShift = 1,
+
+      AsyncCompletionErrorParameterIndexMask = 0x1FF000u,
+      AsyncCompletionErrorParameterIndexShift = 12,
+
+      AsyncCompletionErrorFlagParameterIndexMask = 0x3FE00000u,
+      AsyncCompletionErrorFlagParameterIndexShift = 21,
+
+      AsyncCompletionErrorFlagParameterPolarityMask = 0x40000000u,
+      AsyncCompletionErrorFlagParameterPolarityShift = 30,
+    };
+
   public:
     enum ForeignKind {
       IsNotForeign,
@@ -304,16 +318,24 @@ class AbstractionPattern {
 
     EncodedForeignInfo(Async_t,
                        unsigned completionParameterIndex,
-                       Optional<unsigned> completionErrorParameterIndex)
-      : Value(1
-              + (unsigned(IsAsync) - 1)
-              + (unsigned(completionParameterIndex) << 1)
-              + ((completionErrorParameterIndex ? *completionErrorParameterIndex + 1
-                                                : 0) << 21)) {
-        assert(getKind() == IsAsync);
-        assert(getAsyncCompletionHandlerParamIndex() == completionParameterIndex);
-        assert(getAsyncCompletionHandlerErrorParamIndex() == completionErrorParameterIndex);
-      }
+                       Optional<unsigned> completionErrorParameterIndex,
+                       Optional<unsigned> completionErrorFlagParameterIndex,
+                       bool completionErrorFlagIsZeroOnError)
+    : Value(1
+      + (unsigned(IsAsync) - 1)
+      + (unsigned(completionParameterIndex) << AsyncCompletionParameterIndexShift)
+      + ((completionErrorParameterIndex ? *completionErrorParameterIndex + 1
+                                        : 0) << AsyncCompletionErrorParameterIndexShift)
+      + ((completionErrorFlagParameterIndex ? *completionErrorFlagParameterIndex + 1
+                                            : 0) << AsyncCompletionErrorFlagParameterIndexShift)
+      + (unsigned(completionErrorFlagIsZeroOnError) << AsyncCompletionErrorFlagParameterPolarityShift)){
+
+      assert(getKind() == IsAsync);
+      assert(getAsyncCompletionHandlerParamIndex() == completionParameterIndex);
+      assert(getAsyncCompletionHandlerErrorParamIndex() == completionErrorParameterIndex);
+      assert(getAsyncCompletionHandlerErrorFlagParamIndex() == completionErrorFlagParameterIndex);
+      assert(isCompletionErrorFlagZeroOnError() == completionErrorFlagIsZeroOnError);
+    }
 
   public:
     static EncodedForeignInfo
@@ -345,18 +367,37 @@ class AbstractionPattern {
 
     unsigned getAsyncCompletionHandlerParamIndex() const {
       assert(getKind() == IsAsync);
-      return ((Value - 1) >> 1) & 0xFFFFFu;
+      return ((Value - 1) & AsyncCompletionParameterIndexMask)
+        >> AsyncCompletionParameterIndexShift;
     }
 
     Optional<unsigned> getAsyncCompletionHandlerErrorParamIndex() const {
       assert(getKind() == IsAsync);
 
-      unsigned encodedValue = (Value - 1) >> 21;
+      unsigned encodedValue = ((Value - 1) & AsyncCompletionErrorParameterIndexMask)
+        >> AsyncCompletionErrorParameterIndexShift;
+      if (encodedValue == 0) {
+        return llvm::None;
+      }
+      return encodedValue - 1;
+    }
+
+    Optional<unsigned> getAsyncCompletionHandlerErrorFlagParamIndex() const {
+      assert(getKind() == IsAsync);
+
+      unsigned encodedValue = ((Value - 1) & AsyncCompletionErrorFlagParameterIndexMask)
+        >> AsyncCompletionErrorFlagParameterIndexShift;
       if (encodedValue == 0) {
         return llvm::None;
       }
       return encodedValue - 1;
     }
+
+    bool isCompletionErrorFlagZeroOnError() const {
+      assert(getKind() == IsAsync);
+
+      return (Value - 1) & AsyncCompletionErrorFlagParameterPolarityMask;
+    }
 
     unsigned getForeignParamIndex() const {
       switch (getKind()) {

lib/AST/ASTMangler.cpp

@@ -399,13 +399,17 @@ std::string ASTMangler::mangleObjCAsyncCompletionHandlerImpl(
                                                    CanSILFunctionType BlockType,
                                                    CanType ResultType,
                                                    CanGenericSignature Sig,
+                                                   Optional<bool> ErrorOnZero,
                                                    bool predefined) {
   beginMangling();
   appendType(BlockType);
   appendType(ResultType);
   if (Sig)
     appendGenericSignature(Sig);
-  appendOperator(predefined ? "TZ" : "Tz");
+  if (ErrorOnZero)
+    appendOperator(predefined ? "TZ" : "Tz", Index(*ErrorOnZero + 1));
+  else
+    appendOperator(predefined ? "TZ" : "Tz", Index(0));
   return finalize();
 }
 

lib/ClangImporter/ImportName.cpp

@@ -1497,8 +1497,23 @@ ImportedName NameImporter::importNameImpl(const clang::NamedDecl *D,
           swiftAsyncAttr->getCompletionHandlerIndex().getASTIndex();
     }
 
-    // TODO: Check for the swift_async_error attribute here when Clang
-    // implements it
+    if (const auto *asyncErrorAttr = D->getAttr<clang::SwiftAsyncErrorAttr>()) {
+      switch (auto convention = asyncErrorAttr->getConvention()) {
+      // No flag parameter in these cases.
+      case clang::SwiftAsyncErrorAttr::NonNullError:
+      case clang::SwiftAsyncErrorAttr::None:
+        break;
+
+      // Get the flag argument index and polarity from the attribute.
+      case clang::SwiftAsyncErrorAttr::NonZeroArgument:
+      case clang::SwiftAsyncErrorAttr::ZeroArgument:
+        // NB: Attribute is 1-based rather than 0-based.
+        completionHandlerFlagParamIndex = asyncErrorAttr->getHandlerParamIdx() - 1;
+        completionHandlerFlagIsZeroOnError =
+          convention == clang::SwiftAsyncErrorAttr::ZeroArgument;
+        break;
+      }
+    }
   }
 
   // FIXME: ugly to check here, instead perform unified check up front in

lib/ClangImporter/ImportType.cpp

@@ -2120,11 +2120,14 @@ static Type decomposeCompletionHandlerType(
     if (param.isInOut() || param.isVariadic())
       return Type();
 
-    // If there is an error parameter to the completion handler, it is
+    // If there are error-related parameters to the completion handler, they are
     // not part of the result type of the asynchronous function.
     if (info.completionHandlerErrorParamIndex() &&
         paramIdx == *info.completionHandlerErrorParamIndex())
       continue;
+    if (info.completionHandlerFlagParamIndex() &&
+        paramIdx == *info.completionHandlerFlagParamIndex())
+      continue;
 
     resultTypeElts.push_back(param.getPlainType());
   }

lib/Demangling/Demangler.cpp

@@ -2301,13 +2301,14 @@ NodePointer Demangler::demangleThunkOrSpecialization() {
     }
     case 'z':
     case 'Z': {
+      NodePointer flagMode = demangleIndexAsNode();
       NodePointer sig = popNode(Node::Kind::DependentGenericSignature);
       NodePointer resultType = popNode(Node::Kind::Type);
       NodePointer implType = popNode(Node::Kind::Type);
       auto node = createWithChildren(c == 'z'
                                   ? Node::Kind::ObjCAsyncCompletionHandlerImpl
                                   : Node::Kind::PredefinedObjCAsyncCompletionHandlerImpl,
-                                implType, resultType);
+                                implType, resultType, flagMode);
       if (sig)
         addChild(node, sig);
       return node;

lib/Demangling/NodePrinter.cpp

@@ -2742,9 +2742,24 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
     LLVM_FALLTHROUGH;
   case Node::Kind::ObjCAsyncCompletionHandlerImpl:
     Printer << "@objc completion handler block implementation for ";
+    if (Node->getNumChildren() >= 4)
+      print(Node->getChild(3));
     print(Node->getChild(0));
     Printer << " with result type ";
     print(Node->getChild(1));
+    switch (Node->getChild(2)->getIndex()) {
+    case 0:
+      break;
+    case 1:
+      Printer << " nonzero on error";
+      break;
+    case 2:
+      Printer << " zero on error";
+      break;
+    default:
+      Printer << " <invalid error flag>";
+      break;
+    }
     return nullptr;
   case Node::Kind::CanonicalPrespecializedGenericTypeCachingOnceToken:
     Printer << "flag for loading of canonical specialized generic type "

lib/Demangling/Remangler.cpp

@@ -820,8 +820,12 @@ void Remangler::manglePredefinedObjCAsyncCompletionHandlerImpl(Node *node) {
 }
 
 void Remangler::mangleObjCAsyncCompletionHandlerImpl(Node *node) {
-  mangleChildNodes(node);
+  mangleChildNode(node, 0);
+  mangleChildNode(node, 1);
+  if (node->getNumChildren() == 4)
+    mangleChildNode(node, 3);
   Buffer << "Tz";
+  mangleChildNode(node, 2);
 }
 
 void Remangler::mangleDeallocator(Node *node) {

lib/SIL/IR/AbstractionPattern.cpp

@@ -132,7 +132,9 @@ AbstractionPattern::EncodedForeignInfo::encode(
   if (foreignAsync.hasValue()) {
     return EncodedForeignInfo(EncodedForeignInfo::Async,
                               foreignAsync->completionHandlerParamIndex(),
-                              foreignAsync->completionHandlerErrorParamIndex());
+                              foreignAsync->completionHandlerErrorParamIndex(),
+                              foreignAsync->completionHandlerFlagParamIndex(),
+                              foreignAsync->completionHandlerFlagIsErrorOnZero());
   } else if (foreignError.hasValue()) {
     return EncodedForeignInfo(EncodedForeignInfo::Error,
                               foreignError->getErrorParameterIndex(),
@@ -404,11 +406,17 @@ AbstractionPattern::getTupleElementType(unsigned index) const {
 
   case Kind::ObjCCompletionHandlerArgumentsType: {
     // Match up the tuple element with the parameter from the Clang block type,
-    // skipping the error parameter index if any.
+    // skipping the error parameter and flag indexes if any.
     auto callback = cast<clang::FunctionProtoType>(getClangType());
     auto errorIndex = getEncodedForeignInfo()
       .getAsyncCompletionHandlerErrorParamIndex();
-    unsigned paramIndex = index + (errorIndex && index >= *errorIndex);
+    auto flagIndex = getEncodedForeignInfo()
+      .getAsyncCompletionHandlerErrorFlagParamIndex();
+    unsigned paramIndex = index;
+    if (errorIndex && paramIndex >= *errorIndex)
+      ++paramIndex;
+    if (flagIndex && paramIndex >= *flagIndex)
+      ++paramIndex;
     return AbstractionPattern(getGenericSignature(),
                               getCanTupleElementType(getType(), index),
                               callback->getParamType(paramIndex).getTypePtr());
@@ -544,12 +552,18 @@ AbstractionPattern AbstractionPattern::getFunctionResultType() const {
       // any.
 
       auto callbackErrorIndex = getEncodedForeignInfo()
-                                  .getAsyncCompletionHandlerErrorParamIndex();
+                                    .getAsyncCompletionHandlerErrorParamIndex();
+      auto callbackErrorFlagIndex = getEncodedForeignInfo()
+                                .getAsyncCompletionHandlerErrorFlagParamIndex();
       assert((!callbackErrorIndex.hasValue()
               || callbackParamTy->getNumParams() > *callbackErrorIndex)
              && "completion handler has invalid error param index?!");
+      assert((!callbackErrorFlagIndex.hasValue()
+              || callbackParamTy->getNumParams() > *callbackErrorFlagIndex)
+             && "completion handler has invalid error param index?!");
       unsigned numNonErrorParams
-        = callbackParamTy->getNumParams() - callbackErrorIndex.hasValue();
+        = callbackParamTy->getNumParams() - callbackErrorIndex.hasValue()
+                                          - callbackErrorFlagIndex.hasValue();
 
       switch (numNonErrorParams) {
       case 0:
@@ -560,8 +574,13 @@ AbstractionPattern AbstractionPattern::getFunctionResultType() const {
       case 1: {
         // If there's a single argument, abstract it according to its formal type
         // in the ObjC signature.
-        unsigned callbackResultIndex
-          = callbackErrorIndex && *callbackErrorIndex == 0;
+        unsigned callbackResultIndex = 0;
+        if (callbackErrorIndex && callbackResultIndex >= *callbackErrorIndex)
+          ++callbackResultIndex;
+        if (callbackErrorFlagIndex
+            && callbackResultIndex >= *callbackErrorFlagIndex)
+          ++callbackResultIndex;
+
         auto clangResultType = callbackParamTy
           ->getParamType(callbackResultIndex)
           .getTypePtr();
@@ -1031,7 +1050,14 @@ void AbstractionPattern::print(raw_ostream &out) const {
     case EncodedForeignInfo::IsAsync:
       out << ", completionHandlerParameter=" << errorInfo.getAsyncCompletionHandlerParamIndex();
       if (auto errorParam = errorInfo.getAsyncCompletionHandlerErrorParamIndex()) {
-        out << " (errorParam=" << *errorParam << ')';
+        out << " (errorParam=" << *errorParam;
+        if (auto errorFlag = errorInfo.getAsyncCompletionHandlerErrorFlagParamIndex()) {
+          out << ", errorFlagParam=" << *errorFlag
+              << (errorInfo.isCompletionErrorFlagZeroOnError()
+                    ? ", zeroOnError"
+                    : ", nonzeroOnError");
+        }
+        out << ')';
       }
     }
     out << ", ";