AST: Fix derivation of conformances in Substitution::subst() in the presence of same-type constraints

Applying a substitution list to a Substitution is done in two
steps:

1) First, apply the substitution list to the original replacement
   type R to get the substituted replacement type R'.

2) Second, for each abstract conformance of R from the original
   substitution, look up a concrete conformance of R' from the
   substitution list.

With minimized generic signatures, we would omit conformances of
a nested type T.A if they could be derived from some other
conformance already in the substitution list.

However, the derivation process was incorrect, because it would
only look at the canonical parent type T, and not any other parent
type that had a child A which was same-type equivalent to T.A.

For example, consider the following code:

protocol P1 { associatedtype A : P3 }
protocol P2 { associatedtype A : P4 }

protocol P3 {}
protocol P4 {}

func doSomething<T : P4>(...) {}
func doSomethingElse<T1 : P1, T2 : P2>(...) where T1.A == T2.A {
  doSomething(...)
}

If we specialize doSomethingElse() with a pair of concrete types
to replace T1 and T2, we would need to find a concrete conformance
to replace the abstract conformance T2.A : P4 in the call to
doSomething().

Since the conformance of T2.A : P4 is a redundant requirement,
it does not appear in the conformance map; furthermore, T1.A and
T2.A are same-type equivalent, so they map to the same archetype.

We would therefore look at the canonical parent of T2.A, which is
T1, and look up the conformance of T1 : P1 in the substitution list.

However, the only requirement P1 imposes on A is the conformance
A : P3. There's no conformance A : P4 inside T1 : P1, so we would
hit an assertion.

Indeed, the conformance T1.A : P4 must be recovered from the
conformance of T2 : P2, because P2 requires that A : P4.

This patch ensures that the conformance can be found by changing
the ArchetypeConformanceMap from a simple mapping of archetypes
to conformances into a structure that records same-type constraints
as well.

So instead of just looking at the canonical parent of the archetype
T1.A, we consult the structure to check all archetypes that have
T1.A as a child, in this case both T1 and T2.

T2 : P2 contains the conformance we need, allowing the above code
to be specialized correctly.

include/swift/AST/Substitution.h

@@ -19,6 +19,7 @@
 
 #include "swift/AST/Type.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
 
 namespace llvm {
   class raw_ostream;
@@ -29,10 +30,29 @@ namespace swift {
   class GenericEnvironment;
   class ProtocolConformanceRef;
 
-/// DenseMap type used internally by Substitution::subst to track conformances
-/// applied to archetypes.
-using ArchetypeConformanceMap
-  = llvm::DenseMap<ArchetypeType*, ArrayRef<ProtocolConformanceRef>>;
+/// Data structure type used internally by Substitution::subst to track
+/// conformances applied to archetypes.
+class ArchetypeConformanceMap {
+  using ParentType = std::pair<ArchetypeType *, AssociatedTypeDecl *>;
+
+  llvm::DenseMap<ArchetypeType *, ArrayRef<ProtocolConformanceRef>> map;
+  llvm::DenseMap<ArchetypeType *, SmallVector<ParentType, 1>> parents;
+
+  Optional<ProtocolConformanceRef>
+  lookupArchetypeConformance(ProtocolDecl *proto,
+                             ArrayRef<ProtocolConformanceRef> conformances) const;
+
+public:
+  Optional<ProtocolConformanceRef>
+  lookupArchetypeConformance(ArchetypeType *replacement,
+                             ProtocolDecl *proto) const;
+
+  void addArchetypeConformances(ArchetypeType *replacement,
+                                ArrayRef<ProtocolConformanceRef> conformances);
+
+  void addArchetypeParent(ArchetypeType *replacement, ArchetypeType *parent,
+                          AssociatedTypeDecl *assocType);
+};
 
 /// Substitution - A substitution into a generic specialization.
 class Substitution {

lib/AST/GenericEnvironment.cpp

@@ -98,7 +98,38 @@ getSubstitutionMap(ModuleDecl *mod,
 
     // Record the replacement type and its conformances.
     subsMap[archetype] = sub.getReplacement();
-    conformanceMap[archetype] = sub.getConformances();
+    conformanceMap.addArchetypeConformances(archetype, sub.getConformances());
+  }
+
+  for (auto reqt : sig->getRequirements()) {
+    if (reqt.getKind() != RequirementKind::SameType)
+      continue;
+
+    auto first = reqt.getFirstType()->getAs<DependentMemberType>();
+    auto second = reqt.getSecondType()->getAs<DependentMemberType>();
+
+    if (!first || !second)
+      continue;
+
+    auto archetype = mapTypeIntoContext(mod, first)->getAs<ArchetypeType>();
+    if (!archetype)
+      continue;
+
+    auto firstBase = first->getBase();
+    auto secondBase = second->getBase();
+
+    auto firstBaseArchetype = mapTypeIntoContext(mod, firstBase)->getAs<ArchetypeType>();
+    auto secondBaseArchetype = mapTypeIntoContext(mod, secondBase)->getAs<ArchetypeType>();
+
+    if (!firstBaseArchetype || !secondBaseArchetype)
+      continue;
+
+    if (archetype->getParent() != firstBaseArchetype)
+      conformanceMap.addArchetypeParent(archetype, firstBaseArchetype,
+                                        first->getAssocType());
+    if (archetype->getParent() != secondBaseArchetype)
+      conformanceMap.addArchetypeParent(archetype, secondBaseArchetype,
+                                        second->getAssocType());
   }
 
   assert(subs.empty() && "did not use all substitutions?!");

lib/AST/Substitution.cpp

@@ -24,6 +24,84 @@
 
 using namespace swift;
 
+Optional<ProtocolConformanceRef> ArchetypeConformanceMap::
+lookupArchetypeConformance(ProtocolDecl *proto,
+                           ArrayRef<ProtocolConformanceRef> conformances) const {
+  for (ProtocolConformanceRef found : conformances) {
+    auto foundProto = found.getRequirement();
+    if (foundProto == proto) {
+      return found;
+    } else if (foundProto->inheritsFrom(proto)) {
+      if (found.isConcrete()) {
+        return ProtocolConformanceRef(
+          found.getConcrete()->getInheritedConformance(proto));
+      }
+
+      return found;
+    }
+  }
+
+  return None;
+}
+
+Optional<ProtocolConformanceRef> ArchetypeConformanceMap::
+lookupArchetypeConformance(ArchetypeType *replacement,
+                           ProtocolDecl *proto) const {
+  // Check for conformances for the type that apply to the original
+  // substituted archetype.
+  auto foundReplacement = map.find(replacement);
+  if (foundReplacement != map.end()) {
+    auto substReplacement = foundReplacement->second;
+    if (auto conformance = lookupArchetypeConformance(proto, substReplacement))
+      return conformance;
+  }
+
+  // Check if we have substitutions from one of our parent archetypes.
+  auto foundParents = parents.find(replacement);
+  if (foundParents == parents.end())
+    return None;
+
+  for (auto parent : foundParents->second) {
+    auto *parentProto = parent.second->getProtocol();
+    auto conformance =
+        lookupArchetypeConformance(parent.first, parentProto);
+
+    if (conformance) {
+      if (!conformance->isConcrete())
+        return ProtocolConformanceRef(proto);
+
+      auto sub = conformance->getConcrete()->getTypeWitnessSubstAndDecl(
+          parent.second, nullptr).first;
+
+      if (auto result = lookupArchetypeConformance(proto, sub.getConformances()))
+        return result;
+    }
+  }
+
+  return None;
+}
+
+void ArchetypeConformanceMap::
+addArchetypeConformances(ArchetypeType *replacement,
+                         ArrayRef<ProtocolConformanceRef> conformances) {
+  assert(replacement);
+
+  auto result = map.insert(std::make_pair(replacement, conformances));
+  assert(result.second);
+  (void) result;
+
+  if (auto *parent = replacement->getParent())
+    addArchetypeParent(replacement, parent, replacement->getAssocType());
+}
+
+void ArchetypeConformanceMap::
+addArchetypeParent(ArchetypeType *replacement,
+                   ArchetypeType *parent,
+                   AssociatedTypeDecl *assocType) {
+  assert(replacement && parent && assocType);
+  parents[replacement].push_back(std::make_pair(parent, assocType));
+}
+
 bool Substitution::operator==(const Substitution &other) const {
   // The archetypes may be missing, but we can compare them directly
   // because archetypes are always canonical.
@@ -53,60 +131,6 @@ Substitution Substitution::subst(Module *module,
   return subst(module, subMap, conformanceMap);
 }
 
-static Optional<ProtocolConformanceRef>
-lookupArchetypeConformance(ProtocolDecl *proto,
-                           ArrayRef<ProtocolConformanceRef> conformances) {
-  for (ProtocolConformanceRef found : conformances) {
-    auto foundProto = found.getRequirement();
-    if (foundProto == proto) {
-      return found;
-    } else if (foundProto->inheritsFrom(proto)) {
-      if (found.isConcrete()) {
-        return ProtocolConformanceRef(
-          found.getConcrete()->getInheritedConformance(proto));
-      }
-
-      return found;
-    }
-  }
-
-  return None;
-}
-
-static Optional<ProtocolConformanceRef>
-lookupArchetypeConformance(ArchetypeType *replacement,
-                           ProtocolDecl *proto,
-                           ArchetypeConformanceMap &conformanceMap) {
-  // Check for conformances for the type that apply to the original
-  // substituted archetype.
-  auto it = conformanceMap.find(replacement);
-  if (it != conformanceMap.end()) {
-    if (auto conformance = lookupArchetypeConformance(proto, it->second)) {
-      return conformance;
-    }
-  }
-
-  // Check if we have substitutions for the parent.
-  if (auto *parent = replacement->getParent()) {
-    auto *assocType = replacement->getAssocType();
-    auto *parentProto = assocType->getProtocol();
-    auto conformance =
-        lookupArchetypeConformance(parent, parentProto, conformanceMap);
-
-    if (conformance) {
-      if (!conformance->isConcrete())
-        return ProtocolConformanceRef(proto);
-
-      auto sub = conformance->getConcrete()->getTypeWitnessSubstAndDecl(
-          assocType, nullptr).first;
-
-      return lookupArchetypeConformance(proto, sub.getConformances());
-    }
-  }
-
-  return None;
-}
-
 Substitution Substitution::subst(Module *module,
                                  TypeSubstitutionMap &subMap,
                                  ArchetypeConformanceMap &conformanceMap) const {
@@ -143,8 +167,8 @@ Substitution Substitution::subst(Module *module,
 
     // If the original type was an archetype, check the conformance map.
     if (auto replacementArch = Replacement->getAs<ArchetypeType>()) {
-      conformance = lookupArchetypeConformance(replacementArch, proto,
-                                               conformanceMap);
+      conformance = conformanceMap.lookupArchetypeConformance(
+          replacementArch, proto);
     }
 
     // If that didn't find anything, we can still synthesize AnyObject

test/SILOptimizer/specialize_same_type_constraint.swift

@@ -0,0 +1,53 @@
+// RUN: %target-swift-frontend  -O -emit-sil -primary-file %s | %FileCheck %s
+
+protocol FirstChild {}
+
+protocol FirstParent {
+  associatedtype Child : FirstChild
+
+  var firstChild: Child { get }
+}
+
+protocol SecondChild {}
+
+protocol SecondParent {
+  associatedtype Child : SecondChild
+
+  var secondChild: Child { get }
+}
+
+@_semantics("optimize.sil.never")
+func takesFirstChild<T : FirstChild>(t: T) {}
+
+@_semantics("optimize.sil.never")
+func takesSecondChild<T : SecondChild>(t: T) {}
+
+@inline(never)
+func doStuff<First : FirstParent, Second : SecondParent>(f: First, s: Second)
+    where First.Child == Second.Child {
+  takesFirstChild(t: f.firstChild)
+  takesSecondChild(t: f.firstChild)
+
+  takesFirstChild(t: s.secondChild)
+  takesSecondChild(t: s.secondChild)
+}
+
+struct ConcreteChild : FirstChild, SecondChild {}
+
+struct ConcreteFirstParent<T> : FirstParent {
+  var firstChild: ConcreteChild { return ConcreteChild() }
+}
+
+struct ConcreteSecondParent<T> : SecondParent {
+  var secondChild: ConcreteChild { return ConcreteChild() }
+}
+
+doStuff(f: ConcreteFirstParent<ConcreteChild>(),
+        s: ConcreteSecondParent<ConcreteChild>())
+
+// CHECK-LABEL: sil shared [noinline] @_TTSf4d_d___TTSg5GV31specialize_same_type_constraint19ConcreteFirstParentVS_13ConcreteChild_GS0_S1__S_11FirstParentS__GVS_20ConcreteSecondParentS1__GS3_S1__S_12SecondParentS____TF31specialize_same_type_constraint7doStuffu0_RxS_11FirstParent_S_12SecondParentwx5Childzw_5ChildrFT1fx1sq__T_
+// CHECK: [[FIRST:%.*]] = function_ref @_TF31specialize_same_type_constraint15takesFirstChilduRxS_10FirstChildrFT1tx_T_
+// CHECK: apply [[FIRST]]<ConcreteChild>({{.*}}) : $@convention(thin) <τ_0_0 where τ_0_0 : FirstChild> (@in τ_0_0) -> ()
+// CHECK: [[SECOND:%.*]] = function_ref @_TF31specialize_same_type_constraint16takesSecondChilduRxS_11SecondChildrFT1tx_T_
+// CHECK: apply [[SECOND]]<ConcreteChild>({{.*}}) : $@convention(thin) <τ_0_0 where τ_0_0 : SecondChild> (@in τ_0_0) -> ()
+// CHECK: return