[Mangling] Support function specializations that remove async

Author: kavon

docs/ABI/Mangling.rst

@@ -1150,6 +1150,7 @@ Function Specializations
   specialization ::= type '_' type* 'Ts' SPEC-INFO     // Generic re-abstracted prespecialization
   specialization ::= type '_' type* 'TG' SPEC-INFO     // Generic not re-abstracted specialization
   specialization ::= type '_' type* 'Ti' SPEC-INFO     // Inlined function with generic substitutions.
+  specialization ::= type '_' type* 'Ta' SPEC-INFO     // Non-async specialization
 
 The types are the replacement types of the substitution list.
 
@@ -1172,19 +1173,23 @@ Some kinds need arguments, which precede ``Tf``.
   spec-arg ::= identifier
   spec-arg ::= type
 
-  SPEC-INFO ::= MT-REMOVED? FRAGILE? PASSID
+  SPEC-INFO ::= MT-REMOVED? FRAGILE? ASYNC-REMOVED? PASSID
 
   PASSID ::= '0'                             // AllocBoxToStack,
   PASSID ::= '1'                             // ClosureSpecializer,
   PASSID ::= '2'                             // CapturePromotion,
   PASSID ::= '3'                             // CapturePropagation,
   PASSID ::= '4'                             // FunctionSignatureOpts,
   PASSID ::= '5'                             // GenericSpecializer,
+  PASSID ::= '6'                             // MoveDiagnosticInOutToOut,
+  PASSID ::= '7'                             // AsyncDemotion,
 
   MT-REMOVED ::= 'm'                         // non-generic metatype arguments are removed in the specialized function
 
   FRAGILE ::= 'q'
 
+  ASYNC-REMOVED ::= 'a'                      // async effect removed
+
   ARG-SPEC-KIND ::= 'n'                      // Unmodified argument
   ARG-SPEC-KIND ::= 'c'                      // Consumes n 'type' arguments which are closed over types in argument order
                                              // and one 'identifier' argument which is the closure symbol name

include/swift/Demangling/Demangle.h

@@ -139,16 +139,24 @@ enum class MangledDifferentiabilityKind : char {
 /// that two passes that generate similar changes do not yield the same
 /// mangling. This currently cannot happen, so this is just a safety measure
 /// that creates separate name spaces.
+///
+/// The number of entries is limited! See `Demangler::demangleSpecAttributes`.
+/// If you exceed the max, you'll need to upgrade the mangling.
 enum class SpecializationPass : uint8_t {
-  AllocBoxToStack,
+  AllocBoxToStack = 0,
   ClosureSpecializer,
   CapturePromotion,
   CapturePropagation,
   FunctionSignatureOpts,
   GenericSpecializer,
   MoveDiagnosticInOutToOut,
+  AsyncDemotion,
+  LAST = AsyncDemotion
 };
 
+constexpr uint8_t MAX_SPECIALIZATION_PASS = 10;
+static_assert((uint8_t)SpecializationPass::LAST < MAX_SPECIALIZATION_PASS);
+
 static inline char encodeSpecializationPass(SpecializationPass Pass) {
   return char(uint8_t(Pass)) + '0';
 }

include/swift/Demangling/DemangleNodes.def

@@ -368,5 +368,8 @@ NODE(HasSymbolQuery)
 NODE(OpaqueReturnTypeIndex)
 NODE(OpaqueReturnTypeParent)
 
+// Added in Swift 5.9 + 1
+NODE(AsyncRemoved)
+
 #undef CONTEXT_NODE
 #undef NODE

include/swift/SIL/GenericSpecializationMangler.h

@@ -14,6 +14,7 @@
 #define SWIFT_SIL_UTILS_GENERICSPECIALIZATIONMANGLER_H
 
 #include "swift/AST/ASTMangler.h"
+#include "swift/AST/Effects.h"
 #include "swift/Basic/NullablePtr.h"
 #include "swift/Demangling/Demangler.h"
 #include "swift/SIL/SILFunction.h"
@@ -45,6 +46,9 @@ class SpecializationMangler : public Mangle::ASTMangler {
   llvm::SmallVector<char, 32> ArgOpStorage;
   llvm::raw_svector_ostream ArgOpBuffer;
 
+  // Effects that are removed from the original function in this specialization.
+  PossibleEffects RemovedEffects;
+
 protected:
   SpecializationMangler(SpecializationPass P, IsSerialized_t Serialized,
                         SILFunction *F)

include/swift/SILOptimizer/Utils/SpecializationMangler.h

@@ -42,7 +42,7 @@ class PartialSpecializationMangler : public SpecializationMangler {
   std::string mangle();
 };
 
-// The mangler for functions where arguments are specialized.
+// The mangler for functions where arguments or effects are specialized.
 class FunctionSignatureSpecializationMangler : public SpecializationMangler {
 
   using ReturnValueModifierIntBase = uint16_t;
@@ -95,6 +95,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   FunctionSignatureSpecializationMangler(SpecializationPass Pass,
                                          IsSerialized_t Serialized,
                                          SILFunction *F);
+  // For arguments / return values
   void setArgumentConstantProp(unsigned OrigArgIdx, SILInstruction *constInst);
   void appendStringAsIdentifier(StringRef str);
 
@@ -111,6 +112,9 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   void setArgumentInOutToOut(unsigned OrigArgIdx);
   void setReturnValueOwnedToUnowned();
 
+  // For effects
+  void setRemovedEffect(EffectKind effect);
+
   std::string mangle();
 
 private:

lib/Demangling/Demangler.cpp

@@ -3286,10 +3286,13 @@ NodePointer Demangler::addFuncSpecParamNumber(NodePointer Param,
 NodePointer Demangler::demangleSpecAttributes(Node::Kind SpecKind) {
   bool metatypeParamsRemoved = nextIf('m');
   bool isSerialized = nextIf('q');
+  bool asyncRemoved = nextIf('a');
 
   int PassID = (int)nextChar() - '0';
-  if (PassID < 0 || PassID > 9)
+  if (PassID < 0 || PassID >= MAX_SPECIALIZATION_PASS) {
+    assert(false && "unexpected pass id");
     return nullptr;
+  }
 
   NodePointer SpecNd = createNode(SpecKind);
 
@@ -3300,6 +3303,10 @@ NodePointer Demangler::demangleSpecAttributes(Node::Kind SpecKind) {
     SpecNd->addChild(createNode(Node::Kind::IsSerialized),
                      *this);
 
+  if (asyncRemoved)
+    SpecNd->addChild(createNode(Node::Kind::AsyncRemoved),
+                     *this);
+
   SpecNd->addChild(createNode(Node::Kind::SpecializationPassID, PassID),
                    *this);
   return SpecNd;

lib/Demangling/NodePrinter.cpp

@@ -353,6 +353,7 @@ class NodePrinter {
     case Node::Kind::AssociatedTypeDescriptor:
     case Node::Kind::AssociatedTypeMetadataAccessor:
     case Node::Kind::AssociatedTypeWitnessTableAccessor:
+    case Node::Kind::AsyncRemoved:
     case Node::Kind::AutoClosureType:
     case Node::Kind::BaseConformanceDescriptor:
     case Node::Kind::BaseWitnessTableAccessor:
@@ -1307,6 +1308,10 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     Printer << "static ";
     print(Node->getChild(0), depth + 1);
     return nullptr;
+  case Node::Kind::AsyncRemoved:
+    Printer << "async demotion of ";
+    print(Node->getChild(0), depth + 1);
+    return nullptr;
   case Node::Kind::CurryThunk:
     Printer << "curry thunk of ";
     print(Node->getChild(0), depth + 1);

lib/Demangling/OldRemangler.cpp

@@ -446,6 +446,11 @@ ManglingError Remangler::mangleIsSerialized(Node *node, unsigned depth) {
   return ManglingError::Success;
 }
 
+ManglingError Remangler::mangleAsyncRemoved(Node *node, unsigned depth) {
+  Buffer << "a";
+  return ManglingError::Success;
+}
+
 ManglingError Remangler::mangleMetatypeParamsRemoved(Node *node, unsigned depth) {
   Buffer << "m";
   return ManglingError::Success;

lib/Demangling/Remangler.cpp

@@ -2886,6 +2886,11 @@ ManglingError Remangler::mangleIsSerialized(Node *node, unsigned depth) {
   return ManglingError::Success;
 }
 
+ManglingError Remangler::mangleAsyncRemoved(Node *node, unsigned depth) {
+  Buffer << 'a';
+  return ManglingError::Success;
+}
+
 ManglingError Remangler::mangleMetatypeParamsRemoved(Node *node, unsigned depth) {
   Buffer << 'm';
   return ManglingError::Success;

lib/SIL/Utils/GenericSpecializationMangler.cpp

@@ -21,8 +21,13 @@ using namespace Mangle;
 
 void SpecializationMangler::beginMangling() {
   ASTMangler::beginManglingWithoutPrefix();
+
   if (Serialized)
     ArgOpBuffer << 'q';
+
+  if (RemovedEffects.contains(EffectKind::Async))
+    ArgOpBuffer << 'a';
+
   ArgOpBuffer << char(uint8_t(Pass) + '0');
 }