Merge pull request #36436 from nate-chandler/concurrency/irgen/nonconstant-async-partial-apply-thunk-mangling

[IRGen] Mangling for async, non-constant partial_apply thunks.

Author: nate-chandler

docs/ABI/Mangling.rst

@@ -147,6 +147,7 @@ Globals
   // TODO check this::
   global ::= mangled-name 'TA'                     // partial application forwarder
   global ::= mangled-name 'Ta'                     // ObjC partial application forwarder
+  global ::= mangled-name 'Tw' index               // async partial apply thunk for a non-constant function
 
   global ::= type 'w' VALUE-WITNESS-KIND // value witness
 

include/swift/Demangling/DemangleNodes.def

@@ -316,6 +316,7 @@ NODE(AutoDiffSubsetParametersThunk)
 NODE(AutoDiffDerivativeVTableThunk)
 NODE(DifferentiabilityWitness)
 NODE(IndexSubset)
+NODE(AsyncNonconstantPartialApplyThunk)
 
 #undef CONTEXT_NODE
 #undef NODE

lib/Demangling/Demangler.cpp

@@ -123,6 +123,7 @@ bool swift::Demangle::isFunctionAttr(Node::Kind kind) {
     case Node::Kind::DynamicallyReplaceableFunctionKey:
     case Node::Kind::DynamicallyReplaceableFunctionVar:
     case Node::Kind::AsyncFunctionPointer:
+    case Node::Kind::AsyncNonconstantPartialApplyThunk:
       return true;
     default:
       return false;
@@ -2520,6 +2521,11 @@ NodePointer Demangler::demangleThunkOrSpecialization() {
         return demangleAutoDiffFunctionOrSimpleThunk(
             Node::Kind::AutoDiffFunction);
       }
+    case 'w': {
+      NodePointer discriminator = demangleIndexAsNode();
+      return createWithChild(Node::Kind::AsyncNonconstantPartialApplyThunk,
+                             discriminator);
+    }
     default:
       return nullptr;
   }

lib/Demangling/NodePrinter.cpp

@@ -573,6 +573,7 @@ class NodePrinter {
     case Node::Kind::AutoDiffFunctionKind:
     case Node::Kind::DifferentiabilityWitness:
     case Node::Kind::IndexSubset:
+    case Node::Kind::AsyncNonconstantPartialApplyThunk:
       return false;
     }
     printer_unreachable("bad node kind");
@@ -2774,6 +2775,12 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
   case Node::Kind::AsyncFunctionPointer:
     Printer << "async function pointer to ";
     return nullptr;
+  case Node::Kind::AsyncNonconstantPartialApplyThunk:
+    Printer << "(";
+    print(Node->getChild(0));
+    Printer << ")";
+    Printer << " thunk for non-constant partial apply in ";
+    return nullptr;
   }
   printer_unreachable("bad node kind!");
 }

lib/Demangling/OldRemangler.cpp

@@ -606,6 +606,10 @@ void Remangler::mangleDynamicallyReplaceableFunctionVar(Node *node) {
   Buffer << "TX";
 }
 
+void Remangler::mangleAsyncNonconstantPartialApplyThunk(Node *node) {
+  unreachable("unsupported");
+}
+
 void Remangler::mangleDirectness(Node *node) {
   auto getChar = [](Directness d) -> char {
     switch (d) {

lib/Demangling/Remangler.cpp

@@ -1396,6 +1396,7 @@ void Remangler::mangleGlobal(Node *node) {
       case Node::Kind::DynamicallyReplaceableFunctionImpl:
       case Node::Kind::DynamicallyReplaceableFunctionVar:
       case Node::Kind::AsyncFunctionPointer:
+      case Node::Kind::AsyncNonconstantPartialApplyThunk:
         mangleInReverseOrder = true;
         break;
       default:
@@ -1904,6 +1905,11 @@ void Remangler::mangleDynamicallyReplaceableFunctionVar(Node *node) {
   Buffer << "TX";
 }
 
+void Remangler::mangleAsyncNonconstantPartialApplyThunk(Node *node) {
+  Buffer << "Tw";
+  mangleChildNode(node, 0);
+}
+
 void Remangler::manglePostfixOperator(Node *node) {
   mangleIdentifierImpl(node, /*isOperator*/ true);
   Buffer << "oP";

lib/IRGen/IRGenMangler.cpp

@@ -79,6 +79,23 @@ std::string IRGenMangler::manglePartialApplyForwarder(StringRef FuncName) {
   return finalize();
 }
 
+std::string
+IRGenMangler::mangleAsyncNonconstantPartialApplyThunk(StringRef FuncName,
+                                                      unsigned index) {
+  if (FuncName.empty()) {
+    beginMangling();
+  } else {
+    if (FuncName.startswith(MANGLING_PREFIX_STR)) {
+      Buffer << FuncName;
+    } else {
+      beginMangling();
+      appendIdentifier(FuncName);
+    }
+  }
+  appendOperator("Tw", Index(index));
+  return finalize();
+}
+
 SymbolicMangling
 IRGenMangler::withSymbolicReferences(IRGenModule &IGM,
                                   llvm::function_ref<void ()> body) {

lib/IRGen/IRGenMangler.h

@@ -574,7 +574,9 @@ class IRGenMangler : public Mangle::ASTMangler {
   }
 
   std::string manglePartialApplyForwarder(StringRef FuncName);
-  
+  std::string mangleAsyncNonconstantPartialApplyThunk(StringRef FuncName,
+                                                      unsigned index);
+
   std::string mangleTypeForForeignMetadataUniquing(Type type) {
     return mangleTypeWithoutPrefix(type);
   }

test/Demangle/Inputs/manglings.txt

@@ -388,3 +388,5 @@ $sS2f8mangling3FooV13TangentVectorVIegydd_SfAESfIegydd_TJOp ---> autodiff self-r
 $s13test_mangling3fooyS2f_S2ftFWJrSpSr ---> reverse-mode differentiability witness for test_mangling.foo(Swift.Float, Swift.Float, Swift.Float) -> Swift.Float with respect to parameters {0} and results {0}
 $s13test_mangling3fooyS2f_xq_t16_Differentiation14DifferentiableR_r0_lFAcDRzAcDR_r0_lWJrUSSpSr ---> reverse-mode differentiability witness for test_mangling.foo<A, B where B: _Differentiation.Differentiable>(Swift.Float, A, B) -> Swift.Float with respect to parameters {1, 2} and results {0} with <A, B where A: _Differentiation.Differentiable, B: _Differentiation.Differentiable>
 $s5async1hyyS2iJXEF ---> async.h(@concurrent (Swift.Int) -> Swift.Int) -> ()
+$s12create_pa_f2Tw_ ---> (0) thunk for non-constant partial apply in create_pa_f2
+$s12create_pa_f2Tw0_ ---> (1) thunk for non-constant partial apply in create_pa_f2