Mangler: Correctly record substitutions when dropping protocols from associated types

Author: slavapestov

include/swift/AST/ASTMangler.h

@@ -76,6 +76,15 @@ class ASTMangler : public Mangler {
   ASTMangler(bool DWARFMangling = false)
     : DWARFMangling(DWARFMangling) {}
 
+  void addTypeSubstitution(Type type) {
+    type = dropProtocolsFromAssociatedTypes(type);
+    addSubstitution(type.getPointer());
+  }
+  bool tryMangleTypeSubstitution(Type type) {
+    type = dropProtocolsFromAssociatedTypes(type);
+    return tryMangleSubstitution(type.getPointer());
+  }
+
   std::string mangleClosureEntity(const AbstractClosureExpr *closure,
                                   SymbolKind SKind);
 
@@ -261,6 +270,9 @@ class ASTMangler : public Mangler {
                                    unsigned initialParamDepth,
                                    ArrayRef<Requirement> requirements);
 
+  DependentMemberType *dropProtocolFromAssociatedType(DependentMemberType *dmt);
+  Type dropProtocolsFromAssociatedTypes(Type type);
+
   void appendAssociatedTypeName(DependentMemberType *dmt);
 
   void appendClosureEntity(const SerializedAbstractClosureExpr *closure);

lib/AST/ASTMangler.cpp

@@ -737,15 +737,15 @@ void ASTMangler::appendType(Type type) {
 
       if (aliasTy->getSubstitutionMap().hasAnySubstitutableParams()) {
         // Try to mangle the entire name as a substitution.
-        if (tryMangleSubstitution(tybase))
+        if (tryMangleTypeSubstitution(tybase))
           return;
 
         appendAnyGenericType(decl);
         bool isFirstArgList = true;
         appendBoundGenericArgs(type, isFirstArgList);
         appendRetroactiveConformances(type);
         appendOperator("G");
-        addSubstitution(type.getPointer());
+        addTypeSubstitution(type);
         return;
       }
 
@@ -840,7 +840,7 @@ void ASTMangler::appendType(Type type) {
         Decl = type->getAnyGeneric();
       if (shouldMangleAsGeneric(type)) {
         // Try to mangle the entire name as a substitution.
-        if (tryMangleSubstitution(tybase))
+        if (tryMangleTypeSubstitution(tybase))
           return;
 
         if (isStdlibType(Decl) && Decl->getName().str() == "Optional") {
@@ -855,7 +855,7 @@ void ASTMangler::appendType(Type type) {
           appendRetroactiveConformances(type);
           appendOperator("G");
         }
-        addSubstitution(type.getPointer());
+        addTypeSubstitution(type);
         return;
       }
       appendAnyGenericType(type->getAnyGeneric());
@@ -901,7 +901,7 @@ void ASTMangler::appendType(Type type) {
 
     case TypeKind::DependentMember: {
       auto *DepTy = cast<DependentMemberType>(tybase);
-      if (tryMangleSubstitution(DepTy))
+      if (tryMangleTypeSubstitution(DepTy))
         return;
 
       bool isAssocTypeAtDepth = false;
@@ -920,7 +920,7 @@ void ASTMangler::appendType(Type type) {
         appendAssociatedTypeName(DepTy);
         appendOperator("qa");
       }
-      addSubstitution(DepTy);
+      addTypeSubstitution(DepTy);
       return;
     }
 
@@ -1633,7 +1633,7 @@ void ASTMangler::appendAnyGenericType(const GenericTypeDecl *decl) {
   }
 
   // Try to mangle the entire name as a substitution.
-  if (tryMangleSubstitution(key.getPointer()))
+  if (tryMangleTypeSubstitution(key))
     return;
 
   // Try to mangle a symbolic reference for a nominal type.
@@ -1642,7 +1642,7 @@ void ASTMangler::appendAnyGenericType(const GenericTypeDecl *decl) {
     if (nominal && (!CanSymbolicReference || CanSymbolicReference(nominal))) {
       appendSymbolicReference(nominal);
       // Substitutions can refer back to the symbolic reference.
-      addSubstitution(key.getPointer());
+      addTypeSubstitution(key);
       return;
     }
   }
@@ -1705,7 +1705,7 @@ void ASTMangler::appendAnyGenericType(const GenericTypeDecl *decl) {
     }
   }
 
-  addSubstitution(key.getPointer());
+  addTypeSubstitution(key);
 }
 
 void ASTMangler::appendFunction(AnyFunctionType *fn, bool isFunctionMangling) {
@@ -1941,7 +1941,7 @@ void ASTMangler::appendRequirement(const Requirement &reqt) {
 
   if (auto *DT = FirstTy->getAs<DependentMemberType>()) {
     bool isAssocTypeAtDepth = false;
-    if (tryMangleSubstitution(DT)) {
+    if (tryMangleTypeSubstitution(DT)) {
       switch (reqt.getKind()) {
         case RequirementKind::Conformance:
           return appendOperator("RQ");
@@ -1957,7 +1957,7 @@ void ASTMangler::appendRequirement(const Requirement &reqt) {
       llvm_unreachable("bad requirement type");
     }
     GenericTypeParamType *gpBase = appendAssocType(DT, isAssocTypeAtDepth);
-    addSubstitution(DT);
+    addTypeSubstitution(DT);
     assert(gpBase);
     switch (reqt.getKind()) {
       case RequirementKind::Conformance:
@@ -2043,24 +2043,54 @@ void ASTMangler::appendGenericSignatureParts(
   appendOperator("r", StringRef(OpStorage.data(), OpStorage.size()));
 }
 
+// If the base type is known to have a single protocol conformance
+// in the current generic context, then we don't need to disambiguate the
+// associated type name by protocol.
+DependentMemberType *
+ASTMangler::dropProtocolFromAssociatedType(DependentMemberType *dmt) {
+  auto baseTy = dmt->getBase();
+  bool unambiguous = (!dmt->getAssocType() ||
+                      CurGenericSignature->getConformsTo(baseTy).size() <= 1);
+
+  if (auto *baseDMT = baseTy->getAs<DependentMemberType>())
+    baseTy = dropProtocolFromAssociatedType(baseDMT);
+
+  if (unambiguous)
+    return DependentMemberType::get(baseTy, dmt->getName());
+
+  return DependentMemberType::get(baseTy, dmt->getAssocType());
+}
+
+Type
+ASTMangler::dropProtocolsFromAssociatedTypes(Type type) {
+  if (!OptimizeProtocolNames || !CurGenericSignature)
+    return type;
+
+  if (!type->hasDependentMember())
+    return type;
+
+  return type.transform([&](Type t) -> Type {
+    if (auto *dmt = dyn_cast<DependentMemberType>(t.getPointer()))
+      return dropProtocolFromAssociatedType(dmt);
+    return t;
+  });
+}
+
 void ASTMangler::appendAssociatedTypeName(DependentMemberType *dmt) {
-  auto assocTy = dmt->getAssocType();
-
-  // If the base type is known to have a single protocol conformance
-  // in the current generic context, then we don't need to disambiguate the
-  // associated type name by protocol.
-  // This can result in getting the same mangled string for different
-  // DependentMemberTypes. This is not a problem but re-mangling might do more
-  // aggressive substitutions, which means that the re-mangled name may differ
-  // from the original mangled name.
-  // FIXME: We ought to be able to get to the generic signature from a
-  // dependent type, but can't yet. Shouldn't need this side channel.
-
-  appendIdentifier(assocTy->getName().str());
-  if (!OptimizeProtocolNames || !CurGenericSignature
-      || CurGenericSignature->getConformsTo(dmt->getBase()).size() > 1) {
-    appendAnyGenericType(assocTy->getProtocol());
+  if (auto assocTy = dmt->getAssocType()) {
+    appendIdentifier(assocTy->getName().str());
+
+    // If the base type is known to have a single protocol conformance
+    // in the current generic context, then we don't need to disambiguate the
+    // associated type name by protocol.
+    if (!OptimizeProtocolNames || !CurGenericSignature
+        || CurGenericSignature->getConformsTo(dmt->getBase()).size() > 1) {
+      appendAnyGenericType(assocTy->getProtocol());
+    }
+    return;
   }
+
+  appendIdentifier(dmt->getName().str());
 }
 
 void ASTMangler::appendClosureEntity(

test/SILGen/keypath_witness_overrides.swift

@@ -13,8 +13,8 @@ where OS.Index: Hashable {
   // CHECK-SAME: id #OverridesSetter.subscript!getter.1
   // CHECK-SAME: getter @$s18protocol_overrides14OriginalGetterPy7ElementQz5IndexQzcipAA15OverridesSetterRzSHAGRQlxxTK
   // CHECK-SAME: setter @$s18protocol_overrides10AddsSetterPy7ElementQz5IndexQzcipAA09OverridesD0RzSHAGRQlxxTk
-  // CHECK-SAME: indices_equals @$s5Index18protocol_overrides14OriginalGetterPQzAB15OverridesSetterRzSHAERQlTH
-  // CHECK-SAME: indices_hash @$s5Index18protocol_overrides14OriginalGetterPQzAB15OverridesSetterRzSHAERQlTh
+  // CHECK-SAME: indices_equals @$s5Index18protocol_overrides14OriginalGetterPQzAB15OverridesSetterRzSHAARpzlTH
+  // CHECK-SAME: indices_hash @$s5Index18protocol_overrides14OriginalGetterPQzAB15OverridesSetterRzSHAARpzlTh
   let keypath = \OS.[index]
   return keypath
 }