Mangling: encode an optional unique ID in the function specialization mangling, instead of appending _unique_suffix

Author: eeckstein

docs/ABI.rst

@@ -1343,7 +1343,7 @@ The type is the function type of the specialized function.
 
 ::
 
-  specialization ::= spec-arg* 'Tf' SPEC-INFO ARG-SPEC-KIND* '_' ARG-SPEC-KIND  // Function signature specialization kind
+  specialization ::= spec-arg* 'Tf' SPEC-INFO UNIQUE-ID? ARG-SPEC-KIND* '_' ARG-SPEC-KIND  // Function signature specialization kind
 
 The ``<ARG-SPEC-KIND>`` describes how arguments are specialized.
 Some kinds need arguments, which precede ``Tf``.
@@ -1364,6 +1364,8 @@ Some kinds need arguments, which precede ``Tf``.
 
   FRAGILE ::= 'q'
 
+  UNIQUE-ID ::= NATURAL                      // Used to make unique function names
+
   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/Basic/Demangler.h

@@ -223,7 +223,8 @@ class Demangler {
   NodePointer addFuncSpecParamNumber(NodePointer Param,
                               FunctionSigSpecializationParamKind Kind);
 
-  NodePointer demangleSpecAttributes(Node::Kind SpecKind);
+  NodePointer demangleSpecAttributes(Node::Kind SpecKind,
+                                     bool demangleUniqueID = false);
 
   NodePointer demangleWitness();
   NodePointer demangleSpecialType();

include/swift/SILOptimizer/Utils/SpecializationMangler.h

@@ -150,7 +150,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   void setArgumentBoxToStack(unsigned ArgNo);
   void setReturnValueOwnedToUnowned();
 
-  std::string mangle();
+  std::string mangle(int UniqueID = 0);
 
 private:
   void mangleConstantProp(LiteralInst *LI);

lib/Basic/Demangler.cpp

@@ -1093,7 +1093,7 @@ NodePointer Demangler::demangleGenericSpecialization(Node::Kind SpecKind) {
 
 NodePointer Demangler::demangleFunctionSpecialization() {
   NodePointer Spec = demangleSpecAttributes(
-                                  Node::Kind::FunctionSignatureSpecialization);
+        Node::Kind::FunctionSignatureSpecialization, /*demangleUniqueID*/ true);
   unsigned ParamIdx = 0;
   while (Spec && !nextIf('_')) {
     Spec = addChild(Spec, demangleFuncSpecParam(ParamIdx));
@@ -1230,15 +1230,27 @@ NodePointer Demangler::addFuncSpecParamNumber(NodePointer Param,
      Node::Kind::FunctionSignatureSpecializationParamPayload, Str));
 }
 
-NodePointer Demangler::demangleSpecAttributes(Node::Kind SpecKind) {
-  NodePointer SpecNd = NodeFactory::create(SpecKind);
-  if (nextIf('q'))
-    SpecNd->addChild(NodeFactory::create(Node::Kind::SpecializationIsFragile));
+NodePointer Demangler::demangleSpecAttributes(Node::Kind SpecKind,
+                                              bool demangleUniqueID) {
+  bool isFragile = nextIf('q');
 
   int PassID = (int)nextChar() - '0';
   if (PassID < 0 || PassID > 9)
     return nullptr;
 
+  int Idx = -1;
+  if (demangleUniqueID)
+    Idx = demangleNatural();
+
+  NodePointer SpecNd;
+  if (Idx >= 0) {
+    SpecNd = NodeFactory::create(SpecKind, Idx);
+  } else {
+    SpecNd = NodeFactory::create(SpecKind);
+  }
+  if (isFragile)
+    SpecNd->addChild(NodeFactory::create(Node::Kind::SpecializationIsFragile));
+
   SpecNd->addChild(NodeFactory::create(Node::Kind::SpecializationPassID,
                                        PassID));
   return SpecNd;

lib/Basic/Mangler.cpp

@@ -269,6 +269,11 @@ void NewMangling::printManglingStats() {
 }
 
 void Mangler::beginMangling() {
+  Storage.clear();
+  Substitutions.clear();
+  StringSubstitutions.clear();
+  lastSubstIdx = -2;
+  Words.clear();
   Buffer << MANGLING_PREFIX_STR;
 }
 

lib/Basic/Remangler.cpp

@@ -860,6 +860,11 @@ void Remangler::mangleFunctionSignatureSpecialization(Node *node) {
       }
     }
     mangle(Child.get());
+
+    if (Child->getKind() == Node::Kind::SpecializationPassID &&
+        node->hasIndex()) {
+      Buffer << node->getIndex();
+    }
   }
   if (!returnValMangled)
     Buffer << "_n";

lib/SILOptimizer/Transforms/FunctionSignatureOpts.cpp

@@ -364,8 +364,16 @@ std::string FunctionSignatureTransform::createOptimizedSILFunctionName() {
   }
 
   OldFM.mangle();
-  std::string Old = OldFM.getMangler().finalize();
-  std::string New = NewFM.mangle();
+  SILModule &M = F->getModule();
+  std::string Old = getUniqueName(OldFM.getMangler().finalize(), M);
+
+  int UniqueID = 0;
+  std::string New;
+  do {
+    New = NewFM.mangle(UniqueID);
+    ++UniqueID;
+  } while (M.lookUpFunction(New));
+
   return NewMangling::selectMangling(Old, New);
 }
 
@@ -472,7 +480,7 @@ CanSILFunctionType FunctionSignatureTransform::createOptimizedSILFunctionType()
 void FunctionSignatureTransform::createFunctionSignatureOptimizedFunction() {
   // Create the optimized function !
   SILModule &M = F->getModule();
-  std::string Name = getUniqueName(createOptimizedSILFunctionName(), M);
+  std::string Name = createOptimizedSILFunctionName();
   SILLinkage linkage = F->getLinkage();
   if (isAvailableExternally(linkage))
     linkage = SILLinkage::Shared;

lib/SILOptimizer/Utils/SpecializationMangler.cpp

@@ -309,9 +309,13 @@ mangleReturnValue(ReturnValueModifierIntBase RetMod) {
   }
 }
 
-std::string FunctionSignatureSpecializationMangler::mangle() {
+std::string FunctionSignatureSpecializationMangler::mangle(int UniqueID) {
+  ArgOpStorage.clear();
   beginMangling();
 
+  if (UniqueID)
+    ArgOpBuffer << UniqueID;
+
   for (unsigned i : indices(Args)) {
     ArgumentModifierIntBase ArgMod;
     NullablePtr<SILInstruction> Inst;