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

Author: nate-chandler

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,8 @@ NodePointer Demangler::demangleThunkOrSpecialization() {
         return demangleAutoDiffFunctionOrSimpleThunk(
             Node::Kind::AutoDiffFunction);
       }
+    case 'w':
+      return createNode(Node::Kind::AsyncNonconstantPartialApplyThunk);
     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,9 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
   case Node::Kind::AsyncFunctionPointer:
     Printer << "async function pointer to ";
     return nullptr;
+  case Node::Kind::AsyncNonconstantPartialApplyThunk:
+    Printer << "thunk for non-constant partial apply in ";
+    return nullptr;
   }
   printer_unreachable("bad node kind!");
 }

lib/Demangling/OldRemangler.cpp

@@ -606,6 +606,11 @@ void Remangler::mangleDynamicallyReplaceableFunctionVar(Node *node) {
   Buffer << "TX";
 }
 
+void Remangler::mangleAsyncNonconstantPartialApplyThunk(Node *node) {
+  Buffer << "Tw";
+  mangleSingleChildNode(node); // global
+}
+
 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) {
+  mangleChildNodesReversed(node);
+  Buffer << "Tw";
+}
+
 void Remangler::manglePostfixOperator(Node *node) {
   mangleIdentifierImpl(node, /*isOperator*/ true);
   Buffer << "oP";

lib/IRGen/IRGenMangler.cpp

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

lib/IRGen/IRGenMangler.h

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

test/Demangle/Inputs/manglings.txt

@@ -388,3 +388,6 @@ $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 ---> thunk for non-constant partial apply in create_pa_f2
+$s0019create_pa_f21_HJEEdTw ---> thunk for non-constant partial apply in create_pa_f2.1
+$s4main3fooyySiFyyXEfU_TA.1 ---> {T:} closure #1 () -> () in main.foo(Swift.Int) -> ()partial apply forwarder with unmangled suffix ".1"