Add mangling support for partial specializations.

Author: swiftix

include/swift/SIL/Mangle.h

@@ -30,6 +30,8 @@ class AbstractClosureExpr;
 enum class SpecializationKind : uint8_t {
   Generic,
   NotReAbstractedGeneric,
+  Partial,
+  NotReAbstractedPartial,
   FunctionSignature,
 };
 
@@ -73,6 +75,12 @@ class SpecializationManglerBase {
     case SpecializationKind::NotReAbstractedGeneric:
       M.append("r");
       break;
+    case SpecializationKind::Partial:
+      M.append("p");
+      break;
+    case SpecializationKind::NotReAbstractedPartial:
+      M.append("P");
+      break;
     case SpecializationKind::FunctionSignature:
       M.append("f");
       break;
@@ -145,14 +153,44 @@ class GenericSpecializationMangler :
   };
 
   GenericSpecializationMangler(Mangle::Mangler &M, SILFunction *F,
-                               SubstitutionList Subs,
+                               SubstitutionList Subs, IsFragile_t Fragile,
+                               ReAbstractionMode isReAbstracted = ReAbstracted)
+      : SpecializationMangler(isReAbstracted == ReAbstracted
+                                  ? SpecializationKind::Generic
+                                  : SpecializationKind::NotReAbstractedGeneric,
+                              SpecializationPass::GenericSpecializer, M,
+                              Fragile, F),
+        Subs(Subs) {
+  }
+
+private:
+  void mangleSpecialization();
+};
+
+class PartialSpecializationMangler :
+  public SpecializationMangler<PartialSpecializationMangler> {
+
+  friend class SpecializationMangler<PartialSpecializationMangler>;
+
+  CanSILFunctionType SpecializedFnTy;
+public:
+
+  enum ReAbstractionMode {
+    ReAbstracted,
+    NotReabstracted
+  };
+
+  PartialSpecializationMangler(Mangle::Mangler &M,
+                               SILFunction *F,
+                               CanSILFunctionType SpecializedFnTy,
                                IsFragile_t Fragile,
                                ReAbstractionMode isReAbstracted = ReAbstracted)
     : SpecializationMangler(isReAbstracted == ReAbstracted ?
-                              SpecializationKind::Generic :
-                              SpecializationKind::NotReAbstractedGeneric,
+                              SpecializationKind::Partial :
+                              SpecializationKind::NotReAbstractedPartial,
                             SpecializationPass::GenericSpecializer,
-                            M, Fragile, F), Subs(Subs) {}
+                            M, Fragile, F), SpecializedFnTy(SpecializedFnTy) {
+    }
 
 private:
   void mangleSpecialization();

lib/SIL/Mangle.cpp

@@ -55,7 +55,7 @@ static void mangleSubstitution(Mangler &M, Substitution Sub) {
 
 void GenericSpecializationMangler::mangleSpecialization() {
   Mangler &M = getMangler();
-
+  // This is a full specialization.
   SILFunctionType *FTy = Function->getLoweredFunctionType();
   CanGenericSignature Sig = FTy->getGenericSignature();
 
@@ -73,6 +73,25 @@ void GenericSpecializationMangler::mangleSpecialization() {
   assert(idx == Subs.size() && "subs not parallel to dependent types");
 }
 
+void PartialSpecializationMangler::mangleSpecialization() {
+  Mangler &M = getMangler();
+  // If the only change to the generic signature during specialization is
+  // addition of new same-type requirements, which happens in case of a
+  // full specialization, it would be enough to mangle only the substitutions.
+  //
+  // If the types of function arguments have not changed, but some new
+  // conformances were added to the generic parameters, e.g. in case of
+  // a pre-specialization, then it would be enough to mangle only the new
+  // generic signature.
+  //
+  // If the types of function arguments have changed as a result of a partial
+  // specialization, we need to mangle the entire new function type.
+
+  // This is a partial specialization.
+  M.mangleType(SpecializedFnTy, 0);
+  M.append("_");
+}
+
 //===----------------------------------------------------------------------===//
 //                      Function Signature Optimizations
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Utils/Generics.cpp

@@ -294,16 +294,35 @@ GenericFuncSpecializer::GenericFuncSpecializer(SILFunction *GenericFunc,
       ReInfo(ReInfo) {
 
   assert(GenericFunc->isDefinition() && "Expected definition to specialize!");
+  auto FnTy = ReInfo.getSpecializedType();
 
-  Mangle::Mangler Mangler;
-  GenericSpecializationMangler OldGenericMangler(Mangler, GenericFunc,
-                                              ParamSubs, Fragile);
-  OldGenericMangler.mangle();
-  std::string Old = Mangler.finalize();
+  std::string Old;
+
+  if (ReInfo.isPartialSpecialization()) {
+    Mangle::Mangler Mangler;
+    PartialSpecializationMangler OldGenericMangler(Mangler, GenericFunc, FnTy,
+                                                   Fragile);
+    OldGenericMangler.mangle();
+    Old = Mangler.finalize();
+  } else {
+    Mangle::Mangler Mangler;
+    GenericSpecializationMangler OldGenericMangler(Mangler, GenericFunc,
+                                                   ParamSubs, Fragile);
+    OldGenericMangler.mangle();
+    Old = Mangler.finalize();
+  }
+
+  std::string New;
+  if (ReInfo.isPartialSpecialization()) {
+    NewMangling::PartialSpecializationMangler NewGenericMangler(
+        GenericFunc, FnTy, Fragile, /*isReAbstracted*/ true);
+    New = NewGenericMangler.mangle();
+  } else {
+    NewMangling::GenericSpecializationMangler NewGenericMangler(
+        GenericFunc, ParamSubs, Fragile, /*isReAbstracted*/ true);
+    New = NewGenericMangler.mangle();
+  }
 
-  NewMangling::GenericSpecializationMangler NewGenericMangler(GenericFunc,
-                          ParamSubs, Fragile, /*isReAbstracted*/ true);
-  std::string New = NewGenericMangler.mangle();
   ClonedName = NewMangling::selectMangling(Old, New);
 
   DEBUG(llvm::dbgs() << "    Specialized function " << ClonedName << '\n');