AST: Fix derivation of conformances in Substitution::subst() in the p…

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;
@@ -28,11 +29,30 @@ namespace swift {
   class ArchetypeType;
   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,8 +98,39 @@ 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,38 +24,9 @@
 
 using namespace swift;
 
-bool Substitution::operator==(const Substitution &other) const {
-  // The archetypes may be missing, but we can compare them directly
-  // because archetypes are always canonical.
-  return
-    Replacement->getCanonicalType() == other.Replacement->getCanonicalType() &&
-    Conformance.equals(other.Conformance);
-}
-
-Substitution::Substitution(Type Replacement,
-                           ArrayRef<ProtocolConformanceRef> Conformance)
-  : Replacement(Replacement), Conformance(Conformance)
-{
-  // The replacement type must be materializable.
-  assert(Replacement->isMaterializable()
-         && "cannot substitute with a non-materializable type");
-}
-
-Substitution Substitution::subst(Module *module,
-                                 GenericSignature *sig,
-                                 GenericEnvironment *env,
-                                 ArrayRef<Substitution> subs) const {
-  TypeSubstitutionMap subMap;
-  ArchetypeConformanceMap conformanceMap;
-
-  assert(sig && env);
-  env->getSubstitutionMap(module, sig, subs, subMap, conformanceMap);
-  return subst(module, subMap, conformanceMap);
-}
-
-static Optional<ProtocolConformanceRef>
+Optional<ProtocolConformanceRef> ArchetypeConformanceMap::
 lookupArchetypeConformance(ProtocolDecl *proto,
-                           ArrayRef<ProtocolConformanceRef> conformances) {
+                           ArrayRef<ProtocolConformanceRef> conformances) const {
   for (ProtocolConformanceRef found : conformances) {
     auto foundProto = found.getRequirement();
     if (foundProto == proto) {
@@ -73,40 +44,93 @@ lookupArchetypeConformance(ProtocolDecl *proto,
   return None;
 }
 
-static Optional<ProtocolConformanceRef>
+Optional<ProtocolConformanceRef> ArchetypeConformanceMap::
 lookupArchetypeConformance(ArchetypeType *replacement,
-                           ProtocolDecl *proto,
-                           ArchetypeConformanceMap &conformanceMap) {
+                           ProtocolDecl *proto) const {
   // 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)) {
+  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 for the parent.
-  if (auto *parent = replacement->getParent()) {
-    auto *assocType = replacement->getAssocType();
-    auto *parentProto = assocType->getProtocol();
+  // 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, parentProto, conformanceMap);
+        lookupArchetypeConformance(parent.first, parentProto);
 
     if (conformance) {
       if (!conformance->isConcrete())
         return ProtocolConformanceRef(proto);
 
       auto sub = conformance->getConcrete()->getTypeWitnessSubstAndDecl(
-          assocType, nullptr).first;
+          parent.second, nullptr).first;
 
-      return lookupArchetypeConformance(proto, sub.getConformances());
+      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.
+  return
+    Replacement->getCanonicalType() == other.Replacement->getCanonicalType() &&
+    Conformance.equals(other.Conformance);
+}
+
+Substitution::Substitution(Type Replacement,
+                           ArrayRef<ProtocolConformanceRef> Conformance)
+  : Replacement(Replacement), Conformance(Conformance)
+{
+  // The replacement type must be materializable.
+  assert(Replacement->isMaterializable()
+         && "cannot substitute with a non-materializable type");
+}
+
+Substitution Substitution::subst(Module *module,
+                                 GenericSignature *sig,
+                                 GenericEnvironment *env,
+                                 ArrayRef<Substitution> subs) const {
+  TypeSubstitutionMap subMap;
+  ArchetypeConformanceMap conformanceMap;
+
+  assert(sig && env);
+  env->getSubstitutionMap(module, sig, subs, subMap, conformanceMap);
+  return subst(module, subMap, conformanceMap);
+}
+
 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