RequirementMachine: Fix silly oversight when resolving concrete types and superclass bounds

The property map stores the concrete type or superclass bound for all
terms whose suffix is equal to some key, so eg if you have

  protocol P {
    associatedtype T where T == U?
    associatedtype U
  }

Then you have this rewrite system

  [P].T => [P:T]
  [P].U => [P:U]
  [P:T].[concrete: Optional<τ_0_0> with <[P:U]>] => [P:T]

Which creates this property map

  [P:T] => { [concrete: Optional<τ_0_0> with <[P:U]>] }

Now if I start with a generic signature like

  <T, U where U : P>

This produces the following rewrite system:

  τ_0_1.[U] => τ_0_1
  τ_0_1.T => τ_0_1.[P:T]
  τ_0_1.U => τ_0_1.[P:U]

Consider the computation of the canonical type of τ_0_1.T. This term
reduces to τ_0_1.[P:T]. The suffix [P:T] has a concrete type symbol in
the property map, [concrete: Optional<τ_0_0> with <[P:U]>].

However it would not be correct to canonicalize τ_0_1.[P:T] to
Optional<τ_0_0>.subst { τ_0_0 => getTypeForTerm([P:T]) }; this
produces the type Optional<τ_0_0.T>, and not Optional<τ_0_1.T> as
expected.

The reason is that the substitution τ_0_0 => getTypeForTerm([P:T])
is "relative" to the protocol Self type of P, since domain([P:T]) = {P}.

Indeed, the right solution here is to note that τ_0_1.[P:T] has a suffix
equal to the key of the property map entry, [P:T]. If we strip off the
suffix, we get τ_0_1. If we then prepend τ_0_1 to the substitution term,
we get the term τ_0_1.[P:U], whose canonical type is τ_0_1.U.

Now, applying the substitution τ_0_0 => τ_0_1.U to Optional<τ_0_0>
produces the desired result, Optional<τ_0_1.U>.

Note that this is the same "prepend a prefix to each substitution of
a concrete type symbol" operation that is performed in checkForOverlap()
and PropertyBag::copyPropertiesFrom(), however we can simplify things
slightly by open-coding it instead of calling the utility method
prependPrefixToConcreteSubstitutions(), since the latter creates a
new symbol, which we don't actually need.

Author: slavapestov

lib/AST/RequirementMachine/GenericSignatureQueries.cpp

@@ -47,12 +47,12 @@ RequirementMachine::getLocalRequirements(
     return result;
 
   if (props->isConcreteType()) {
-    result.concreteType = props->getConcreteType({}, protos, Context);
+    result.concreteType = props->getConcreteType({}, term, protos, Context);
     return result;
   }
 
   if (props->hasSuperclassBound()) {
-    result.superclass = props->getSuperclassBound({}, protos, Context);
+    result.superclass = props->getSuperclassBound({}, term, protos, Context);
   }
 
   for (const auto *proto : props->getConformsToExcludingSuperclassConformances())
@@ -153,7 +153,7 @@ Type RequirementMachine::getSuperclassBound(Type depType) const {
     return Type();
 
   auto &protos = System.getProtocols();
-  return props->getSuperclassBound({ }, protos, Context);
+  return props->getSuperclassBound({ }, term, protos, Context);
 }
 
 bool RequirementMachine::isConcreteType(Type depType) const {
@@ -183,7 +183,7 @@ Type RequirementMachine::getConcreteType(Type depType) const {
     return Type();
 
   auto &protos = System.getProtocols();
-  return props->getConcreteType({ }, protos, Context);
+  return props->getConcreteType({ }, term, protos, Context);
 }
 
 bool RequirementMachine::areSameTypeParameterInContext(Type depType1,
@@ -349,7 +349,8 @@ Type RequirementMachine::getCanonicalTypeInContext(
       if (props) {
         if (props->isConcreteType()) {
           auto concreteType = props->getConcreteType(genericParams,
-                                                     protos, Context);
+                                                     prefix, protos,
+                                                     Context);
           if (!concreteType->hasTypeParameter())
             return concreteType;
 
@@ -364,7 +365,8 @@ Type RequirementMachine::getCanonicalTypeInContext(
         if (props->hasSuperclassBound() &&
             prefix.size() != term.size()) {
           auto superclass = props->getSuperclassBound(genericParams,
-                                                      protos, Context);
+                                                      prefix, protos,
+                                                      Context);
           if (!superclass->hasTypeParameter())
             return superclass;
 

lib/AST/RequirementMachine/PropertyMap.cpp

@@ -144,6 +144,7 @@ static unsigned getGenericParamIndex(Type type) {
 static Type getTypeFromSubstitutionSchema(Type schema,
                                           ArrayRef<Term> substitutions,
                               TypeArrayView<GenericTypeParamType> genericParams,
+                                          const MutableTerm &prefix,
                                           const ProtocolGraph &protos,
                                           RewriteContext &ctx) {
   assert(!schema->isTypeParameter() && "Must have a concrete type here");
@@ -155,41 +156,76 @@ static Type getTypeFromSubstitutionSchema(Type schema,
     if (t->is<GenericTypeParamType>()) {
       auto index = getGenericParamIndex(t);
 
-      return ctx.getTypeForTerm(substitutions[index],
-                                genericParams, protos);
+      // Prepend the prefix of the lookup key to the substitution, skipping
+      // creation of a new MutableTerm in the case where the prefix is empty.
+      if (prefix.empty()) {
+        return ctx.getTypeForTerm(substitutions[index], genericParams, protos);
+      } else {
+        MutableTerm substitution(prefix);
+        substitution.append(substitutions[index]);
+        return ctx.getTypeForTerm(substitution, genericParams, protos);
+      }
     }
 
     assert(!t->isTypeParameter());
     return None;
   });
 }
 
-/// Get the superclass bound for the term represented by this property bag.
+/// Given a term \p lookupTerm whose suffix must equal this property bag's
+/// key, return a new term with that suffix stripped off. Will be empty if
+/// \p lookupTerm exactly equals the key.
+MutableTerm
+PropertyBag::getPrefixAfterStrippingKey(const MutableTerm &lookupTerm) const {
+  assert(lookupTerm.size() >= Key.size());
+  auto prefixBegin = lookupTerm.begin();
+  auto prefixEnd = lookupTerm.end() - Key.size();
+  assert(std::equal(prefixEnd, lookupTerm.end(), Key.begin()) &&
+         "This is not the bag you're looking for");
+  return MutableTerm(prefixBegin, prefixEnd);
+}
+
+/// Get the superclass bound for \p lookupTerm, whose suffix must be the term
+/// represented by this property bag.
+///
+/// The original \p lookupTerm is important in case the concrete type has
+/// substitutions. For example, if \p lookupTerm is [P:A].[U:B], and this
+/// property bag records that the suffix [U:B] has a superclass symbol
+/// [superclass: Cache<τ_0_0> with <[U:C]>], then we actually need to
+/// apply the substitution τ_0_0 := [P:A].[U:C] to the type 'Cache<τ_0_0>'.
 ///
 /// Asserts if this property bag does not have a superclass bound.
 Type PropertyBag::getSuperclassBound(
     TypeArrayView<GenericTypeParamType> genericParams,
+    const MutableTerm &lookupTerm,
     const ProtocolGraph &protos,
     RewriteContext &ctx) const {
+  MutableTerm prefix = getPrefixAfterStrippingKey(lookupTerm);
   return getTypeFromSubstitutionSchema(Superclass->getSuperclass(),
                                        Superclass->getSubstitutions(),
-                                       genericParams,
-                                       protos,
-                                       ctx);
+                                       genericParams, prefix,
+                                       protos, ctx);
 }
 
 /// Get the concrete type of the term represented by this property bag.
 ///
+/// The original \p lookupTerm is important in case the concrete type has
+/// substitutions. For example, if \p lookupTerm is [P:A].[U:B], and this
+/// property bag records that the suffix [U:B] has a concrete type symbol
+/// [concrete: Array<τ_0_0> with <[U:C]>], then we actually need to
+/// apply the substitution τ_0_0 := [P:A].[U:C] to the type 'Array<τ_0_0>'.
+///
 /// Asserts if this property bag is not concrete.
 Type PropertyBag::getConcreteType(
     TypeArrayView<GenericTypeParamType> genericParams,
+    const MutableTerm &lookupTerm,
     const ProtocolGraph &protos,
     RewriteContext &ctx) const {
+  MutableTerm prefix = getPrefixAfterStrippingKey(lookupTerm);
   return getTypeFromSubstitutionSchema(ConcreteType->getConcreteType(),
                                        ConcreteType->getSubstitutions(),
-                                       genericParams,
-                                       protos,
-                                       ctx);
+                                       genericParams, prefix,
+                                       protos, ctx);
 }
 
 /// Computes the term corresponding to a member type access on a substitution.

lib/AST/RequirementMachine/PropertyMap.h

@@ -101,6 +101,7 @@ class PropertyBag {
 
   Type getSuperclassBound(
       TypeArrayView<GenericTypeParamType> genericParams,
+      const MutableTerm &lookupTerm,
       const ProtocolGraph &protos,
       RewriteContext &ctx) const;
 
@@ -114,6 +115,7 @@ class PropertyBag {
 
   Type getConcreteType(
       TypeArrayView<GenericTypeParamType> genericParams,
+      const MutableTerm &lookupTerm,
       const ProtocolGraph &protos,
       RewriteContext &ctx) const;
 
@@ -127,6 +129,8 @@ class PropertyBag {
 
   llvm::TinyPtrVector<const ProtocolDecl *>
   getConformsToExcludingSuperclassConformances() const;
+
+  MutableTerm getPrefixAfterStrippingKey(const MutableTerm &lookupTerm) const;
 };
 
 /// Stores all rewrite rules of the form T.[p] => T, where [p] is a property

test/Generics/concrete_type_property_of_suffix.swift

@@ -0,0 +1,16 @@
+// RUN: %target-typecheck-verify-swift -requirement-machine=verify
+
+protocol P {
+  associatedtype T where T == U?
+  associatedtype U
+}
+
+func sameType<T>(_: T.Type, _: T.Type) {}
+
+func foo<X : P, Y : P>(_: X, _: Y) {
+  // X.T is Optional<X.U>.
+  sameType(X.T.self, X.U?.self)
+
+  // Y.T is Optional<Y.U>.
+  sameType(Y.T.self, Y.U?.self)
+}